From 9fdaea6cd2d66f0e8945ffda49978feab5f1940b Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 25 Jun 2026 11:35:29 +0100 Subject: [PATCH] Pass selected run-scoped documents to workflow agent nodes natively An agent node can list input_documents (short refs like A1, a literal "*" for all inputs, or state variables holding refs produced by an upstream node). The engine resolves each through the run-scoped artifact gate and feeds it to the node's LLM natively when the model supports the format (PDFs included via the provider's image synthesis), otherwise extracts it to text with Docling. A file_passing policy (auto/native/extract) and per-node count + per-file byte caps bound cost; the native decision uses the same supported-types source the provider handler uses, so a document is never silently dropped. Bytes never enter run state. --- .../agents/tools/document_extractor.py | 78 +++ application/agents/workflows/schemas.py | 9 + .../agents/workflows/workflow_engine.py | 187 +++++- application/core/settings.py | 4 + .../agents/test_workflow_node_attachments.py | 560 ++++++++++++++++++ 5 files changed, 833 insertions(+), 5 deletions(-) create mode 100644 tests/agents/test_workflow_node_attachments.py diff --git a/application/agents/tools/document_extractor.py b/application/agents/tools/document_extractor.py index a93f9104..d42f9ee1 100644 --- a/application/agents/tools/document_extractor.py +++ b/application/agents/tools/document_extractor.py @@ -109,6 +109,84 @@ _EXTRACT_PROGRAM = ( ) +def truncate_text_head_tail(text: str, max_bytes: Optional[int] = None) -> str: + """Bound text to a head+tail byte window so a large file can't flood context.""" + cap = int(max_bytes or _MARKDOWN_MAX_BYTES) + if cap <= 0: + return text + encoded = text.encode("utf-8") + if len(encoded) <= cap: + return text + head = cap // 2 + tail = cap - head + dropped = len(encoded) - head - tail + head_text = encoded[:head].decode("utf-8", errors="ignore") + tail_text = encoded[-tail:].decode("utf-8", errors="ignore") + return f"{head_text}\n\n...[truncated {dropped} bytes]...\n\n{tail_text}" + + +def extract_markdown_from_bytes( + data: bytes, + filename: str, + session_id: str, + *, + markdown_max_bytes: Optional[int] = None, +) -> Optional[str]: + """Run the fixed Docling program in a sandbox session and return the document's markdown, or None. + + The bytes ride in as a DATA file the fixed program reads; nothing untrusted is + interpolated into the program. Returns ``None`` when Docling is unavailable or + extraction fails (the caller decides how to degrade). + """ + safe_name = safe_filename(filename) or "document" + token = uuid.uuid4().hex + token_dir = f"extract/{token}" + input_path = f"{token_dir}/inputs/{safe_name}" + params_path = f"{token_dir}/params.json" + result_path = f"{token_dir}/result.json" + params = { + "input_path": input_path, + "markdown_max_bytes": int(markdown_max_bytes or _MARKDOWN_MAX_BYTES), + "max_tables": 0, + } + 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: + logger.exception("extract_markdown_from_bytes: failed to open sandbox session") + return None + try: + manager.put_file(session_id, input_path, 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: + return None + raw = manager.get_file(session_id, result_path) + except Exception: + logger.exception("extract_markdown_from_bytes: extraction failed") + return None + finally: + manager.remove_path(session_id, token_dir) + try: + manager.close(session_id) + except Exception: + logger.exception("extract_markdown_from_bytes: session close failed") + + if not raw: + return None + try: + extracted = json.loads(raw.decode("utf-8")) + except (ValueError, UnicodeDecodeError): + return None + if not isinstance(extracted, dict) or extracted.get("error") or not extracted.get("ok"): + return None + markdown = extracted.get("markdown") + return markdown if isinstance(markdown, str) else None + + class DocumentExtractorTool(Tool): """Document Extractor Convert an input document artifact (pdf/docx/pptx/...) to compact, schema-validated structured JSON via Docling. diff --git a/application/agents/workflows/schemas.py b/application/agents/workflows/schemas.py index d5dda148..4ccd1762 100644 --- a/application/agents/workflows/schemas.py +++ b/application/agents/workflows/schemas.py @@ -49,6 +49,15 @@ class AgentNodeConfig(BaseModel): retriever: str = "" model_id: Optional[str] = None json_schema: Optional[Dict[str, Any]] = None + # Run-scoped documents fed to this node's LLM. Entries are state-var names + # holding artifact refs (single dict or a list of dicts), raw artifact ids, + # short refs (``A1``), or the ``"*"``/``"input_documents"`` token meaning + # "every ref in ``state['input_documents']``". + input_documents: List[str] = Field(default_factory=list) + # How selected documents reach the model: ``auto`` (native when the model + # accepts the mime, else extract to text), ``native`` (force native; raise + # on an unsupported mime), or ``extract`` (always inline extracted text). + file_passing: Literal["auto", "native", "extract"] = "auto" class CodeNodeConfig(BaseModel): diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index fce3a70a..a2008085 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -250,10 +250,10 @@ class WorkflowEngine: ) node_api_key = get_api_key_for_provider(node_llm_name) or self.agent.api_key + # Structured output gates on the model's registry capability flags; + # fetch them only when a node json_schema needs the check. if node_json_schema and node_model_id: - model_capabilities = get_model_capabilities( - node_model_id, user_id=node_user_id - ) + model_capabilities = get_model_capabilities(node_model_id, user_id=node_user_id) if model_capabilities and not model_capabilities.get( "supports_structured_output", False ): @@ -294,6 +294,16 @@ class WorkflowEngine: if getattr(node_agent, "tool_executor", None) is not None: node_agent.tool_executor.workflow_run_id = self.workflow_run_id + # Decide native-eligibility from the SAME supported-types list the provider + # handler filters on at send time, so a mime is never sent native-but-empty + # and then silently dropped. Read post-construction: BaseAgent consumes + # ``self.attachments`` at gen time, so assigning here takes effect. + node_attachments = self._materialize_node_attachments( + node_config, node.title, self._agent_supported_attachment_types(node_agent) + ) + if node_attachments: + node_agent.attachments = node_attachments + full_response_parts: List[str] = [] structured_response_parts: List[str] = [] has_structured_response = False @@ -477,16 +487,183 @@ class WorkflowEngine: loaded.append(f"inputs/{filename}") return loaded + def _materialize_node_attachments( + self, + node_config: AgentNodeConfig, + node_title: str, + supported_types: List[str], + ) -> List[Dict[str, Any]]: + """Resolve a node's selected documents to native/extracted attachment dicts for its LLM.""" + from application.agents.tools.artifact_ref import resolve_artifact_id + from application.core.settings import settings + from application.storage.db.repositories.artifacts import ArtifactsRepository + from application.storage.db.session import db_readonly + + raw_ids = self._resolve_input_artifact_ids(node_config.input_documents) + if not raw_ids: + return [] + + supported = set(supported_types) + supports_images = any(t.startswith("image/") for t in supported) + max_files = int(getattr(settings, "WORKFLOW_NODE_NATIVE_MAX_FILES", 5)) + max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024)) + + # One read-only connection for the whole batch; the resolved-version + # rows are collected, then storage reads happen outside the DB context. + resolved: List[tuple] = [] + with db_readonly() as conn: + repo = ArtifactsRepository(conn) + for raw in raw_ids: + # A short ref (A1/...) resolves to an id within this run only; the + # resolved id is re-checked through the run-scoped gate so a forged + # or cross-run ref can never reach another tenant's bytes. + artifact_id = resolve_artifact_id(repo, raw, workflow_run_id=self.workflow_run_id) + artifact = ( + repo.get_artifact_in_parent(artifact_id, workflow_run_id=self.workflow_run_id) + if artifact_id is not None + else None + ) + if artifact is None: + raise ValueError( + f'Document "{raw}" for node "{node_title}" was not found in this run.' + ) + version = repo.get_version(artifact_id, artifact["current_version"]) + if not version or not version.get("storage_path"): + raise ValueError(f"input document {artifact_id} has no stored content.") + resolved.append((str(artifact_id), artifact, version)) + + attachments: List[Dict[str, Any]] = [] + native_count = 0 + for artifact_id, artifact, version in resolved: + storage_path = version["storage_path"] + mime_type = version.get("mime_type") or "application/octet-stream" + filename = version.get("filename") or artifact.get("title") or artifact_id + size = version.get("size") + if isinstance(size, int) and size > max_bytes: + logger.warning( + "Workflow node %s: document %s (%d bytes) exceeds the %d-byte cap; skipping", + node_title, artifact_id, size, max_bytes, + ) + continue + + native_ok = self._is_native_mime(mime_type, supported, supports_images) + policy = node_config.file_passing + if policy == "native" and not native_ok: + raise ValueError( + f'Model "{node_config.model_id or self.agent.model_id}" cannot read ' + f'"{mime_type}" natively for node "{node_title}".' + ) + go_native = policy == "native" or (policy == "auto" and native_ok) + if go_native and native_count >= max_files: + logger.warning( + "Workflow node %s: native file cap (%d) reached; extracting %s instead", + node_title, max_files, filename, + ) + go_native = False + + if go_native: + # No bytes copied: the provider reads them from storage via ``path``. + attachments.append({"id": artifact_id, "mime_type": mime_type, "path": storage_path}) + native_count += 1 + else: + content = self._extract_attachment_text(storage_path, mime_type, filename, max_bytes) + if content is None: + logger.warning( + "Workflow node %s: could not extract text from %s; skipping", + node_title, filename, + ) + continue + # A non-native mime routes this through ``_append_unsupported_attachments``, + # which inlines ``content`` as text in the system prompt. + attachments.append( + {"id": artifact_id, "mime_type": "text/plain", "content": content} + ) + return attachments + + @staticmethod + def _agent_supported_attachment_types(node_agent: "BaseAgent") -> List[str]: + """Return the provider's authoritative supported attachment mime types (handler's source).""" + llm = getattr(node_agent, "llm", None) + getter = getattr(llm, "get_supported_attachment_types", None) + if not callable(getter): + return [] + types = getter() + return list(types) if isinstance(types, (list, tuple, set)) else [] + + @staticmethod + def _is_native_mime(mime_type: str, supported_types: set, supports_images: bool) -> bool: + """A mime is native if the model accepts it, or it is a PDF a vision model renders to images.""" + if mime_type in supported_types: + return True + return mime_type == "application/pdf" and supports_images + + def _extract_attachment_text( + self, storage_path: str, mime_type: str, filename: str, max_bytes: int + ) -> Optional[str]: + """Get an attachment's text: inline already-text formats, else extract via Docling; None on failure.""" + from application.agents.tools.document_extractor import ( + extract_markdown_from_bytes, + truncate_text_head_tail, + ) + from application.storage.storage_creator import StorageCreator + + try: + data = StorageCreator.get_storage().get_file(storage_path).read() + except Exception: + logger.exception("Workflow node: failed to read document bytes for extraction") + return None + # Defensive size gate: a NULL/missing version size skips the pre-read cap, + # so re-check the actual bytes before inlining or sandbox-extracting. + if len(data) > max_bytes: + logger.warning( + "Workflow node: document at %s (%d bytes) exceeds the %d-byte cap; skipping", + storage_path, len(data), max_bytes, + ) + return None + if self._is_inline_text_mime(mime_type): + try: + text = data.decode("utf-8", errors="replace") + except Exception: + return None + # Bound the inlined text to a head+tail window so a large-but-under-cap + # text file can't blow the context (the Docling branch is already bounded). + return truncate_text_head_tail(text) + return extract_markdown_from_bytes(data, filename, self._session_id()) + + @staticmethod + def _is_inline_text_mime(mime_type: str) -> bool: + """Already-text formats are inlined directly (no Docling round-trip).""" + if mime_type.startswith("text/"): + return True + return mime_type in ("application/json", "application/xml") + def _resolve_input_artifact_ids(self, inputs: List[str]) -> List[str]: - """Resolve node ``inputs`` (state-var names holding refs, or raw ids) to artifact ids.""" + """Resolve node ``inputs`` (refs/ids, ``*`` token, or state vars holding a ref or a list of refs).""" resolved: List[str] = [] for raw in inputs or []: + # ``*`` / ``input_documents`` expands to every run input document. + if isinstance(raw, str) and raw.strip() in ("*", "input_documents"): + resolved.extend(self._input_document_ids()) + continue ref = self.state.get(raw) if isinstance(raw, str) else None if isinstance(ref, dict) and ref.get("artifact_id"): resolved.append(str(ref["artifact_id"])) + elif isinstance(ref, list): + # A state var holding a list of ref dicts (e.g. input_documents). + for item in ref: + if isinstance(item, dict) and item.get("artifact_id"): + resolved.append(str(item["artifact_id"])) elif isinstance(raw, str) and raw.strip(): resolved.append(raw.strip()) - return resolved + # Dedup preserving order so ``["*", "A1"]`` / duplicate refs don't attach twice. + return list(dict.fromkeys(resolved)) + + def _input_document_ids(self) -> List[str]: + """Return the artifact ids of every ref in ``state['input_documents']``.""" + docs = self.state.get("input_documents") + if not isinstance(docs, list): + return [] + return [str(d["artifact_id"]) for d in docs if isinstance(d, dict) and d.get("artifact_id")] def _session_id(self) -> str: """Sanitize the run id into the sandbox-gateway charset for the session key.""" diff --git a/application/core/settings.py b/application/core/settings.py index 2b264d71..e64cf532 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -343,6 +343,10 @@ class Settings(BaseSettings): 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 + # Per-agent-node cap on files passed natively to the node's LLM (vision/doc + # inputs). Files past the cap are extracted to text or dropped, not attached + # natively, to bound context/cost. Re-uses SANDBOX_MAX_INPUT_BYTES per file. + WORKFLOW_NODE_NATIVE_MAX_FILES: int = 5 # 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. diff --git a/tests/agents/test_workflow_node_attachments.py b/tests/agents/test_workflow_node_attachments.py new file mode 100644 index 00000000..3c2e4731 --- /dev/null +++ b/tests/agents/test_workflow_node_attachments.py @@ -0,0 +1,560 @@ +"""Unit tests for native/extracted document passing to workflow AGENT nodes. + +These exercise the engine's ``_materialize_node_attachments`` policy and the +``_execute_agent_node`` wiring with a fake artifacts repo (no live DB / storage / +sandbox / LLM). The run-scoped authz gate, the auto/native/extract policy matrix, +the caps, and the no-selection regression path are all covered here. +""" + +from __future__ import annotations + +import contextlib +import uuid +from types import SimpleNamespace +from typing import Any, Dict + +import pytest + +from application.agents.workflows.node_agent import WorkflowNodeAgentFactory +from application.agents.workflows.schemas import ( + NodeType, + Workflow, + WorkflowGraph, + WorkflowNode, +) +from application.agents.workflows.workflow_engine import WorkflowEngine +from application.storage.db.repositories.artifacts import ArtifactsRepository + +RUN_ID = "11111111-1111-1111-1111-111111111111" + + +# --------------------------------------------------------------------------- +# Fixtures / fakes +# --------------------------------------------------------------------------- + + +def _patch_repo(monkeypatch, artifacts: Dict[str, Dict[str, Any]], run_id: str = RUN_ID) -> None: + """Patch db_readonly + the three repo methods the engine calls so it reads the fake data. + + The repo CLASS is left intact (transitive ``from ... import ArtifactsRepository`` + bindings are untouched); only the methods the resolver uses are redirected, which + monkeypatch reverts cleanly on teardown. + """ + # 1-based position order = insertion order of ids belonging to this run. + positions = [aid for aid, a in artifacts.items() if a.get("run_id") == run_id] + + def _at_position(self, n, *, conversation_id=None, workflow_run_id=None): + if workflow_run_id != run_id or not isinstance(n, int) or n < 1 or n > len(positions): + return None + return positions[n - 1] + + def _in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None): + art = artifacts.get(artifact_id) + if art is None or art.get("run_id") != workflow_run_id: + return None + return {"current_version": art["current_version"], "title": art.get("title")} + + def _version(self, artifact_id, version): + art = artifacts.get(artifact_id) + return art["versions"].get(version) if art is not None else None + + @contextlib.contextmanager + def _fake_db_readonly(): + yield object() + + monkeypatch.setattr("application.storage.db.session.db_readonly", _fake_db_readonly) + monkeypatch.setattr(ArtifactsRepository, "__init__", lambda self, conn=None: None) + monkeypatch.setattr(ArtifactsRepository, "artifact_id_at_position", _at_position) + monkeypatch.setattr(ArtifactsRepository, "get_artifact_in_parent", _in_parent) + monkeypatch.setattr(ArtifactsRepository, "get_version", _version) + + +def _artifact(run_id, mime, *, filename="doc", size=10, version=1, storage_path=None): + """Build a fake artifact record with a single current version.""" + aid = str(uuid.uuid4()) + return aid, { + "run_id": run_id, + "current_version": version, + "title": filename, + "versions": { + version: { + "storage_path": storage_path or f"store/{aid}", + "mime_type": mime, + "filename": filename, + "size": size, + } + }, + } + + +def _engine(monkeypatch, run_id: str = RUN_ID) -> WorkflowEngine: + graph = WorkflowGraph(workflow=Workflow(name="Attach Test"), nodes=[], edges=[]) + agent = SimpleNamespace( + endpoint="stream", + llm_name="openai", + model_id="gpt-4o-mini", + api_key="test-key", + chat_history=[], + user="user-1", + decoded_token={"sub": "user-1"}, + ) + return WorkflowEngine(graph, agent, workflow_run_id=run_id) + + +# Provider supported-types lists, mirroring what ``llm.get_supported_attachment_types()`` +# returns -- the engine now decides native-eligibility from this same source. +VISION_TYPES = ["image/png", "image/jpeg"] +TEXT_TYPES: list = [] + + +# --------------------------------------------------------------------------- +# _resolve_input_artifact_ids: tokens, lists, refs +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_star_token_expands_to_all_input_documents(monkeypatch): + """``*`` expands to every ref in state['input_documents'].""" + eng = _engine(monkeypatch) + eng.state["input_documents"] = [ + {"artifact_id": "aaa"}, {"artifact_id": "bbb"}, {"not": "a ref"}, + ] + assert eng._resolve_input_artifact_ids(["*"]) == ["aaa", "bbb"] + assert eng._resolve_input_artifact_ids(["input_documents"]) == ["aaa", "bbb"] + + +@pytest.mark.unit +def test_state_var_single_ref_and_list_of_refs(monkeypatch): + """A state var may hold a single ref dict or a list of ref dicts.""" + eng = _engine(monkeypatch) + eng.state["one"] = {"artifact_id": "x1"} + eng.state["many"] = [{"artifact_id": "y1"}, {"artifact_id": "y2"}] + assert eng._resolve_input_artifact_ids(["one"]) == ["x1"] + assert eng._resolve_input_artifact_ids(["many"]) == ["y1", "y2"] + + +@pytest.mark.unit +def test_raw_id_and_short_ref_pass_through(monkeypatch): + """A raw id or short ref that is not a state key passes through verbatim.""" + eng = _engine(monkeypatch) + assert eng._resolve_input_artifact_ids(["A1", "abc-id"]) == ["A1", "abc-id"] + + +# --------------------------------------------------------------------------- +# _materialize_node_attachments: policy matrix +# --------------------------------------------------------------------------- + + +def _node_config(**kwargs): + from application.agents.workflows.schemas import AgentNodeConfig + + return AgentNodeConfig(**kwargs) + + +@pytest.mark.unit +def test_auto_image_goes_native_for_vision_model(monkeypatch): + """auto + image mime + vision model -> a native attachment {id, mime_type, path}.""" + aid, rec = _artifact(RUN_ID, "image/png", filename="chart.png") + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + cfg = _node_config(input_documents=[aid], file_passing="auto") + + out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES) + + assert out == [{"id": aid, "mime_type": "image/png", "path": rec["versions"][1]["storage_path"]}] + + +@pytest.mark.unit +def test_auto_pdf_native_when_model_supports_images(monkeypatch): + """auto + PDF + vision (image) model -> native (synthetic PDF-to-image path).""" + aid, rec = _artifact(RUN_ID, "application/pdf", filename="r.pdf") + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + cfg = _node_config(input_documents=[aid], file_passing="auto") + + out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES) + + assert out[0]["mime_type"] == "application/pdf" + assert "path" in out[0] and "content" not in out[0] + + +@pytest.mark.unit +def test_auto_text_only_model_extracts_to_content(monkeypatch): + """auto + text-only model -> extracted/inlined text content, non-native mime.""" + aid, rec = _artifact(RUN_ID, "text/plain", filename="notes.txt") + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + staticmethod(lambda: _FakeStorage(b"hello world")), + ) + cfg = _node_config(input_documents=[aid], file_passing="auto") + + out = eng._materialize_node_attachments(cfg, "Node", TEXT_TYPES) + + assert out == [{"id": aid, "mime_type": "text/plain", "content": "hello world"}] + + +@pytest.mark.unit +def test_extract_always_inlines_text_even_for_vision_model(monkeypatch): + """file_passing=extract inlines text even when the model could take it natively.""" + aid, rec = _artifact(RUN_ID, "text/markdown", filename="r.md") + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + staticmethod(lambda: _FakeStorage(b"# Title")), + ) + cfg = _node_config(input_documents=[aid], file_passing="extract") + + out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES) + + assert out == [{"id": aid, "mime_type": "text/plain", "content": "# Title"}] + + +@pytest.mark.unit +def test_native_on_unsupported_mime_raises(monkeypatch): + """file_passing=native on a mime the model can't read raises a clear error.""" + aid, rec = _artifact(RUN_ID, "application/pdf", filename="r.pdf") + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + cfg = _node_config(input_documents=[aid], file_passing="native", model_id="gpt-4o-mini") + + with pytest.raises(ValueError, match="cannot read"): + eng._materialize_node_attachments(cfg, "MyNode", TEXT_TYPES) + + +@pytest.mark.unit +def test_extract_non_text_uses_docling(monkeypatch): + """A non-text mime under extract routes through the Docling helper.""" + aid, rec = _artifact(RUN_ID, "application/vnd.openxmlformats", filename="r.docx") + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + staticmethod(lambda: _FakeStorage(b"\x00binary")), + ) + monkeypatch.setattr( + "application.agents.tools.document_extractor.extract_markdown_from_bytes", + lambda data, filename, session_id, **kw: "EXTRACTED MD", + ) + cfg = _node_config(input_documents=[aid], file_passing="extract") + + out = eng._materialize_node_attachments(cfg, "Node", TEXT_TYPES) + + assert out == [{"id": aid, "mime_type": "text/plain", "content": "EXTRACTED MD"}] + + +# --------------------------------------------------------------------------- +# Security: cross-run / forged refs are rejected, never attached +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_cross_run_artifact_is_rejected(monkeypatch): + """An artifact belonging to a different run yields None at the gate -> ValueError.""" + aid, rec = _artifact("99999999-9999-9999-9999-999999999999", "image/png") + _patch_repo(monkeypatch, {aid: rec}) # run scope is RUN_ID; artifact is in another run + eng = _engine(monkeypatch) + cfg = _node_config(input_documents=[aid], file_passing="auto") + + with pytest.raises(ValueError, match="not found in this run"): + eng._materialize_node_attachments(cfg, "Node", VISION_TYPES) + + +@pytest.mark.unit +def test_forged_uuid_is_rejected(monkeypatch): + """A forged uuid never matches the run-scoped gate -> ValueError, nothing attached.""" + aid, rec = _artifact(RUN_ID, "image/png") + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + forged = str(uuid.uuid4()) + cfg = _node_config(input_documents=[forged], file_passing="auto") + + with pytest.raises(ValueError, match="not found in this run"): + eng._materialize_node_attachments(cfg, "Node", VISION_TYPES) + + +# --------------------------------------------------------------------------- +# Caps +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_native_file_cap_bounds_native_then_extracts(monkeypatch): + """More than the native cap: the first N go native, the rest are extracted.""" + monkeypatch.setattr( + "application.core.settings.settings.WORKFLOW_NODE_NATIVE_MAX_FILES", 2, raising=False + ) + artifacts = {} + ids = [] + for i in range(4): + aid, rec = _artifact(RUN_ID, "image/png", filename=f"i{i}.png") + artifacts[aid] = rec + ids.append(aid) + _patch_repo(monkeypatch, artifacts) + eng = _engine(monkeypatch) + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + staticmethod(lambda: _FakeStorage(b"img-bytes")), + ) + # extract of a non-text image falls back to Docling; stub it so it returns text. + monkeypatch.setattr( + "application.agents.tools.document_extractor.extract_markdown_from_bytes", + lambda data, filename, session_id, **kw: "fallback text", + ) + cfg = _node_config(input_documents=ids, file_passing="auto") + + out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES) + + native = [a for a in out if "path" in a] + extracted = [a for a in out if "content" in a] + assert len(native) == 2 + assert len(extracted) == 2 + + +@pytest.mark.unit +def test_oversize_file_is_skipped(monkeypatch): + """A file past the per-file byte ceiling is dropped, not attached.""" + monkeypatch.setattr( + "application.core.settings.settings.SANDBOX_MAX_INPUT_BYTES", 5, raising=False + ) + aid, rec = _artifact(RUN_ID, "image/png", size=999) + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + cfg = _node_config(input_documents=[aid], file_passing="auto") + + out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES) + + assert out == [] + + +@pytest.mark.unit +def test_oversize_text_skipped_by_post_read_guard_when_size_missing(monkeypatch): + """A NULL/missing version size skips the pre-read cap; the post-read byte guard still drops it.""" + monkeypatch.setattr( + "application.core.settings.settings.SANDBOX_MAX_INPUT_BYTES", 5, raising=False + ) + # size=None bypasses the ``isinstance(size, int)`` pre-read check; the actual + # bytes (longer than the 5-byte cap) must be rejected after reading. + aid, rec = _artifact(RUN_ID, "text/plain", filename="big.txt", size=None) + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + staticmethod(lambda: _FakeStorage(b"way over the cap")), + ) + cfg = _node_config(input_documents=[aid], file_passing="auto") + + out = eng._materialize_node_attachments(cfg, "Node", TEXT_TYPES) + + assert out == [] + + +@pytest.mark.unit +def test_large_under_cap_text_is_windowed_not_inlined_whole(monkeypatch): + """A large-but-under-cap text file is bounded to a head+tail window, not inlined whole.""" + from application.agents.tools.document_extractor import _MARKDOWN_MAX_BYTES + + big_text = ("A" * (_MARKDOWN_MAX_BYTES * 3)).encode("utf-8") + aid, rec = _artifact(RUN_ID, "text/plain", filename="notes.txt", size=len(big_text)) + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + staticmethod(lambda: _FakeStorage(big_text)), + ) + cfg = _node_config(input_documents=[aid], file_passing="auto") + + out = eng._materialize_node_attachments(cfg, "Node", TEXT_TYPES) + + assert len(out) == 1 + content = out[0]["content"] + assert "...[truncated" in content + # Bounded to roughly the window, far below the whole 3x payload. + assert len(content.encode("utf-8")) < _MARKDOWN_MAX_BYTES + 200 + + +@pytest.mark.unit +def test_duplicate_refs_attach_once(monkeypatch): + """A duplicated/expanded ref resolves once so the same artifact is not attached twice.""" + aid, rec = _artifact(RUN_ID, "image/png", filename="dup.png") + _patch_repo(monkeypatch, {aid: rec}) + eng = _engine(monkeypatch) + eng.state["input_documents"] = [{"artifact_id": aid}] + cfg = _node_config(input_documents=["*", aid], file_passing="auto") + + out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES) + + assert out == [{"id": aid, "mime_type": "image/png", "path": rec["versions"][1]["storage_path"]}] + + +# --------------------------------------------------------------------------- +# _execute_agent_node wiring +# --------------------------------------------------------------------------- + + +class _StubLLM: + """Minimal LLM stub exposing the provider's supported-types list (the native source).""" + + def __init__(self, supported_types): + self._supported_types = supported_types + + def get_supported_attachment_types(self): + return self._supported_types + + +class _StubNodeAgent: + """Captures the factory kwargs and yields a fixed answer (no LLM call).""" + + def __init__(self, *, supported_types=None, **kwargs): + self.kwargs = kwargs + self.tool_executor = None + self.attachments = [] + self.llm = _StubLLM(supported_types if supported_types is not None else []) + + def gen(self, _prompt): + yield {"answer": "ok"} + + +def _agent_node(input_documents=None, file_passing="auto", node_id="agent_1") -> WorkflowNode: + config = { + "agent_type": "classic", + "system_prompt": "You are a helpful assistant.", + "prompt_template": "", + "stream_to_user": False, + "tools": [], + } + if input_documents is not None: + config["input_documents"] = input_documents + config["file_passing"] = file_passing + return WorkflowNode( + id=node_id, + workflow_id="wf-1", + type=NodeType.AGENT, + title="Agent", + position={"x": 0, "y": 0}, + config=config, + ) + + +def _patch_capabilities(monkeypatch): + """Stub provider/api-key resolution (capabilities are only fetched for json_schema nodes).""" + monkeypatch.setattr( + "application.core.model_utils.get_model_capabilities", + lambda model_id, user_id=None: None, + ) + monkeypatch.setattr( + "application.core.model_utils.get_provider_from_model_id", + lambda model_id, user_id=None: "openai", + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", lambda _p: "k" + ) + + +@pytest.mark.unit +def test_execute_agent_node_assigns_native_attachment_from_provider_types(monkeypatch): + """A node with input_documents assigns a native attachment, decided from the provider's types.""" + aid, rec = _artifact(RUN_ID, "image/png", filename="c.png") + _patch_repo(monkeypatch, {aid: rec}) + _patch_capabilities(monkeypatch) + eng = _engine(monkeypatch) + + captured: dict = {} + created: list = [] + + def _create(**kwargs): + captured.update(kwargs) + agent = _StubNodeAgent(supported_types=VISION_TYPES, **kwargs) + created.append(agent) + return agent + + monkeypatch.setattr(WorkflowNodeAgentFactory, "create", staticmethod(_create)) + + node = _agent_node(input_documents=[aid], file_passing="auto") + list(eng._execute_agent_node(node)) + + # Attachments are assigned post-construction (BaseAgent reads them at gen time), + # not passed through factory_kwargs. + assert "attachments" not in captured + assert created[0].attachments == [ + {"id": aid, "mime_type": "image/png", "path": rec["versions"][1]["storage_path"]} + ] + + +@pytest.mark.unit +def test_execute_agent_node_native_decision_tracks_provider_types(monkeypatch): + """A mime the provider does NOT support falls back to extracted text under ``auto``.""" + aid, rec = _artifact(RUN_ID, "image/png", filename="c.png") + _patch_repo(monkeypatch, {aid: rec}) + _patch_capabilities(monkeypatch) + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + staticmethod(lambda: _FakeStorage(b"image-bytes")), + ) + monkeypatch.setattr( + "application.agents.tools.document_extractor.extract_markdown_from_bytes", + lambda data, filename, session_id, **kw: "EXTRACTED", + ) + eng = _engine(monkeypatch) + + created: list = [] + + def _create(**kwargs): + # Provider reports NO native types, so the registry-agnostic decision + # must extract rather than send a native-but-empty attachment. + agent = _StubNodeAgent(supported_types=[], **kwargs) + created.append(agent) + return agent + + monkeypatch.setattr(WorkflowNodeAgentFactory, "create", staticmethod(_create)) + + node = _agent_node(input_documents=[aid], file_passing="auto") + list(eng._execute_agent_node(node)) + + assert created[0].attachments == [ + {"id": aid, "mime_type": "text/plain", "content": "EXTRACTED"} + ] + + +@pytest.mark.unit +def test_execute_agent_node_without_documents_has_no_attachments(monkeypatch): + """A node without input_documents leaves attachments empty (no regression).""" + _patch_capabilities(monkeypatch) + eng = _engine(monkeypatch) + + captured: dict = {} + created: list = [] + + def _create(**kwargs): + captured.update(kwargs) + agent = _StubNodeAgent(supported_types=VISION_TYPES, **kwargs) + created.append(agent) + return agent + + monkeypatch.setattr(WorkflowNodeAgentFactory, "create", staticmethod(_create)) + + node = _agent_node(input_documents=None) + list(eng._execute_agent_node(node)) + + assert "attachments" not in captured + assert created[0].attachments == [] + + +class _FakeStorage: + """Minimal storage stub: get_file(path).read() -> fixed bytes.""" + + def __init__(self, data: bytes) -> None: + self._data = data + + def get_file(self, _path): + return _FakeFile(self._data) + + +class _FakeFile: + def __init__(self, data: bytes) -> None: + self._data = data + + def read(self) -> bytes: + return self._data