Add document extraction tool using Docling

Add a document extractor tool with an extract_document action that converts an
input artifact (PDF, docx, pptx, ...) into structured JSON by running a fixed
Docling program in the sandbox (the document and parameters travel as data, so
nothing in them can execute). Output is a compact payload bounded by an input
size cap, a head-and-tail markdown window, and per-table caps, can be validated
against a JSON schema, and is persisted as a data artifact by reference. Docling
is heavy, so it stays out of the base image and ships in an opt-in sandbox image
variant.

application/agents/tools/document_extractor.py

@@ -0,0 +1,422 @@
+"""Document Extractor tool: convert an input artifact to schema-validated structured JSON via Docling.
+
+The ``extract_document`` action stages a parent-scoped input artifact into the
+sandbox workspace, then runs a FIXED program that uses Docling (MIT) to convert
+the document (pdf/docx/pptx/...) into a compact structured payload (markdown +
+tables + a structured dict). The program reads its parameters from a JSON DATA
+file (``json.load``) and never string-interpolates untrusted content into code,
+so a malicious filename/parameter is treated as literal data, not executed. When
+a ``json_schema`` is supplied the compact ``structured`` payload is validated
+through the existing jsonschema path; the extracted JSON may also be persisted as
+a ``data`` artifact by reference.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import uuid
+from typing import Any, Dict, List, Optional
+
+from application.agents.tools.base import Tool
+from application.core.json_schema_utils import (
+    JsonSchemaValidationError,
+    normalize_json_schema_payload,
+)
+from application.core.settings import settings
+from application.sandbox.artifacts_capture import persist_new_artifact
+from application.sandbox.sandbox_creator import SandboxCreator
+from application.storage.db.repositories.artifacts import ArtifactsRepository
+from application.storage.db.session import db_readonly
+from application.storage.storage_creator import StorageCreator
+from application.utils import safe_filename
+
+logger = logging.getLogger(__name__)
+
+try:
+    import jsonschema
+except Exception:  # pragma: no cover - jsonschema is a declared dependency
+    jsonschema = None  # type: ignore[assignment]
+
+# Cap the extracted payload returned to the LLM so a huge document can't flood
+# context; the full extraction is still persisted as a ``data`` artifact. When the
+# markdown exceeds the cap a head+tail window is kept so both the document's
+# beginning AND end (e.g. totals/signatures) survive within the byte budget.
+_MARKDOWN_MAX_BYTES = 8000
+_MAX_TABLES_RETURNED = 20
+# Per-table content caps so a single giant table can't bloat context.
+_MAX_TABLE_ROWS = 50
+_MAX_CELL_CHARS = 200
+
+# Mime + kind for the persisted extraction artifact (JSON by reference).
+_EXTRACT_MIME = "application/json"
+_EXTRACT_KIND = "data"
+
+# FIXED extraction program. It reads ``params.json`` (the input path + caps) as
+# DATA and writes ``result.json`` to the workspace. The params are NEVER
+# string-interpolated into the program; ``{params_path}``/``{result_path}`` are
+# server-controlled path literals only. A missing Docling install yields a clean
+# ``{"error": ...}`` result rather than a traceback.
+_EXTRACT_PROGRAM = (
+    "import json\n"
+    "params = json.load(open({params_path!r}))\n"
+    "result_path = {result_path!r}\n"
+    "def _write(obj):\n"
+    "    with open(result_path, 'w') as fh:\n"
+    "        json.dump(obj, fh)\n"
+    "try:\n"
+    "    from docling.document_converter import DocumentConverter\n"
+    "except Exception as exc:\n"
+    "    _write({{'error': 'docling is not available in the sandbox runner: '\n"
+    "            + type(exc).__name__ + ': ' + str(exc)}})\n"
+    "    raise SystemExit(0)\n"
+    "src = params['input_path']\n"
+    "md_cap = int(params.get('markdown_max_bytes', 0)) or None\n"
+    "table_cap = int(params.get('max_tables', 0)) or None\n"
+    "try:\n"
+    "    converter = DocumentConverter()\n"
+    "    doc = converter.convert(src).document\n"
+    "    markdown = doc.export_to_markdown()\n"
+    "    structured = doc.export_to_dict()\n"
+    "    tables = []\n"
+    "    for tbl in getattr(doc, 'tables', []) or []:\n"
+    "        try:\n"
+    "            df = tbl.export_to_dataframe()\n"
+    "            tables.append({{'columns': [str(c) for c in df.columns],\n"
+    "                           'rows': df.astype(str).values.tolist()}})\n"
+    "        except Exception:\n"
+    "            try:\n"
+    "                tables.append({{'markdown': tbl.export_to_markdown()}})\n"
+    "            except Exception:\n"
+    "                continue\n"
+    "        if table_cap is not None and len(tables) >= table_cap:\n"
+    "            break\n"
+    "    page_count = len(getattr(doc, 'pages', {{}}) or {{}})\n"
+    "    md_truncated = False\n"
+    "    if md_cap is not None and len(markdown) > md_cap:\n"
+    "        head = md_cap // 2\n"
+    "        tail = md_cap - head\n"
+    "        dropped = len(markdown) - head - tail\n"
+    "        markdown = (markdown[:head]\n"
+    "                    + '\\n\\n...[truncated ' + str(dropped) + ' chars]...\\n\\n'\n"
+    "                    + markdown[-tail:])\n"
+    "        md_truncated = True\n"
+    "    _write({{'ok': True, 'markdown': markdown, 'markdown_truncated': md_truncated,\n"
+    "            'tables': tables, 'page_count': page_count, 'structured': structured}})\n"
+    "except Exception as exc:\n"
+    "    _write({{'error': 'docling extraction failed: ' + type(exc).__name__ + ': ' + str(exc)}})\n"
+)
+
+
+class DocumentExtractorTool(Tool):
+    """Document Extractor
+    Convert an input document artifact (pdf/docx/pptx/...) to compact, schema-validated structured JSON via Docling.
+    """
+
+    def __init__(self, tool_config: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None) -> None:
+        """Bind the tool to the invoker and its conversation/run-scoped sandbox session."""
+        self.config: Dict[str, Any] = tool_config or {}
+        self.user_id: Optional[str] = user_id
+        self.tool_id: Optional[str] = self.config.get("tool_id")
+        self.conversation_id: Optional[str] = self.config.get("conversation_id")
+        self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id")
+        self._last_artifact_id: Optional[str] = None
+
+    # ------------------------------------------------------------------
+    # Tool ABC
+    # ------------------------------------------------------------------
+    def get_actions_metadata(self) -> List[Dict[str, Any]]:
+        """Return JSON metadata describing the ``extract_document`` action for tool schemas."""
+        return [
+            {
+                "name": "extract_document",
+                "description": (
+                    "Extract a document artifact (pdf/docx/pptx/...) into compact structured JSON "
+                    "(markdown + tables + structure) using Docling. Optionally validate the result "
+                    "against a json_schema and persist it as a downloadable data artifact."
+                ),
+                "active": True,
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "input": {
+                            "type": "string",
+                            "description": "Artifact id (from this conversation/run) of the document to extract.",
+                        },
+                        "json_schema": {
+                            "type": "object",
+                            "description": "Optional JSON schema the extracted 'structured' payload must satisfy.",
+                        },
+                        "persist": {
+                            "type": "boolean",
+                            "description": "Persist the extracted JSON as a downloadable data artifact (default true).",
+                        },
+                    },
+                    "required": ["input"],
+                },
+            }
+        ]
+
+    def get_config_requirements(self) -> Dict[str, Any]:
+        """Return configuration requirements (none beyond the deployment sandbox backend)."""
+        return {}
+
+    def get_artifact_id(self, action_name: str, **kwargs: Any) -> Optional[str]:
+        """Return the persisted extraction artifact id so the UI artifact rail lights up."""
+        return self._last_artifact_id
+
+    # ------------------------------------------------------------------
+    # Dispatch
+    # ------------------------------------------------------------------
+    def execute_action(self, action_name: str, **kwargs: Any) -> Dict[str, Any]:
+        """Dispatch a tool action; only ``extract_document`` is supported."""
+        self._last_artifact_id = None
+        if action_name != "extract_document":
+            return {"status": "error", "error": f"unknown action: {action_name}"}
+        if not self.user_id:
+            return {"status": "error", "error": "document_extractor requires a valid user_id."}
+        if self.conversation_id is None and self.workflow_run_id is None:
+            return {"status": "error", "error": "document_extractor requires a conversation_id or workflow_run_id."}
+        return self._extract(**kwargs)
+
+    # ------------------------------------------------------------------
+    # Extraction
+    # ------------------------------------------------------------------
+    def _extract(self, **kwargs: Any) -> Dict[str, Any]:
+        """Stage the parent-scoped input, run the fixed Docling program, validate, and (optionally) persist."""
+        input_id = kwargs.get("input")
+        json_schema = kwargs.get("json_schema")
+        should_persist = kwargs.get("persist", True)
+
+        if not isinstance(input_id, str) or not input_id.strip():
+            return {"status": "error", "error": "input artifact id is required."}
+        if json_schema is not None:
+            schema_err = self._check_schema(json_schema)
+            if schema_err is not None:
+                return schema_err
+
+        session_id = self._resolve_session_id()
+        if session_id is None:
+            return {"status": "error", "error": "document_extractor requires a conversation_id or workflow_run_id."}
+
+        loaded = self._load_input(input_id.strip())
+        if loaded.get("error"):
+            return {"status": "error", "error": loaded["error"]}
+
+        max_input = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024))
+        if len(loaded["data"]) > max_input:
+            return {
+                "status": "error",
+                "error": f"input artifact is too large: {len(loaded['data'])} bytes exceeds the "
+                f"{max_input}-byte sandbox input cap.",
+            }
+
+        token = uuid.uuid4().hex
+        input_path = f"extract/{token}/inputs/{loaded['filename']}"
+        params_path = f"extract/{token}/params.json"
+        result_path = f"extract/{token}/result.json"
+        params = {
+            "input_path": input_path,
+            "markdown_max_bytes": _MARKDOWN_MAX_BYTES,
+            "max_tables": _MAX_TABLES_RETURNED,
+        }
+        program = _EXTRACT_PROGRAM.format(params_path=params_path, result_path=result_path)
+        timeout = float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60))
+
+        manager = SandboxCreator.get_manager()
+        try:
+            manager.open(session_id, ttl=timeout)
+        except Exception as exc:
+            logger.exception("document_extractor: failed to open sandbox session")
+            return {"status": "error", "error": f"sandbox unavailable: {type(exc).__name__}: {exc}"}
+        try:
+            # The document bytes and the params ride in as DATA files the program
+            # reads; neither is interpolated into the program, so a hostile
+            # filename or document content stays inert data.
+            manager.put_file(session_id, input_path, loaded["data"])
+            manager.put_file(session_id, params_path, json.dumps(params).encode("utf-8"))
+            result = manager.exec(session_id, program, timeout=timeout)
+            if not result.ok:
+                detail = (
+                    f"{result.error_name}: {result.error_value}"
+                    if result.error_name
+                    else (result.error_value or "extraction failed")
+                )
+                return {"status": "error", "error": f"extraction failed: {detail}"}
+            raw = manager.get_file(session_id, result_path)
+        except Exception as exc:
+            logger.exception("document_extractor: extraction failed")
+            return {"status": "error", "error": f"extraction failed: {type(exc).__name__}: {exc}"}
+        finally:
+            try:
+                manager.close(session_id)
+            except Exception:
+                logger.exception("document_extractor: session close failed")
+
+        return self._finish(raw, loaded, json_schema, should_persist)
+
+    def _finish(
+        self,
+        raw: bytes,
+        loaded: Dict[str, Any],
+        json_schema: Any,
+        should_persist: Any,
+    ) -> Dict[str, Any]:
+        """Parse the program result, validate against json_schema, and shape the compact payload."""
+        if not raw:
+            return {"status": "error", "error": "extractor produced no result."}
+        try:
+            extracted = json.loads(raw.decode("utf-8"))
+        except (ValueError, UnicodeDecodeError):
+            return {"status": "error", "error": "extractor produced an invalid result."}
+        if isinstance(extracted, dict) and extracted.get("error"):
+            return {"status": "error", "error": str(extracted["error"])}
+        if not isinstance(extracted, dict) or not extracted.get("ok"):
+            return {"status": "error", "error": "extractor produced an unexpected result."}
+
+        structured = extracted.get("structured")
+        if json_schema is not None:
+            valid = self._validate(json_schema, structured)
+            if valid is not None:
+                return valid
+
+        compact = self._compact_payload(extracted)
+        payload: Dict[str, Any] = {"status": "ok", "structured": compact}
+        if should_persist:
+            ref = self._persist(extracted, loaded["title"])
+            if ref is not None:
+                self._last_artifact_id = ref["artifact_id"]
+                payload["artifact"] = ref
+        return payload
+
+    # ------------------------------------------------------------------
+    # Input / payload helpers
+    # ------------------------------------------------------------------
+    def _load_input(self, artifact_id: str) -> Dict[str, Any]:
+        """Fetch the parent-scoped input artifact's current-version bytes; never cross-tenant."""
+        try:
+            with db_readonly() as conn:
+                repo = ArtifactsRepository(conn)
+                artifact = repo.get_artifact_in_parent(
+                    artifact_id,
+                    conversation_id=self.conversation_id,
+                    workflow_run_id=self.workflow_run_id,
+                )
+                if artifact is None:
+                    return {"error": f"input artifact {artifact_id} not found in this conversation/run."}
+                version = repo.get_version(artifact_id, artifact["current_version"])
+        except Exception:
+            logger.exception("document_extractor: failed to load input artifact")
+            return {"error": f"failed to load input artifact {artifact_id}."}
+        if not version or not version.get("storage_path"):
+            return {"error": f"input artifact {artifact_id} has no stored content."}
+        display_name = version.get("filename") or artifact.get("title") or artifact_id
+        filename = safe_filename(display_name)
+        try:
+            file_obj = StorageCreator.get_storage().get_file(version["storage_path"])
+            data = file_obj.read()
+        except Exception:
+            logger.exception("document_extractor: failed to read input artifact bytes")
+            return {"error": f"failed to read input artifact {artifact_id}."}
+        return {"data": data, "filename": filename, "title": display_name}
+
+    def _compact_payload(self, extracted: Dict[str, Any]) -> Dict[str, Any]:
+        """Shape the LLM-facing payload: markdown window + bounded tables + structure summary, no raw bytes."""
+        structured = extracted.get("structured")
+        return {
+            "markdown": extracted.get("markdown", ""),
+            "markdown_truncated": bool(extracted.get("markdown_truncated")),
+            "tables": [self._compact_table(tbl) for tbl in extracted.get("tables", [])],
+            "page_count": extracted.get("page_count", 0),
+            "summary": self._structure_summary(structured),
+        }
+
+    @staticmethod
+    def _compact_table(table: Any) -> Any:
+        """Bound a single table's rows and cell sizes so one giant table can't bloat context."""
+        if not isinstance(table, dict):
+            return table
+
+        def _cell(value: Any) -> Any:
+            if isinstance(value, str) and len(value) > _MAX_CELL_CHARS:
+                return value[:_MAX_CELL_CHARS] + "...[truncated]"
+            return value
+
+        rows = table.get("rows")
+        if not isinstance(rows, list):
+            return table
+        capped = [[_cell(c) for c in row] if isinstance(row, list) else _cell(row) for row in rows[:_MAX_TABLE_ROWS]]
+        compact = dict(table)
+        compact["rows"] = capped
+        if len(rows) > _MAX_TABLE_ROWS:
+            compact["rows_truncated"] = True
+            compact["total_rows"] = len(rows)
+        return compact
+
+    @staticmethod
+    def _structure_summary(structured: Any) -> Dict[str, Any]:
+        """Summarize the Docling structured dict by top-level element counts (keeps context compact)."""
+        if not isinstance(structured, dict):
+            return {}
+        counts: Dict[str, int] = {}
+        for key in ("texts", "tables", "pictures", "groups", "pages"):
+            value = structured.get(key)
+            if isinstance(value, (list, dict)):
+                counts[key] = len(value)
+        return counts
+
+    def _persist(self, extracted: Dict[str, Any], title: str) -> Optional[Dict[str, Any]]:
+        """Persist the full extraction JSON as a ``data`` artifact by reference; return its reference."""
+        try:
+            data = json.dumps(extracted).encode("utf-8")
+        except (TypeError, ValueError):
+            logger.exception("document_extractor: extraction is not JSON-serializable")
+            return None
+        filename = f"{safe_filename(title) or 'extract'}.extract.json"
+        ref = persist_new_artifact(
+            user_id=self.user_id,
+            kind=_EXTRACT_KIND,
+            data=data,
+            filename=filename,
+            mime_type=_EXTRACT_MIME,
+            title=f"{title} (extracted)",
+            conversation_id=self.conversation_id,
+            workflow_run_id=self.workflow_run_id,
+            produced_by={"tool": "document_extractor", "action": "extract_document", "tool_id": self.tool_id},
+        )
+        return ref
+
+    # ------------------------------------------------------------------
+    # Schema validation
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _check_schema(json_schema: Any) -> Optional[Dict[str, Any]]:
+        """Return an error payload when ``json_schema`` itself is malformed, else None."""
+        try:
+            normalize_json_schema_payload(json_schema)
+        except JsonSchemaValidationError as exc:
+            return {"status": "error", "error": f"invalid json_schema: {exc}"}
+        return None
+
+    @staticmethod
+    def _validate(json_schema: Any, instance: Any) -> Optional[Dict[str, Any]]:
+        """Validate ``instance`` against the (already-normalized) json_schema; error payload on mismatch."""
+        if jsonschema is None:
+            return {"status": "error", "error": "jsonschema is required for json_schema validation."}
+        schema = normalize_json_schema_payload(json_schema)
+        try:
+            jsonschema.validate(instance=instance, schema=schema)
+        except jsonschema.exceptions.ValidationError as exc:
+            return {"status": "error", "error": f"extracted structure did not match json_schema: {exc.message}"}
+        return None
+
+    # ------------------------------------------------------------------
+    # Misc helpers
+    # ------------------------------------------------------------------
+    def _resolve_session_id(self) -> Optional[str]:
+        """Derive the sandbox session id from the bound conversation/run; sanitize to the gateway charset."""
+        raw = self.conversation_id or self.workflow_run_id
+        if not raw:
+            return None
+        sanitized = "".join(c if c.isalnum() or c in "-_" else "-" for c in str(raw))
+        return sanitized or None

