smart redaction (#6195)

Co-authored-by: James Brunton <[email protected]>

engine/src/stirling/agents/orchestrator.py

@@ -86,7 +86,7 @@ class OrchestratorAgent:
             system_prompt=(
                 "You are the top-level orchestrator. "
                 "Choose exactly one output function that best handles the request. "
-                "Use delegate_pdf_edit for requested modifications of single or multiple PDFs. "
+                "Use delegate_pdf_edit for any requested modification of one or more PDFs. "
                 "Use delegate_pdf_question for questions about the contents of the attached PDFs. "
                 "Use delegate_user_spec for requests to create or define an agent spec. "
                 "Use delegate_pdf_review when the user wants the PDF returned with review"

engine/src/stirling/agents/pdf_edit.py

@@ -39,11 +39,39 @@ class PdfEditPlanSelection(ApiModel):
     summary: str
 
 
-type PdfEditPlanOutput = PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse | NeedContentResponse
+class PdfEditNeedContentSelection(ApiModel):
+    """LLM-facing variant of need_content: the model signals it needs document content and gives
+    a reason. File objects are resolved by Python from request.files by matching names, so the
+    LLM can't fabricate file ids — it only selects by the display names it sees in the prompt.
+    """
+
+    outcome: Literal["need_content"] = "need_content"
+    reason: str
+    file_names: list[str] | None = Field(
+        default=None,
+        description=(
+            "Names of files whose content is needed. "
+            "Use the exact names shown in the prompt. "
+            "Omit or leave empty to request content from all files."
+        ),
+    )
+    max_pages: int | None = None
+    max_characters: int | None = None
+
+
+type PdfEditPlanOutput = (
+    PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse | PdfEditNeedContentSelection
+)
 
 
 class PdfEditSelectionAgent:
-    def __init__(self, runtime: AppRuntime, base_system_prompt: str, *, allow_need_content: bool) -> None:
+    def __init__(
+        self,
+        runtime: AppRuntime,
+        base_system_prompt: str,
+        *,
+        allow_need_content: bool,
+    ) -> None:
         self.runtime = runtime
         output_types: list[type[PdfEditPlanOutput]] = [
             PdfEditPlanSelection,
@@ -52,11 +80,12 @@ class PdfEditSelectionAgent:
         ]
         system_prompt = base_system_prompt
         if allow_need_content:
-            output_types.append(NeedContentResponse)
+            output_types.append(PdfEditNeedContentSelection)
             system_prompt += (
                 " Return need_content when planning a correct answer requires inspecting the actual PDF "
                 "page text (e.g. 'split after every page that says NEW PAGE', "
-                "'rotate pages that mention draft')."
+                "'rotate pages that mention draft'). "
+                "Set file_names to only the files that need to be read; omit it to read all files."
             )
         self.agent = Agent(
             model=runtime.smart_model,
@@ -101,14 +130,18 @@ class PdfEditParameterSelector:
         parameter_result = await self.agent.run(
             prompt,
             output_type=NativeOutput(parameter_model),
-            instructions=(
-                f"Generate only the parameters for the PDF operation `{operation_id.name}`. "
-                "Do not include fields from any other operation."
-            ),
+            instructions=self._get_operation_instructions(operation_id),
         )
         logger.debug("[pdf-edit params %s] output: %s", operation_id.name, Pretty(parameter_result.output))
         return parameter_result.output
 
+    @staticmethod
+    def _get_operation_instructions(operation_id: ToolEndpoint) -> str:
+        return (
+            f"Generate only the parameters for the PDF operation `{operation_id.name}`. "
+            "Do not include fields from any other operation."
+        )
+
     def _build_parameter_prompt(
         self,
         request: PdfEditRequest,
@@ -173,8 +206,7 @@ class PdfEditAgent:
             request.enabled_endpoints,
             request.user_message,
         )
-        supported_operations = self._get_supported_operations(request)
-        unavailable_operations = self._get_unavailable_operations(supported_operations)
+        supported_operations, unavailable_operations = self._classify_operations(request)
         if not supported_operations:
             return EditCannotDoResponse(reason="No PDF edit operations are available on this server.")
         selection = await self._select_plan(
@@ -183,9 +215,9 @@ class PdfEditAgent:
         if isinstance(selection, EditClarificationRequest | EditCannotDoResponse):
             logger.info("[pdf-edit] selection -> %s: %s", selection.outcome, Pretty(selection))
             return selection
-        if isinstance(selection, NeedContentResponse):
+        if isinstance(selection, PdfEditNeedContentSelection):
             logger.info("[pdf-edit] selection -> need_content: %s", selection.reason)
-            return self._fill_need_content_defaults(selection, request)
+            return self._build_need_content_response(selection, request)
         enabled = set(supported_operations)
         unsupported = [op for op in selection.operations if op not in enabled]
         if unsupported:
@@ -228,7 +260,9 @@ class PdfEditAgent:
     ) -> PdfEditPlanOutput:
         can_request_content = allow_need_content and not has_page_text(request.page_text)
         agent = self._build_selection_agent(
-            supported_operations, unavailable_operations, allow_need_content=can_request_content
+            supported_operations,
+            unavailable_operations,
+            allow_need_content=can_request_content,
         )
         return await agent.select(self._build_selection_prompt(request, supported_operations, unavailable_operations))
 
@@ -263,9 +297,11 @@ class PdfEditAgent:
                 "merging, or extracting pages then re-inserting them). "
                 "Only return cannot_do when no sequence of the supported operations could achieve the request. "
                 "Do not produce operation parameters in this stage. "
-                "Return need_clarification when the request is genuinely ambiguous. "
                 "Return plan when a reasonable multi-step plan can be created. "
-                "Never return partial plans."
+                "Never return partial plans. "
+                "Return need_clarification only when the request is genuinely ambiguous in a way "
+                "that no reasonable interpretation could produce a correct plan — do not ask to "
+                "confirm details that are already clear from the user's message."
             ),
             allow_need_content=allow_need_content,
         )
@@ -291,13 +327,31 @@ class PdfEditAgent:
             f"Extracted page text:\n{format_page_text(request.page_text)}"
         )
 
