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