application/agents/tools/tool_manager.py

@@ -39,6 +39,7 @@ class ToolManager:
                         "remote_device",
                         "code_executor",
                         "artifact_generator",
+                        "document_extractor",
                     }
                     and user_id
                 ):
@@ -60,6 +61,7 @@ class ToolManager:
                 "remote_device",
                 "code_executor",
                 "artifact_generator",
+                "document_extractor",
             }
             and user_id
         ):

application/core/settings.py

@@ -337,6 +337,7 @@ class Settings(BaseSettings):
     SANDBOX_HTTP_TIMEOUT: int = 10  # fixed cap (s) for REST control calls (create/delete/alive/interrupt)
     SANDBOX_MAX_OUTPUT_BYTES: int = 8 * 1024 * 1024  # cap on buffered stdout+stderr per exec
     SANDBOX_MAX_FILE_BYTES: int = 10 * 1024 * 1024  # cap on get_file size routed through stdout
+    SANDBOX_MAX_INPUT_BYTES: int = 25 * 1024 * 1024  # cap on an input document staged into a sandbox session
     # Runner container resource caps — consumed by the docsgpt-sandbox compose
     # service (deployment/sandbox), not by the app client. cgroup CPU/mem caps
     # are part of the untrusted-code security boundary.

deployment/sandbox/Dockerfile