-    def _get_supported_operations(self, request: PdfEditRequest) -> Iterable[ToolEndpoint]:
-        return request.enabled_endpoints
+    # Endpoints that exist on the server and are callable via the direct API or the manual UI,
+    # but are never offered to the AI agent as a routing option.
+    #
+    # Why: REDACT_EXECUTE is the preferred AI-driven redaction route. AUTO_REDACT and REDACT are
+    # legacy endpoints that remain fully functional for human callers (the manual redact UI, direct
+    # API consumers, pipelines) but would produce a worse experience if the AI routed to them —
+    # they accept a simpler, less expressive schema and pre-date the unified operation model.
+    # Hiding them here channels all AI redaction traffic through REDACT_EXECUTE without disabling
+    # the legacy endpoints for anyone else.
+    #
+    # How to reuse: add an endpoint here whenever a legacy endpoint has a preferred replacement
+    # that the AI should use exclusively. The endpoint remains live on the server; only the AI
+    # planner is prevented from selecting it.
+    _AGENT_HIDDEN_ENDPOINTS: frozenset[ToolEndpoint] = frozenset({ToolEndpoint.AUTO_REDACT, ToolEndpoint.REDACT})
 
-    @staticmethod
-    def _get_unavailable_operations(supported_operations: Iterable[ToolEndpoint]) -> Iterable[ToolEndpoint]:
-        supported_set = set(supported_operations)
-        return [op for op in OPERATIONS if op not in supported_set]
+    def _classify_operations(self, request: PdfEditRequest) -> tuple[list[ToolEndpoint], list[ToolEndpoint]]:
+        """Split the universe of operations into (supported, unavailable) from the agent's
+        point of view. Endpoints in `_AGENT_HIDDEN_ENDPOINTS` are filtered out regardless
+        of enabled state — they exist on the server but only callers outside the AI
+        pipeline (the manual redact UI, direct API consumers) can invoke them.
+        """
+        enabled_set = set(request.enabled_endpoints)
+        supported = [op for op in request.enabled_endpoints if op not in self._AGENT_HIDDEN_ENDPOINTS]
+        unavailable = [op for op in OPERATIONS if op not in enabled_set and op not in self._AGENT_HIDDEN_ENDPOINTS]
+        return supported, unavailable
 
     @staticmethod
     def _get_operations_prompt(operations: Iterable[ToolEndpoint]) -> str:
@@ -326,18 +380,29 @@ class PdfEditAgent:
                     lines.append(f"    {name}")
         return "\n".join(lines)
 
-    def _fill_need_content_defaults(
+    def _build_need_content_response(
         self,
-        selection: NeedContentResponse,
+        selection: PdfEditNeedContentSelection,
         request: PdfEditRequest,
     ) -> NeedContentResponse:
-        files = selection.files or [
-            NeedContentFileRequest(file=file, content_types=[PdfContentType.PAGE_TEXT]) for file in request.files
-        ]
+        # File objects are resolved here by matching names against request.files so the LLM
+        # can't fabricate file ids — it selects by display name, Python provides the AiFile.
+        if selection.file_names:
+            requested = set(selection.file_names)
+            files = [f for f in request.files if f.name in requested]
+            if not files:
+                # Names didn't match anything; fall back to all files rather than sending nothing.
+                logger.warning(
+                    "[pdf-edit] need_content file_names %s matched no request files — using all",
+                    selection.file_names,
+                )
+                files = request.files
+        else:
+            files = request.files
         return NeedContentResponse(
             resume_with=SupportedCapability.PDF_EDIT,
             reason=selection.reason,
-            files=files,
+            files=[NeedContentFileRequest(file=file, content_types=[PdfContentType.PAGE_TEXT]) for file in files],
             max_pages=selection.max_pages or self.runtime.settings.max_pages,
             max_characters=selection.max_characters or self.runtime.settings.max_characters,
         )

engine/src/stirling/contracts/__init__.py

@@ -116,7 +116,6 @@ from .progress import (
 )
 
 __all__ = [
-    "AiFile",
     "AgentDraft",
     "AgentDraftRequest",
     "AgentDraftResponse",
@@ -127,6 +126,7 @@ __all__ = [
     "AgentRevisionWorkflowResponse",
     "AgentSpec",
     "AgentSpecStep",
+    "AiFile",
     "AiToolAgentStep",
     "ArtifactKind",
     "CannotContinueExecutionAction",

engine/src/stirling/models/tool_models.py

@@ -542,6 +542,18 @@ class HtmlToPdfParams(ApiModel):
     zoom: float = Field(1, description="Zoom level for displaying the website. Default is '1'.")
 
 
+class ImageBox(ApiModel):
+    """
+    Rectangular areas to black out, each defined by a page number and bounding box coordinates.
+    """
+
+    page_index: int = Field(..., description="0-indexed page number (first page = 0).")
+    x1: float = Field(..., description="Left x coordinate of the redaction rectangle in PDF user-space points.")
+    x2: float = Field(..., description="Right x coordinate of the redaction rectangle in PDF user-space points.")
+    y1: float = Field(..., description="Top y coordinate of the redaction rectangle in PDF user-space points.")
+    y2: float = Field(..., description="Bottom y coordinate of the redaction rectangle in PDF user-space points.")
+
+
 class ColorType(StrEnum):
     """
     The color type of the output image(s)
@@ -984,6 +996,27 @@ class RearrangePagesParams(ApiModel):
     )
 
 
+class Strategy(StrEnum):
+    """
+    Execution strategy hint for the redaction pipeline
+    """
+
+    auto = "AUTO"
+    overlay_only = "OVERLAY_ONLY"
+    image_finalize = "IMAGE_FINALIZE"
+
+
+class Style(ApiModel):
+    """
+    Redaction style options
+    """
+
+    color: str = Field("#000000", description="Hex redaction box color")
+    convert_to_image: bool = Field(False, description="Rasterize output to prevent text extraction")
+    padding: float = Field(0, description="Extra padding around each box in points")
+    strategy: Strategy = Field(Strategy.auto, description="Execution strategy hint for the redaction pipeline")
+
+
 class RedactionArea(ApiModel):
     """
     A list of areas that should be redacted
@@ -1279,6 +1312,23 @@ class SvgToPdfParams(ApiModel):
     )
 
 
+class TextRange(ApiModel):
+    """
+    Text ranges to redact by specifying a start and end anchor phrase. All content between the two phrases (inclusive) is redacted. Anchors work best when short and unique. They must appear verbatim in the document.
+    """
+
+    end_string: str = Field(
+        ...,
+        description="A short, distinctive phrase (5–15 words) that marks where redaction ends (inclusive). Must appear verbatim in the document. Shorter phrases match more reliably.",
+        min_length=1,
+    )
+    start_string: str = Field(
+        ...,
+        description="A short, distinctive phrase (5–15 words) that marks where redaction begins (inclusive). Must appear verbatim in the document — e.g. a section heading or a unique sentence fragment.",
+        min_length=1,
+    )
+
+
 class TimestampPdfParams(ApiModel):
     tsa_url: str = Field(
         "http://timestamp.digicert.com",
@@ -1346,6 +1396,32 @@ class VectorToPdfParams(ApiModel):
     prepress: Prepress = Field(Prepress.boolean_false, description="Apply Ghostscript prepress settings")
 
 
+class RedactExecuteParams(ApiModel):
+    image_boxes: list[ImageBox] | None = Field(
+        None, description="Rectangular areas to black out, each defined by a page number and bounding box coordinates."
+    )
+    ranges: list[TextRange] | None = Field(
+        None,
+        description="Text ranges to redact by specifying a start and end anchor phrase. All content between the two phrases (inclusive) is redacted. Anchors work best when short and unique. They must appear verbatim in the document.",
+    )
+    redact_image_pages: list[int] | None = Field(
+        None,
+        description="1-indexed page numbers to redact all detected images from. Pass an empty list to redact images from every page. Omit or pass null to skip image redaction entirely.",
+    )
+    regex_patterns: list[str] | None = Field(
+        None,
+        description="Regex patterns to match and redact. Each match anywhere in the document is blacked out. Uses Java/PCRE regex syntax. Well-suited for strings that follow known patterns, like phone numbers, email addresses, national ID numbers, or dates (which can appear with different separators, optional country codes, etc.). For fixed known strings such as names, use textValues instead.",
+    )
+    style: Style | None = Field(None, description="Redaction style options")
+    text_values: list[str] | None = Field(
+        None,
+        description="Exact strings to find and black out. One entry per phrase to redact. Best for known names, identifiers, and specific text found in the document.",
+    )
+    wipe_pages: list[int] | None = Field(
+        None, description="1-indexed page numbers to wipe entirely (all content removed from those pages)."
+    )
+
+
 class RedactParams(ApiModel):
     convert_pdf_to_image: bool = Field(False, description="Convert the redacted PDF to an image")
     page_numbers: str = Field(
@@ -1420,6 +1496,7 @@ class Model(
         | SessionsParams
         | ValidateCertificateParams
         | RedactParams
+        | RedactExecuteParams
         | RemovePasswordParams
         | SanitizePdfParams
         | TimestampPdfParams
@@ -1488,6 +1565,7 @@ class Model(
         | SessionsParams
         | ValidateCertificateParams
         | RedactParams
+        | RedactExecuteParams
         | RemovePasswordParams
         | SanitizePdfParams
         | TimestampPdfParams
@@ -1557,6 +1635,7 @@ type ParamToolModel = (
     | SessionsParams
     | ValidateCertificateParams
     | RedactParams
+    | RedactExecuteParams
     | RemovePasswordParams
     | SanitizePdfParams
     | TimestampPdfParams
@@ -1627,6 +1706,7 @@ class ToolEndpoint(StrEnum):
     SESSIONS = "/api/v1/security/cert-sign/sessions"
     VALIDATE_CERTIFICATE = "/api/v1/security/cert-sign/validate-certificate"
     REDACT = "/api/v1/security/redact"
+    REDACT_EXECUTE = "/api/v1/security/redact-execute"
     REMOVE_PASSWORD = "/api/v1/security/remove-password"
     SANITIZE_PDF = "/api/v1/security/sanitize-pdf"
     TIMESTAMP_PDF = "/api/v1/security/timestamp-pdf"
@@ -1695,6 +1775,7 @@ OPERATIONS: dict[ToolEndpoint, ParamToolModelType] = {
     ToolEndpoint.SESSIONS: SessionsParams,
     ToolEndpoint.VALIDATE_CERTIFICATE: ValidateCertificateParams,
     ToolEndpoint.REDACT: RedactParams,
+    ToolEndpoint.REDACT_EXECUTE: RedactExecuteParams,
     ToolEndpoint.REMOVE_PASSWORD: RemovePasswordParams,
     ToolEndpoint.SANITIZE_PDF: SanitizePdfParams,
     ToolEndpoint.TIMESTAMP_PDF: TimestampPdfParams,

engine/tests/test_pdf_edit_agent.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass
 import pytest
 
 from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
-from stirling.agents.pdf_edit import PdfEditPlanOutput
+from stirling.agents.pdf_edit import PdfEditNeedContentSelection, PdfEditPlanOutput
 from stirling.contracts import (
     AiFile,
     EditCannotDoResponse,
@@ -82,12 +82,12 @@ class StubPdfEditAgent(PdfEditAgent):
         if parameter_selector is not None:
             self.parameter_selector = parameter_selector
 
-    def _get_supported_operations(self, request: PdfEditRequest) -> Iterable[ToolEndpoint]:
+    def _classify_operations(self, request: PdfEditRequest) -> tuple[list[ToolEndpoint], list[ToolEndpoint]]:
         # Tests construct requests without `enabled_endpoints`; pretend everything is enabled
         # unless the test explicitly supplies an enabled set.
-        if request.enabled_endpoints:
-            return request.enabled_endpoints
-        return OPERATIONS
+        if not request.enabled_endpoints:
+            return list(OPERATIONS), []
+        return super()._classify_operations(request)
 
     async def _select_plan(
         self,
@@ -194,12 +194,8 @@ async def test_pdf_edit_agent_returns_need_content_without_building_plan(runtime
     parameter_selector = RecordingParameterSelector()
     agent = StubPdfEditAgent(
         runtime,
-        NeedContentResponse(
-            resume_with=SupportedCapability.PDF_EDIT,
+        PdfEditNeedContentSelection(
             reason="Need page text to locate the NEW PAGE markers.",
-            files=[],
-            max_pages=0,
-            max_characters=0,
         ),
         parameter_selector=parameter_selector,
     )
@@ -274,8 +270,8 @@ async def test_pdf_edit_selection_agent_excludes_need_content_from_schema_when_n
     can_request = PdfEditSelectionAgent(runtime, "base", allow_need_content=True)
     cannot_request = PdfEditSelectionAgent(runtime, "base", allow_need_content=False)
 
-    assert NeedContentResponse in _agent_output_types(can_request)
-    assert NeedContentResponse not in _agent_output_types(cannot_request)
+    assert PdfEditNeedContentSelection in _agent_output_types(can_request)
+    assert PdfEditNeedContentSelection not in _agent_output_types(cannot_request)
 
 
 def _agent_output_types(agent: object) -> list[type]:
@@ -336,7 +332,7 @@ async def test_pdf_edit_agent_supported_operations_defaults_to_empty(
     runtime: AppRuntime,
 ) -> None:
     agent = PdfEditAgent(runtime)
-    supported = agent._get_supported_operations(PdfEditRequest(user_message="hi"))
+    supported, _ = agent._classify_operations(PdfEditRequest(user_message="hi"))
 
     assert list(supported) == []
 
@@ -351,7 +347,7 @@ async def test_pdf_edit_agent_supported_operations_uses_provided_list(
         enabled_endpoints=[ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF],
     )
 
-    supported = agent._get_supported_operations(request)
+    supported, _ = agent._classify_operations(request)
 
     assert list(supported) == [ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF]
 
@@ -372,8 +368,7 @@ def test_pdf_edit_selection_prompt_includes_unavailable_operations(runtime: AppR
         user_message="Run OCR.",
         enabled_endpoints=[ToolEndpoint.FLATTEN],
     )
-    supported = agent._get_supported_operations(request)
-    unavailable = agent._get_unavailable_operations(supported)
+    supported, unavailable = agent._classify_operations(request)
 
     prompt = agent._build_selection_prompt(request, supported, unavailable)