@@ -32,6 +32,16 @@ RUN pip install --no-cache-dir \
     pandas==2.2.3 \
     matplotlib==3.9.2
 
+# Docling (MIT) for the document_extractor tool — OFF by default because it pulls
+# torch + models and makes the image multi-GB. Build the "extract" variant with
+# `--build-arg INSTALL_DOCLING=true` (see README) when document extraction is
+# needed. Docling is MIT and uses its own PDF backend; PyMuPDF (AGPL) is NOT
+# installed here. The base image and the app's requirements stay docling-free.
+ARG INSTALL_DOCLING=false
+RUN if [ "$INSTALL_DOCLING" = "true" ]; then \
+        pip install --no-cache-dir docling==2.8.3; \
+    fi
+
 USER sandbox
 WORKDIR /home/sandbox
 

deployment/sandbox/README.md

@@ -43,6 +43,25 @@ The `docsgpt-sandbox` service is defined in `deployment/docker-compose.yaml` on
 an internal-only network. The backend and worker reach it at
 `http://docsgpt-sandbox:8888`.
 
+## Document extraction variant (Docling)
+
+The `document_extractor` tool runs Docling (MIT) inside the runner to convert
+documents (pdf/docx/pptx/...) to schema-validated JSON. Docling pulls `torch` and
+ML models, which makes the image multi-GB, so it is **off by default**: it is not
+in the base image, not in `application/requirements.txt`, and not required in the
+dev `.venv`. Build the extract variant only where extraction is needed:
+
+```bash
+docker build -t docsgpt-sandbox-extract \
+  --build-arg INSTALL_DOCLING=true deployment/sandbox
+docker run --rm -p 8888:8888 docsgpt-sandbox-extract
+```
+
+Point the app at this runner via `SANDBOX_GATEWAY_URL` exactly as the base image.
+Docling uses its own permissive PDF backend; **PyMuPDF (AGPL) is intentionally not
+installed.** If the base (non-extract) image is used, `extract_document` returns a
+clean "docling is not available in the sandbox runner" error rather than crashing.
+
 ## Hardening (separate slice — not in this image)
 
 Egress/SSRF blocks (drop RFC1918, link-local, `169.254.169.254`), the gVisor

tests/agents/tools/test_document_extractor_live.py

@@ -0,0 +1,56 @@
+"""Optional live Docling extraction test, SKIPPED unless RUN_DOCLING_LIVE=1.
+
+Default CI/dev runs never import docling (which pulls torch + models), so this is
+skipped unless explicitly opted in. It runs the FIXED extraction program in-process
+against a tiny generated PDF and asserts the compact result shape, exercising the
+exact program string the tool ships to the sandbox.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+
+import pytest
+
+from application.agents.tools.document_extractor import _EXTRACT_PROGRAM
+
+RUN_LIVE = os.environ.get("RUN_DOCLING_LIVE") == "1"
+
+pytestmark = pytest.mark.skipif(not RUN_LIVE, reason="set RUN_DOCLING_LIVE=1 to run live Docling extraction")
+
+
+def _make_pdf(path: str) -> None:
+    """Write a one-line PDF using reportlab (kept out of the import path unless live)."""
+    from reportlab.pdfgen import canvas
+
+    c = canvas.Canvas(path)
+    c.drawString(72, 720, "Compliance balance: 1000")
+    c.showPage()
+    c.save()
+
+
+def test_live_docling_extracts_pdf():
+    pytest.importorskip("docling")
+    pytest.importorskip("reportlab")
+
+    workdir = tempfile.mkdtemp()
+    input_path = os.path.join(workdir, "inputs", "doc.pdf")
+    os.makedirs(os.path.dirname(input_path), exist_ok=True)
+    _make_pdf(input_path)
+
+    params_path = os.path.join(workdir, "params.json")
+    result_path = os.path.join(workdir, "result.json")
+    with open(params_path, "w") as fh:
+        json.dump({"input_path": input_path, "markdown_max_bytes": 8000, "max_tables": 20}, fh)
+
+    program = _EXTRACT_PROGRAM.format(params_path=params_path, result_path=result_path)
+    namespace: dict = {}
+    exec(compile(program, "<extractor>", "exec"), namespace, namespace)  # noqa: S102
+
+    with open(result_path) as fh:
+        result = json.load(fh)
+    assert result.get("ok") is True, result
+    assert "balance" in result["markdown"].lower()
+    assert isinstance(result["structured"], dict)

tests/agents/tools/test_document_extractor_unit.py

@@ -0,0 +1,538 @@
+"""Unit tests for DocumentExtractorTool: staging, fixed-program shaping, schema validation, and injection safety.
+
+Docling and the sandbox exec are MOCKED so no heavy import (torch/models) or real
+kernel is touched. These cover: a parent-scoped input is staged and extracted into
+a compact payload; json_schema validation (pass + fail); a cross-tenant input id is
+denied; and the extraction program is FIXED with params travelling as data so a
+malicious filename/parameter is never executed.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Dict, List, Optional
+
+import os
+import sys
+import tempfile
+import types
+
+import application.agents.tools.document_extractor as de
+from application.agents.tools.document_extractor import (
+    _EXTRACT_PROGRAM,
+    _MARKDOWN_MAX_BYTES,
+    _MAX_CELL_CHARS,
+    _MAX_TABLE_ROWS,
+    DocumentExtractorTool,
+)
+from application.sandbox.base import ExecResult
+
+# A canned Docling-style extraction result the mocked program "writes" to result.json.
+_CANNED_EXTRACT = {
+    "ok": True,
+    "markdown": "# Statement\n\nBalance: 1000",
+    "markdown_truncated": False,
+    "tables": [{"columns": ["item", "amount"], "rows": [["fee", "10"]]}],
+    "page_count": 2,
+    "structured": {
+        "schema_name": "DoclingDocument",
+        "texts": [{"text": "Balance: 1000"}],
+        "tables": [{"data": "..."}],
+        "pages": {"1": {}, "2": {}},
+    },
+}
+
+
+class _FakeManager:
+    """In-memory sandbox stand-in: records put_file calls and the exact exec program/code."""
+
+    def __init__(self, result: ExecResult, result_file: bytes) -> None:
+        self._result = result
+        self._result_file = result_file
+        self.put_files: Dict[str, bytes] = {}
+        self.exec_programs: List[str] = []
+        self.opened: List[Any] = []
+        self.closed: List[str] = []
+
+    def open(self, session_id, ttl=None):
+        self.opened.append((session_id, ttl))
+        return session_id
+
+    def put_file(self, session_id, dest_path, data):
+        self.put_files[dest_path] = data
+
+    def exec(self, session_id, code, timeout=None):
+        self.exec_programs.append(code)
+        return self._result
+
+    def get_file(self, session_id, path):
+        return self._result_file
+
+    def close(self, session_id):
+        self.closed.append(session_id)
+
+
+class _FakeVersion(dict):
+    pass
+
+
+def _stub_repo(monkeypatch, *, found: bool, conv: Optional[str], run: Optional[str]):
+    """Patch db_readonly + ArtifactsRepository so input scoping is exercised without a DB."""
+
+    class _Repo:
+        def __init__(self, conn):
+            pass
+
+        def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None):
+            # Mirror the real scope gate: the artifact resolves only for the right parent.
+            if not found:
+                return None
+            if conv is not None and conversation_id != conv:
+                return None
+            if run is not None and workflow_run_id != run:
+                return None
+            return {"id": artifact_id, "current_version": 1, "title": "statement.pdf"}
+
+        def get_version(self, artifact_id, version):
+            return _FakeVersion(
+                {"filename": "statement.pdf", "storage_path": f"inputs/u/artifacts/{artifact_id}/v1/statement.pdf"}
+            )
+
+    class _Conn:
+        def __enter__(self):
+            return object()
+
+        def __exit__(self, *exc):
+            return False
+
+    monkeypatch.setattr(de, "db_readonly", lambda: _Conn())
+    monkeypatch.setattr(de, "ArtifactsRepository", _Repo)
+
+
+class _FakeStorage:
+    """Returns canned input bytes for any storage path."""
+
+    def __init__(self, data: bytes = b"%PDF-1.4 fake") -> None:
+        self._data = data
+
+    def get_file(self, path):
+        import io
+
+        return io.BytesIO(self._data)
+
+
+def _patch_storage(monkeypatch):
+    monkeypatch.setattr(de.StorageCreator, "get_storage", staticmethod(lambda: _FakeStorage()))
+
+
+def _patch_manager(monkeypatch, manager):
+    monkeypatch.setattr(de.SandboxCreator, "get_manager", lambda: manager)
+
+
+def _patch_no_persist(monkeypatch):
+    # Default the persist path to a no-op so most tests don't touch the artifact store.
+    monkeypatch.setattr(de, "persist_new_artifact", lambda **kwargs: None)
+
+
+def _tool(**config) -> DocumentExtractorTool:
+    base = {"conversation_id": "conv-1", "tool_id": "t-1"}
+    base.update(config)
+    return DocumentExtractorTool(tool_config=base, user_id="u-1")
+
+
+def _manager_with_extract(extract: Dict[str, Any]) -> _FakeManager:
+    return _FakeManager(ExecResult(status="ok"), json.dumps(extract).encode("utf-8"))
+
+
+# ---------------------------------------------------------------------------
+# Guards
+# ---------------------------------------------------------------------------
+def test_unknown_action_rejected():
+    out = _tool().execute_action("nope", input="a")
+    assert out["status"] == "error" and "unknown action" in out["error"]
+
+
+def test_requires_user_and_parent():
+    no_user = DocumentExtractorTool({"conversation_id": "c"}, user_id=None)
+    out = no_user.execute_action("extract_document", input="a")
+    assert out["status"] == "error" and "user_id" in out["error"]
+
+    no_parent = DocumentExtractorTool({}, user_id="u")
+    out2 = no_parent.execute_action("extract_document", input="a")
+    assert out2["status"] == "error" and "conversation_id" in out2["error"]
+
+
+def test_input_id_required():
+    out = _tool().execute_action("extract_document", input="   ")
+    assert out["status"] == "error" and "input artifact id is required" in out["error"]
+
+
+def test_action_metadata_shape():
+    meta = _tool().get_actions_metadata()[0]
+    assert meta["name"] == "extract_document"
+    assert meta["parameters"]["required"] == ["input"]
+    assert "json_schema" in meta["parameters"]["properties"]
+
+
+# ---------------------------------------------------------------------------
+# Happy path: staging + compact payload shaping
+# ---------------------------------------------------------------------------
+def test_extract_stages_input_and_shapes_compact_payload(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    manager = _manager_with_extract(_CANNED_EXTRACT)
+    _patch_manager(monkeypatch, manager)
+
+    out = _tool().execute_action("extract_document", input="art-1", persist=False)
+
+    assert out["status"] == "ok"
+    structured = out["structured"]
+    # Compact payload: markdown + tables + page_count + a structure summary, no raw bytes.
+    assert structured["markdown"].startswith("# Statement")
+    assert structured["tables"] == _CANNED_EXTRACT["tables"]
+    assert structured["page_count"] == 2
+    assert structured["summary"] == {"texts": 1, "tables": 1, "pages": 2}
+    assert "bytes" not in out
+
+    # The input document was staged into the workspace as a DATA file (not code).
+    staged = [p for p in manager.put_files if p.endswith("statement.pdf")]
+    assert staged, "input document should be staged into the workspace"
+    assert manager.put_files[staged[0]] == b"%PDF-1.4 fake"
+    # The session is closed after a one-shot extraction.
+    assert manager.closed == ["conv-1"]
+
+
+def test_extract_persists_data_artifact_by_reference(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    manager = _manager_with_extract(_CANNED_EXTRACT)
+    _patch_manager(monkeypatch, manager)
+
+    captured: Dict[str, Any] = {}
+
+    def _fake_persist(**kwargs):
+        captured.update(kwargs)
+        return {"artifact_id": "new-art", "version": 1, "filename": "x.json",
+                "mime_type": "application/json", "size": 10}
+
+    monkeypatch.setattr(de, "persist_new_artifact", _fake_persist)
+
+    tool = _tool()
+    out = tool.execute_action("extract_document", input="art-1")  # persist defaults to true
+    assert out["status"] == "ok"
+    assert out["artifact"]["artifact_id"] == "new-art"
+    assert tool.get_artifact_id("extract_document") == "new-art"
+    # Persisted as a JSON data artifact carrying the full extraction (by reference).
+    assert captured["kind"] == "data"
+    assert captured["mime_type"] == "application/json"
+    assert json.loads(captured["data"].decode("utf-8")) == _CANNED_EXTRACT
+
+
+# ---------------------------------------------------------------------------
+# Cross-tenant scope
+# ---------------------------------------------------------------------------
+def test_cross_tenant_input_is_denied(monkeypatch):
+    # Tool is bound to conv-1; the artifact only exists under conv-OTHER, so the
+    # parent-scoped fetch returns None and the extraction is refused.
+    _stub_repo(monkeypatch, found=True, conv="conv-OTHER", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    manager = _manager_with_extract(_CANNED_EXTRACT)
+    _patch_manager(monkeypatch, manager)
+
+    out = _tool().execute_action("extract_document", input="art-1", persist=False)
+    assert out["status"] == "error"
+    assert "not found in this conversation/run" in out["error"]
+    # Nothing was staged or executed for a denied input.
+    assert manager.exec_programs == []
+
+
+def test_missing_input_artifact_is_denied(monkeypatch):
+    _stub_repo(monkeypatch, found=False, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    manager = _manager_with_extract(_CANNED_EXTRACT)
+    _patch_manager(monkeypatch, manager)
+
+    out = _tool().execute_action("extract_document", input="ghost", persist=False)
+    assert out["status"] == "error" and "not found" in out["error"]
+    assert manager.exec_programs == []
+
+
+# ---------------------------------------------------------------------------
+# json_schema validation (pass + fail)
+# ---------------------------------------------------------------------------
+def _schema_requiring_texts() -> Dict[str, Any]:
+    return {
+        "type": "object",
+        "properties": {"texts": {"type": "array"}, "schema_name": {"type": "string"}},
+        "required": ["texts", "schema_name"],
+    }
+
+
+def test_json_schema_validation_passes(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    _patch_manager(monkeypatch, _manager_with_extract(_CANNED_EXTRACT))
+
+    out = _tool().execute_action(
+        "extract_document", input="art-1", json_schema=_schema_requiring_texts(), persist=False
+    )
+    assert out["status"] == "ok"
+
+
+def test_json_schema_validation_fails_with_clean_error(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    # structured payload lacks the required "amount" field the schema demands.
+    schema = {"type": "object", "required": ["amount"], "properties": {"amount": {"type": "number"}}}
+    _patch_manager(monkeypatch, _manager_with_extract(_CANNED_EXTRACT))
+
+    out = _tool().execute_action("extract_document", input="art-1", json_schema=schema, persist=False)
+    assert out["status"] == "error"
+    assert "did not match json_schema" in out["error"]
+
+
+def test_malformed_json_schema_rejected_before_run(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    manager = _manager_with_extract(_CANNED_EXTRACT)
+    _patch_manager(monkeypatch, manager)
+
+    # No "type"/"schema" key -> normalize_json_schema_payload rejects it up front.
+    out = _tool().execute_action("extract_document", input="art-1", json_schema={"properties": {}}, persist=False)
+    assert out["status"] == "error" and "invalid json_schema" in out["error"]
+    assert manager.exec_programs == []
+
+
+# ---------------------------------------------------------------------------
+# Docling-unavailable / extractor error surfacing
+# ---------------------------------------------------------------------------
+def test_docling_unavailable_surfaces_clean_error(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    # The program's own "docling not available" branch writes an {"error": ...} result.
+    err_result = {
+        "error": "docling is not available in the sandbox runner: ModuleNotFoundError: No module named 'docling'"
+    }
+    _patch_manager(monkeypatch, _manager_with_extract(err_result))
+
+    out = _tool().execute_action("extract_document", input="art-1", persist=False)
+    assert out["status"] == "error"
+    assert "docling is not available" in out["error"]
+
+
+def test_exec_error_surfaces_clean_error(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    manager = _FakeManager(
+        ExecResult(status="error", error_name="TimeoutError", error_value="exceeded 60s"), b""
+    )
+    _patch_manager(monkeypatch, manager)
+
+    out = _tool().execute_action("extract_document", input="art-1", persist=False)
+    assert out["status"] == "error"
+    assert "TimeoutError: exceeded 60s" in out["error"]
+
+
+# ---------------------------------------------------------------------------
+# Injection safety: the extraction program is FIXED; params travel as data
+# ---------------------------------------------------------------------------
+def test_extraction_program_is_fixed_and_params_are_data(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    manager = _manager_with_extract(_CANNED_EXTRACT)
+    _patch_manager(monkeypatch, manager)
+
+    _tool().execute_action("extract_document", input="art-1", persist=False)
+
+    program = manager.exec_programs[0]
+    # The executed program is the fixed template with only server-controlled path
+    # literals substituted — no input id, filename, or document content appears.
+    assert "json.load(open(" in program
+    assert "DocumentConverter" in program
+    assert "art-1" not in program
+    assert "statement.pdf" not in program
+    # Params reached the program as a JSON DATA file, not via interpolation.
+    params_files = [p for p in manager.put_files if p.endswith("params.json")]
+    assert params_files
+    params = json.loads(manager.put_files[params_files[0]].decode("utf-8"))
+    assert params["input_path"].endswith("statement.pdf")
+
+
+def test_malicious_param_value_is_not_executed(monkeypatch):
+    """A hostile filename ends up only in a DATA file the program json.loads, never in the program text."""
+    # The artifact's stored filename is attacker-controlled in the worst case;
+    # prove a code-shaped filename never lands in the executed program.
+    payload = "__import__('os').system('echo PWNED'); x = 'a.pdf"
+
+    class _Repo:
+        def __init__(self, conn):
+            pass
+
+        def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None):
+            return {"id": artifact_id, "current_version": 1, "title": payload}
+
+        def get_version(self, artifact_id, version):
+            return {"filename": payload, "storage_path": "inputs/u/artifacts/x/v1/f"}
+
+    class _Conn:
+        def __enter__(self):
+            return object()
+
+        def __exit__(self, *exc):
+            return False
+
+    monkeypatch.setattr(de, "db_readonly", lambda: _Conn())
+    monkeypatch.setattr(de, "ArtifactsRepository", _Repo)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+    manager = _manager_with_extract(_CANNED_EXTRACT)
+    _patch_manager(monkeypatch, manager)
+
+    out = _tool().execute_action("extract_document", input="art-1", persist=False)
+    assert out["status"] == "ok"
+
+    program = manager.exec_programs[0]
+    # The injection payload is absent from the program; it (sanitized) only rides
+    # in the params DATA file, so it is parsed as a string, never executed.
+    assert "PWNED" not in program
+    assert "__import__" not in program
+    params_files = [p for p in manager.put_files if p.endswith("params.json")]
+    params = json.loads(manager.put_files[params_files[0]].decode("utf-8"))
+    # The path is sanitized but even unsanitized it is pure data inside params.json.
+    assert "input_path" in params
+
+
+# ---------------------------------------------------------------------------
+# The fixed program template itself
+# ---------------------------------------------------------------------------
+def test_program_template_has_no_unbound_interpolation_points():
+    # Only the two server-controlled path placeholders may be formatted in; the
+    # template must format cleanly with just those, proving nothing else is a hole.
+    rendered = _EXTRACT_PROGRAM.format(params_path="extract/x/params.json", result_path="extract/x/result.json")
+    assert "params.json" in rendered
+    assert "result.json" in rendered
+    # Compiles as valid Python (no broken escapes from the doubled braces).
+    compile(rendered, "<extractor>", "exec")
+
+
+# ---------------------------------------------------------------------------
+# Markdown head+tail window (finding: the END of compliance docs must survive)
+# ---------------------------------------------------------------------------
+def _run_program_with_fake_docling(monkeypatch, markdown: str, md_cap: int) -> Dict[str, Any]:
+    """Exec the REAL extraction program with docling stubbed so we test the markdown windowing."""
+
+    class _FakeDoc:
+        tables: List[Any] = []
+        pages: Dict[str, Any] = {}
+
+        def export_to_markdown(self):
+            return markdown
+
+        def export_to_dict(self):
+            return {"texts": []}
+
+    class _FakeConverter:
+        def convert(self, src):
+            return types.SimpleNamespace(document=_FakeDoc())
+
+    fake_mod = types.ModuleType("docling")
+    fake_sub = types.ModuleType("docling.document_converter")
+    fake_sub.DocumentConverter = _FakeConverter
+    fake_mod.document_converter = fake_sub
+    monkeypatch.setitem(sys.modules, "docling", fake_mod)
+    monkeypatch.setitem(sys.modules, "docling.document_converter", fake_sub)
+
+    workdir = tempfile.mkdtemp()
+    params_path = os.path.join(workdir, "params.json")
+    result_path = os.path.join(workdir, "result.json")
+    with open(params_path, "w") as fh:
+        json.dump({"input_path": "x", "markdown_max_bytes": md_cap, "max_tables": 20}, fh)
+
+    program = _EXTRACT_PROGRAM.format(params_path=params_path, result_path=result_path)
+    ns: Dict[str, Any] = {}
+    exec(compile(program, "<extractor>", "exec"), ns, ns)  # noqa: S102
+    with open(result_path) as fh:
+        return json.load(fh)
+
+
+def test_markdown_window_keeps_head_and_tail_when_truncated(monkeypatch):
+    head_marker = "HEAD_START_OF_DOC"
+    tail_marker = "TAIL_END_OF_DOC_TOTALS"
+    body = "x" * (_MARKDOWN_MAX_BYTES * 2)
+    markdown = head_marker + body + tail_marker
+
+    result = _run_program_with_fake_docling(monkeypatch, markdown, _MARKDOWN_MAX_BYTES)
+
+    assert result["markdown_truncated"] is True
+    out = result["markdown"]
+    # Both the beginning AND the end of the document survive the byte budget.
+    assert head_marker in out
+    assert tail_marker in out
+    assert "...[truncated" in out
+    # The window respects the budget (plus the small truncation marker).
+    assert len(out) <= _MARKDOWN_MAX_BYTES + 64
+
+
+def test_markdown_full_content_when_under_cap(monkeypatch):
+    markdown = "# Short doc\n\nBalance: 1000\n\nSignature: Jane"
+    result = _run_program_with_fake_docling(monkeypatch, markdown, _MARKDOWN_MAX_BYTES)
+
+    assert result["markdown_truncated"] is False
+    assert result["markdown"] == markdown
+
+
+# ---------------------------------------------------------------------------
+# Per-table content cap (a single giant table must not bloat context)
+# ---------------------------------------------------------------------------
+def test_compact_payload_bounds_table_rows_and_cell_bytes(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_storage(monkeypatch)
+    _patch_no_persist(monkeypatch)
+
+    big_cell = "y" * (_MAX_CELL_CHARS * 3)
+    huge_table = {
+        "columns": ["a", "b"],
+        "rows": [[str(i), big_cell] for i in range(_MAX_TABLE_ROWS * 4)],
+    }
+    extract = dict(_CANNED_EXTRACT)
+    extract["tables"] = [huge_table]
+    _patch_manager(monkeypatch, _manager_with_extract(extract))
+
+    out = _tool().execute_action("extract_document", input="art-1", persist=False)
+    assert out["status"] == "ok"
+    table = out["structured"]["tables"][0]
+    # Rows are capped and flagged with the original total.
+    assert len(table["rows"]) == _MAX_TABLE_ROWS
+    assert table["rows_truncated"] is True
+    assert table["total_rows"] == _MAX_TABLE_ROWS * 4
+    # Long cell strings are truncated.
+    assert len(table["rows"][0][1]) <= _MAX_CELL_CHARS + len("...[truncated]")
+    assert table["rows"][0][1].endswith("...[truncated]")
+
+
+# ---------------------------------------------------------------------------
+# Input size cap (don't load an unbounded artifact into backend + kernel memory)
+# ---------------------------------------------------------------------------
+def test_oversized_input_is_rejected_before_exec(monkeypatch):
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    _patch_no_persist(monkeypatch)
+    monkeypatch.setattr(de.StorageCreator, "get_storage", staticmethod(lambda: _FakeStorage(b"P" * 4096)))
+    monkeypatch.setattr(de.settings, "SANDBOX_MAX_INPUT_BYTES", 1024, raising=False)
+    manager = _manager_with_extract(_CANNED_EXTRACT)
+    _patch_manager(monkeypatch, manager)
+
+    out = _tool().execute_action("extract_document", input="art-1", persist=False)
+    assert out["status"] == "error"
+    assert "too large" in out["error"]
+    # Nothing was staged or executed for an oversized input.
+    assert manager.exec_programs == []
+    assert manager.opened == []