From 0d67734fdd42a8aed555cca118cd797b5e4ac548 Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 25 Jun 2026 17:11:45 +0100 Subject: [PATCH] Let sandbox tools use a chat attachment, bridged on first reference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When code_executor or read_document is given an input that doesn't resolve to an existing artifact, it now falls back to the caller's own chat attachments, bridges the referenced one into a conversation-scoped artifact (idempotent, server-computed size/sha256, quota-respected), and stages it. This is lazy — plain chat uploads are never bridged, only a file a tool actually references — and conversation-scoped only (workflow nodes bridge attachments up front). The match is confined to the request's own attachments and re-verified against the user-scoped attachments repo, so a model-supplied name/id can only ever reach the caller's own files. Only the current request's attachments are reachable (the attachments table has no conversation column); prior-turn files are a follow-up. --- application/agents/base.py | 4 + application/agents/tool_executor.py | 19 +- application/agents/tools/attachment_bridge.py | 132 +++++ application/agents/tools/code_executor.py | 40 +- application/agents/tools/read_document.py | 31 +- .../storage/db/repositories/artifacts.py | 21 + tests/agents/tools/test_attachment_bridge.py | 463 ++++++++++++++++++ 7 files changed, 703 insertions(+), 7 deletions(-) create mode 100644 application/agents/tools/attachment_bridge.py create mode 100644 tests/agents/tools/test_attachment_bridge.py diff --git a/application/agents/base.py b/application/agents/base.py index 8001f53a..0f6ea2b0 100644 --- a/application/agents/base.py +++ b/application/agents/base.py @@ -340,6 +340,10 @@ class BaseAgent(ABC): self.tools = self.tool_executor.prepare_tools_for_llm(tools_dict) def _execute_tool_action(self, tools_dict, call): + # Mirror the request's attachments onto the executor so sandbox tools + # can lazily bridge a referenced chat attachment to a conversation + # artifact; only the caller's own (user-scoped) attachments are passed. + self.tool_executor.attachments = self.attachments return self.tool_executor.execute( tools_dict, call, self.llm.__class__.__name__ ) diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py index d168ba9a..e97a7699 100644 --- a/application/agents/tool_executor.py +++ b/application/agents/tool_executor.py @@ -202,6 +202,10 @@ class ToolExecutor: # workflow run rather than a conversation. self.workflow_run_id: Optional[str] = None self.message_id: Optional[str] = None + # The request's own (already user-scoped) chat attachments, stamped onto + # sandbox tools so a referenced attachment can be lazily bridged to a + # conversation-scoped artifact at tool-use time. + self.attachments: List[Dict] = [] self.client_tools: Optional[List[Dict]] = None self._name_to_tool: Dict[str, Tuple[str, str]] = {} self._tool_to_name: Dict[Tuple[str, str], str] = {} @@ -885,7 +889,15 @@ class ToolExecutor: """Load a tool, using cache when possible.""" cache_key = f"{tool_data['name']}:{tool_id}:{self.user or ''}" if cache_key in self._loaded_tools: - return self._loaded_tools[cache_key] + cached = self._loaded_tools[cache_key] + # A tool cached on an earlier turn carries that turn's attachments; + # refresh them so a chat attachment added this turn is bridgeable. + cached_config = getattr(cached, "config", None) + if isinstance(cached_config, dict) and self.conversation_id: + # Refresh unconditionally so a turn with no attachments clears the + # prior turn's list (no stale carryover within the session). + cached_config["attachments"] = self.attachments or [] + return cached tm = ToolManager(config={}) @@ -949,6 +961,11 @@ class ToolExecutor: tool_config["tool_id"] = str(row_id) if self.conversation_id: tool_config["conversation_id"] = self.conversation_id + # Carry the request's own attachments so sandbox tools can + # lazily bridge a referenced chat attachment (conversation + # scope only; workflow nodes bridge attachments up front). + if self.attachments: + tool_config["attachments"] = self.attachments # Workflow agent nodes run-scope their artifact tools so a short # ref (A1) and edit_artifact resolve against the workflow run. if self.workflow_run_id: diff --git a/application/agents/tools/attachment_bridge.py b/application/agents/tools/attachment_bridge.py new file mode 100644 index 00000000..c3887cf8 --- /dev/null +++ b/application/agents/tools/attachment_bridge.py @@ -0,0 +1,132 @@ +"""Lazily bridge a chat attachment into a conversation-scoped artifact when a tool references it. + +A chat attachment lives in the ``attachments`` table (parsed to text for the LLM context); it is +not an artifact and so cannot be fed to ``code_executor`` / ``read_document`` directly. When one of +those tools references an attachment by id or filename, this module materializes it into a +conversation-scoped artifact on demand — only the request's own (already user-scoped) attachments are +reachable, and an already-bridged attachment is reused so repeated references never burn extra quota. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +from application.sandbox.artifacts_capture import QuotaExceeded, persist_new_artifact +from application.storage.db.repositories.artifacts import ArtifactsRepository +from application.storage.db.repositories.attachments import AttachmentsRepository +from application.storage.db.session import db_readonly +from application.storage.storage_creator import StorageCreator + +logger = logging.getLogger(__name__) + + +class AttachmentBridgeError(Exception): + """Raised when a matched attachment cannot be bridged (e.g. quota, unreadable bytes).""" + + +def _normalize_name(value: Any) -> str: + """Lowercase + strip a filename for tolerant matching.""" + return str(value or "").strip().lower() + + +def match_attachment( + attachments: Optional[List[Dict[str, Any]]], raw_ref: str, user_id: str +) -> Optional[Dict[str, Any]]: + """Match a model-supplied id/filename against the caller's OWN request attachments; None otherwise. + + Matching is confined to ``attachments`` (already user-scoped when loaded) so a forged id/name can + never reach another user's or conversation's attachment. An id match is re-verified against + ``AttachmentsRepository.get_any(id, user_id)`` so only the owner's row is ever bridged. When two + attachments share a filename the first is chosen; reference by id to disambiguate. + """ + if not attachments or not raw_ref: + return None + ref = raw_ref.strip() + if not ref: + return None + ref_norm = _normalize_name(ref) + by_filename: Optional[Dict[str, Any]] = None + for attachment in attachments: + if not isinstance(attachment, dict): + continue + ids = { + str(attachment.get(key)) + for key in ("id", "_id", "legacy_mongo_id") + if attachment.get(key) is not None + } + if ref in ids: + return _verify_owner(attachment, user_id) + filename = attachment.get("filename") + if by_filename is None and filename and _normalize_name(filename) == ref_norm: + by_filename = attachment + if by_filename is not None: + return _verify_owner(by_filename, user_id) + return None + + +def _verify_owner(attachment: Dict[str, Any], user_id: str) -> Optional[Dict[str, Any]]: + """Re-confirm the attachment belongs to ``user_id`` via the user-scoped repo; in-memory dict on hit.""" + attachment_id = attachment.get("id") or attachment.get("_id") or attachment.get("legacy_mongo_id") + if attachment_id is None: + return None + try: + with db_readonly() as conn: + owned = AttachmentsRepository(conn).get_any(str(attachment_id), user_id) + except Exception: + logger.exception("attachment_bridge: ownership re-check failed") + return None + # Prefer the DB row (authoritative upload_path/mime) but only when it confirms ownership. + return owned if owned is not None else None + + +def bridge_attachment( + attachment: Dict[str, Any], *, user_id: str, conversation_id: str +) -> str: + """Return the conversation artifact id for ``attachment``, reusing an existing bridge or creating one. + + Idempotent (best-effort): an artifact already derived from this attachment in this conversation + (matched via its version ``produced_by.attachment_id``) is reused, so a second reference never + consumes a new quota slot. The reuse is a read-then-write across transactions, so two concurrent + references to the same not-yet-bridged attachment may each create one. Otherwise the attachment + bytes are read server-side and persisted as a conversation-scoped ``file`` artifact (server-computed + size/sha256/storage key). + """ + attachment_id = str(attachment.get("id") or attachment.get("_id") or attachment.get("legacy_mongo_id")) + try: + with db_readonly() as conn: + existing = ArtifactsRepository(conn).find_bridged_attachment( + attachment_id, conversation_id=conversation_id + ) + except Exception: + logger.exception("attachment_bridge: idempotency lookup failed") + existing = None + if existing is not None: + return str(existing["id"]) + + upload_path = attachment.get("upload_path") or attachment.get("path") + if not upload_path: + raise AttachmentBridgeError(f"attachment {attachment_id} has no stored content.") + filename = attachment.get("filename") or "attachment" + mime_type = attachment.get("mime_type") or "application/octet-stream" + try: + data = StorageCreator.get_storage().get_file(upload_path).read() + except Exception as exc: + logger.exception("attachment_bridge: failed to read attachment bytes") + raise AttachmentBridgeError(f"failed to read attachment {attachment_id}.") from exc + try: + ref = persist_new_artifact( + user_id=user_id, + kind="file", + data=data, + filename=filename, + mime_type=mime_type, + title=filename, + conversation_id=conversation_id, + produced_by={"attachment_id": attachment_id, "source": "chat_attachment"}, + ) + except QuotaExceeded as exc: + raise AttachmentBridgeError(str(exc)) from exc + if ref is None: + raise AttachmentBridgeError(f"failed to bridge attachment {attachment_id}.") + return str(ref["artifact_id"]) diff --git a/application/agents/tools/code_executor.py b/application/agents/tools/code_executor.py index 1cadb86f..896ffd9a 100644 --- a/application/agents/tools/code_executor.py +++ b/application/agents/tools/code_executor.py @@ -7,6 +7,11 @@ import re from typing import Any, Dict, List, Optional, Tuple from application.agents.tools.artifact_ref import resolve_artifact_id +from application.agents.tools.attachment_bridge import ( + AttachmentBridgeError, + bridge_attachment, + match_attachment, +) from application.agents.tools.base import Tool from application.core.settings import settings from application.sandbox.artifacts_capture import ( @@ -97,9 +102,9 @@ class CodeExecutorTool(Tool): "inputs": { "type": "array", "items": {"type": "string"}, - "description": "Artifacts (from this conversation/run) to materialize into the " - "workspace; each accepts the short ref like `A1` returned by a previous " - "artifact action, or the full artifact id.", + "description": "Files to materialize into the workspace; each accepts the short " + "ref like `A1` returned by a previous artifact action, a full artifact id, or " + "the name/id of a file the user attached to this conversation.", }, "timeout": { "type": "integer", @@ -257,7 +262,20 @@ class CodeExecutorTool(Tool): else None ) if artifact is None: - return {"error": f"input artifact {raw} not found in this conversation/run."} + # Conversation scope only: a raw ref that is not an artifact + # may name a chat attachment; bridge it on demand. Workflows + # bridge attachments up front, so never double-bridge there. + bridged_id = self._bridge_chat_attachment(raw) + if isinstance(bridged_id, dict): + return bridged_id # error payload + if bridged_id is None: + return {"error": f"input artifact {raw} not found in this conversation/run."} + artifact_id = bridged_id + artifact = repo.get_artifact_in_parent( + artifact_id, conversation_id=self.conversation_id + ) + if artifact is None: + return {"error": f"input artifact {raw} not found in this conversation/run."} version = repo.get_version(artifact_id, artifact["current_version"]) except Exception: logger.exception("code_executor: failed to load input artifact") @@ -281,6 +299,20 @@ class CodeExecutorTool(Tool): loaded.append(f"inputs/{filename}") return {"loaded": loaded} + def _bridge_chat_attachment(self, raw: str) -> Any: + """Bridge a referenced chat attachment to a conversation artifact id; None on miss, error dict on failure.""" + if not self.conversation_id or not self.user_id: + return None + attachment = match_attachment(self.config.get("attachments"), raw, self.user_id) + if attachment is None: + return None + try: + return bridge_attachment( + attachment, user_id=self.user_id, conversation_id=self.conversation_id + ) + except AttachmentBridgeError as exc: + return {"error": f"failed to attach {raw}: {exc}"} + # Cap the per-run capture work so a workspace full of pre-existing files # can't turn one exec into an unbounded read+persist sweep. _MAX_CAPTURED_FILES = MAX_CAPTURED_FILES diff --git a/application/agents/tools/read_document.py b/application/agents/tools/read_document.py index 0d19ea77..0f538245 100644 --- a/application/agents/tools/read_document.py +++ b/application/agents/tools/read_document.py @@ -15,6 +15,11 @@ import logging from typing import Any, Dict, List, Optional from application.agents.tools.artifact_ref import resolve_artifact_id +from application.agents.tools.attachment_bridge import ( + AttachmentBridgeError, + bridge_attachment, + match_attachment, +) from application.agents.tools.base import Tool from application.core.json_schema_utils import ( JsonSchemaValidationError, @@ -66,8 +71,9 @@ class ReadDocumentTool(Tool): "properties": { "input": { "type": "string", - "description": "Document (from this conversation/run) to read; accepts the short " - "ref like `A1` returned by a previous artifact action, or the full artifact id.", + "description": "Document to read; accepts the short ref like `A1` returned by a " + "previous artifact action, a full artifact id, or the name/id of a file the user " + "attached to this conversation.", }, "output": { "type": "string", @@ -232,9 +238,30 @@ class ReadDocumentTool(Tool): logger.exception("read_document: failed to resolve input artifact") return {"status": "error", "error": f"failed to resolve input artifact {raw_id}."} if artifact is None: + # Conversation scope only: a raw ref that is not an artifact may name a + # chat attachment; bridge it on demand. Workflows bridge up front. + bridged_id = self._bridge_chat_attachment(raw_id) + if isinstance(bridged_id, dict): + return bridged_id + if bridged_id is not None: + return bridged_id return {"status": "error", "error": f"input artifact {raw_id} not found in this conversation/run."} return str(artifact_id) + def _bridge_chat_attachment(self, raw_id: str) -> Any: + """Bridge a referenced chat attachment to a conversation artifact id; None on miss, error dict on failure.""" + if not self.conversation_id or not self.user_id: + return None + attachment = match_attachment(self.config.get("attachments"), raw_id, self.user_id) + if attachment is None: + return None + try: + return bridge_attachment( + attachment, user_id=self.user_id, conversation_id=self.conversation_id + ) + except AttachmentBridgeError as exc: + return {"status": "error", "error": f"failed to attach {raw_id}: {exc}"} + def _parent(self) -> Dict[str, Any]: """Build the run-scoped parent dict passed to the worker for its independent re-resolve.""" if self.conversation_id is not None: diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py index 4cc06da7..e703d3b3 100644 --- a/application/storage/db/repositories/artifacts.py +++ b/application/storage/db/repositories/artifacts.py @@ -136,6 +136,27 @@ class ArtifactsRepository: row = result.fetchone() return _artifact_to_dict(row) if row is not None else None + def find_bridged_attachment( + self, + attachment_id: str, + *, + conversation_id: str, + ) -> Optional[dict]: + """Return the conversation artifact already bridged from ``attachment_id``, or None (idempotency gate).""" + result = self._conn.execute( + text( + "SELECT a.* FROM artifacts a " + "JOIN artifact_versions v " + " ON v.artifact_id = a.id AND v.version = a.current_version " + "WHERE a.conversation_id = CAST(:conversation_id AS uuid) " + " AND v.produced_by ->> 'attachment_id' = :attachment_id " + "ORDER BY a.created_at ASC, a.id ASC LIMIT 1" + ), + {"conversation_id": conversation_id, "attachment_id": str(attachment_id)}, + ) + row = result.fetchone() + return _artifact_to_dict(row) if row is not None else None + def list_artifacts( self, conversation_id: Optional[str] = None, diff --git a/tests/agents/tools/test_attachment_bridge.py b/tests/agents/tools/test_attachment_bridge.py new file mode 100644 index 00000000..3e38b864 --- /dev/null +++ b/tests/agents/tools/test_attachment_bridge.py @@ -0,0 +1,463 @@ +"""Unit tests for the lazy chat-attachment -> conversation-artifact bridge. + +Covers the matcher (caller-only, by id and filename), the idempotent bridge +(reuse an already-bridged attachment), the quota-error surface, and the wiring +into ``code_executor._materialize_inputs`` / ``read_document._resolve_input``. +No live DB / storage / sandbox / LLM is touched: the repos, storage, and +``persist_new_artifact`` are stubbed. +""" + +from __future__ import annotations + +import contextlib +import uuid +from typing import Any, Dict, List, Optional + +import pytest + +import application.agents.tools.attachment_bridge as bridge_mod +from application.agents.tools.attachment_bridge import ( + AttachmentBridgeError, + bridge_attachment, + match_attachment, +) +from application.sandbox.artifacts_capture import QuotaExceeded + +CONV = "11111111-1111-1111-1111-111111111111" +USER = "user-1" + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- +class _FakeFile: + def __init__(self, data: bytes) -> None: + self._data = data + + def read(self) -> bytes: + return self._data + + +class _FakeStorage: + def __init__(self, data: bytes = b"hello bytes") -> None: + self._data = data + self.requested: List[str] = [] + + def get_file(self, path): + self.requested.append(path) + return _FakeFile(self._data) + + +class _FakeAttachmentsRepo: + """Confirms ownership: returns the row only for the owner user_id.""" + + rows: Dict[str, Dict[str, Any]] = {} + + def __init__(self, conn=None) -> None: + pass + + def get_any(self, attachment_id: str, user_id: str) -> Optional[Dict[str, Any]]: + row = self.rows.get(str(attachment_id)) + if row is None or row.get("user_id") != user_id: + return None + return row + + +class _FakeArtifactsRepo: + """Idempotency gate + parent-scoped artifact lookups for the bridge + tools.""" + + bridged: Dict[str, str] = {} # attachment_id -> existing artifact id + artifacts: Dict[str, Dict[str, Any]] = {} # artifact_id -> {conversation_id, current_version, versions} + + def __init__(self, conn=None) -> None: + pass + + def find_bridged_attachment(self, attachment_id: str, *, conversation_id: str): + aid = self.bridged.get(str(attachment_id)) + if aid is None: + return None + return {"id": aid} + + def artifact_id_at_position(self, n, *, conversation_id=None, workflow_run_id=None): + return None + + def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None): + art = self.artifacts.get(str(artifact_id)) + if art is None or art.get("conversation_id") != conversation_id: + return None + return {"id": artifact_id, "current_version": art["current_version"]} + + def get_version(self, artifact_id, version): + art = self.artifacts.get(str(artifact_id)) + return art["versions"].get(version) if art is not None else None + + +@contextlib.contextmanager +def _fake_db(): + yield object() + + +def _attachment(*, aid=None, filename="report.pdf", user_id=USER, mime="application/pdf", path="up/report.pdf"): + return { + "id": aid or str(uuid.uuid4()), + "filename": filename, + "user_id": user_id, + "mime_type": mime, + "upload_path": path, + } + + +def _patch_bridge(monkeypatch, *, storage=None, persisted_id="art-new"): + """Patch the bridge module's repos/storage/persist so no infra is touched.""" + storage = storage or _FakeStorage() + monkeypatch.setattr(bridge_mod, "db_readonly", _fake_db) + monkeypatch.setattr(bridge_mod, "AttachmentsRepository", _FakeAttachmentsRepo) + monkeypatch.setattr(bridge_mod, "ArtifactsRepository", _FakeArtifactsRepo) + monkeypatch.setattr( + bridge_mod.StorageCreator, "get_storage", staticmethod(lambda: storage) + ) + calls: List[Dict[str, Any]] = [] + + def _persist(**kwargs): + calls.append(kwargs) + return {"artifact_id": persisted_id, "version": 1, "filename": kwargs["filename"]} + + monkeypatch.setattr(bridge_mod, "persist_new_artifact", _persist) + return storage, calls + + +@pytest.fixture(autouse=True) +def _reset_state(): + _FakeAttachmentsRepo.rows = {} + _FakeArtifactsRepo.bridged = {} + _FakeArtifactsRepo.artifacts = {} + yield + + +# --------------------------------------------------------------------------- +# Matcher: caller-only, by id and by filename +# --------------------------------------------------------------------------- +@pytest.mark.unit +def test_match_by_id_confirms_owner(monkeypatch): + _patch_bridge(monkeypatch) + att = _attachment() + _FakeAttachmentsRepo.rows = {att["id"]: att} + out = match_attachment([att], att["id"], USER) + assert out is not None and out["id"] == att["id"] + + +@pytest.mark.unit +def test_match_by_filename_exact_and_normalized(monkeypatch): + _patch_bridge(monkeypatch) + att = _attachment(filename="Report.PDF") + _FakeAttachmentsRepo.rows = {att["id"]: att} + assert match_attachment([att], "report.pdf", USER) is not None + assert match_attachment([att], " Report.PDF ", USER) is not None + + +@pytest.mark.unit +def test_no_match_returns_none(monkeypatch): + _patch_bridge(monkeypatch) + att = _attachment(filename="a.pdf") + _FakeAttachmentsRepo.rows = {att["id"]: att} + assert match_attachment([att], "other.pdf", USER) is None + assert match_attachment([], "a.pdf", USER) is None + assert match_attachment(None, "a.pdf", USER) is None + + +@pytest.mark.unit +def test_match_rejects_when_owner_check_fails(monkeypatch): + """An attachment dict present in the request but NOT owned by user is rejected at the owner re-check.""" + _patch_bridge(monkeypatch) + att = _attachment(user_id="someone-else") + # repo has no row for USER -> ownership re-check fails even though the dict was supplied. + _FakeAttachmentsRepo.rows = {} + assert match_attachment([att], att["id"], USER) is None + + +# --------------------------------------------------------------------------- +# Bridge: create + idempotent reuse + quota +# --------------------------------------------------------------------------- +@pytest.mark.unit +def test_bridge_creates_conversation_artifact(monkeypatch): + storage, calls = _patch_bridge(monkeypatch, persisted_id="art-A") + att = _attachment() + art_id = bridge_attachment(att, user_id=USER, conversation_id=CONV) + assert art_id == "art-A" + assert len(calls) == 1 + kw = calls[0] + assert kw["conversation_id"] == CONV and kw["user_id"] == USER and kw["kind"] == "file" + assert kw["produced_by"] == {"attachment_id": att["id"], "source": "chat_attachment"} + assert storage.requested == ["up/report.pdf"] + + +@pytest.mark.unit +def test_bridge_reuses_existing(monkeypatch): + _storage, calls = _patch_bridge(monkeypatch) + att = _attachment() + _FakeArtifactsRepo.bridged = {att["id"]: "art-existing"} + art_id = bridge_attachment(att, user_id=USER, conversation_id=CONV) + assert art_id == "art-existing" + assert calls == [] # no second persist -> no extra quota slot + + +@pytest.mark.unit +def test_bridge_surfaces_quota_as_clean_error(monkeypatch): + _patch_bridge(monkeypatch) + att = _attachment() + + def _persist(**kwargs): + raise QuotaExceeded("artifact count quota reached (5)") + + monkeypatch.setattr(bridge_mod, "persist_new_artifact", _persist) + with pytest.raises(AttachmentBridgeError, match="quota"): + bridge_attachment(att, user_id=USER, conversation_id=CONV) + + +@pytest.mark.unit +def test_bridge_missing_upload_path_errors(monkeypatch): + _patch_bridge(monkeypatch) + att = _attachment(path=None) + att["upload_path"] = None + att["path"] = None + with pytest.raises(AttachmentBridgeError, match="no stored content"): + bridge_attachment(att, user_id=USER, conversation_id=CONV) + + +# --------------------------------------------------------------------------- +# code_executor wiring: fallback fires, stages bytes, succeeds +# --------------------------------------------------------------------------- +def _patch_code_executor_repo(monkeypatch): + import application.agents.tools.code_executor as ce + + monkeypatch.setattr(ce, "db_readonly", _fake_db) + monkeypatch.setattr(ce, "ArtifactsRepository", _FakeArtifactsRepo) + + +class _Manager: + def __init__(self): + self.staged: Dict[str, bytes] = {} + + def put_file(self, session_id, path, data): + self.staged[path] = data + + +@pytest.mark.unit +def test_code_executor_bridges_referenced_attachment_by_name(monkeypatch): + from application.agents.tools.code_executor import CodeExecutorTool + + storage, calls = _patch_bridge(monkeypatch, storage=_FakeStorage(b"PDFDATA"), persisted_id="art-A") + _patch_code_executor_repo(monkeypatch) + att = _attachment(filename="statement.pdf", path="up/statement.pdf") + _FakeAttachmentsRepo.rows = {att["id"]: att} + # After the bridge creates art-A, the tool re-reads it parent-scoped. + _FakeArtifactsRepo.artifacts = { + "art-A": { + "conversation_id": CONV, + "current_version": 1, + "versions": {1: {"storage_path": "store/art-A", "filename": "statement.pdf"}}, + } + } + tool = CodeExecutorTool( + tool_config={"conversation_id": CONV, "attachments": [att]}, user_id=USER + ) + mgr = _Manager() + out = tool._materialize_inputs(mgr, "sess", ["statement.pdf"]) + assert out.get("error") is None + assert out["loaded"] == ["inputs/statement.pdf"] + assert mgr.staged["inputs/statement.pdf"] == b"PDFDATA" + assert len(calls) == 1 # bridged once + + +@pytest.mark.unit +def test_code_executor_bridges_referenced_attachment_by_id(monkeypatch): + from application.agents.tools.code_executor import CodeExecutorTool + + _patch_bridge(monkeypatch, storage=_FakeStorage(b"X"), persisted_id="art-B") + _patch_code_executor_repo(monkeypatch) + att = _attachment(filename="data.csv", path="up/data.csv") + _FakeAttachmentsRepo.rows = {att["id"]: att} + _FakeArtifactsRepo.artifacts = { + "art-B": { + "conversation_id": CONV, + "current_version": 1, + "versions": {1: {"storage_path": "store/art-B", "filename": "data.csv"}}, + } + } + tool = CodeExecutorTool( + tool_config={"conversation_id": CONV, "attachments": [att]}, user_id=USER + ) + mgr = _Manager() + out = tool._materialize_inputs(mgr, "sess", [att["id"]]) + assert out["loaded"] == ["inputs/data.csv"] + + +@pytest.mark.unit +def test_code_executor_idempotent_reuse_no_second_persist(monkeypatch): + from application.agents.tools.code_executor import CodeExecutorTool + + _storage, calls = _patch_bridge(monkeypatch, persisted_id="art-A") + _patch_code_executor_repo(monkeypatch) + att = _attachment(filename="doc.pdf", path="up/doc.pdf") + _FakeAttachmentsRepo.rows = {att["id"]: att} + # Pretend it was already bridged once -> reuse path. + _FakeArtifactsRepo.bridged = {att["id"]: "art-A"} + _FakeArtifactsRepo.artifacts = { + "art-A": { + "conversation_id": CONV, + "current_version": 1, + "versions": {1: {"storage_path": "store/art-A", "filename": "doc.pdf"}}, + } + } + tool = CodeExecutorTool( + tool_config={"conversation_id": CONV, "attachments": [att]}, user_id=USER + ) + mgr = _Manager() + # Reference the same attachment twice in one call. + out = tool._materialize_inputs(mgr, "sess", ["doc.pdf", att["id"]]) + assert out.get("error") is None + assert calls == [] # never persisted -> reused the existing artifact both times + + +@pytest.mark.unit +def test_code_executor_rejects_foreign_attachment(monkeypatch): + from application.agents.tools.code_executor import CodeExecutorTool + + _patch_bridge(monkeypatch) + _patch_code_executor_repo(monkeypatch) + foreign = _attachment(filename="secret.pdf", user_id="attacker") + # The model names another user's file, but it is not in THIS request's attachments + # and ownership re-check has no row for USER. + _FakeAttachmentsRepo.rows = {} + tool = CodeExecutorTool( + tool_config={"conversation_id": CONV, "attachments": []}, user_id=USER + ) + out = tool._materialize_inputs(_Manager(), "sess", [foreign["id"]]) + assert "not found in this conversation/run" in out["error"] + + +@pytest.mark.unit +def test_code_executor_workflow_scope_does_not_bridge(monkeypatch): + from application.agents.tools.code_executor import CodeExecutorTool + + _storage, calls = _patch_bridge(monkeypatch) + _patch_code_executor_repo(monkeypatch) + att = _attachment(filename="wf.pdf") + _FakeAttachmentsRepo.rows = {att["id"]: att} + # workflow_run_id set, no conversation_id -> attachment fallback must NOT fire. + tool = CodeExecutorTool( + tool_config={"workflow_run_id": CONV, "attachments": [att]}, user_id=USER + ) + out = tool._materialize_inputs(_Manager(), "sess", ["wf.pdf"]) + assert "not found in this conversation/run" in out["error"] + assert calls == [] # no double-bridge in workflow scope + + +@pytest.mark.unit +def test_code_executor_unresolvable_ref_still_errors(monkeypatch): + from application.agents.tools.code_executor import CodeExecutorTool + + _patch_bridge(monkeypatch) + _patch_code_executor_repo(monkeypatch) + tool = CodeExecutorTool( + tool_config={"conversation_id": CONV, "attachments": []}, user_id=USER + ) + out = tool._materialize_inputs(_Manager(), "sess", ["ghost.bin"]) + assert "not found in this conversation/run" in out["error"] + + +@pytest.mark.unit +def test_code_executor_existing_artifact_still_used(monkeypatch): + """A ref that resolves to a real artifact uses it directly (no bridge regression).""" + from application.agents.tools.code_executor import CodeExecutorTool + + storage, calls = _patch_bridge(monkeypatch) + _patch_code_executor_repo(monkeypatch) + real_id = str(uuid.uuid4()) + _FakeArtifactsRepo.artifacts = { + real_id: { + "conversation_id": CONV, + "current_version": 1, + "versions": {1: {"storage_path": "store/real", "filename": "real.txt"}}, + } + } + monkeypatch.setattr( + "application.agents.tools.code_executor.StorageCreator.get_storage", + staticmethod(lambda: storage), + ) + tool = CodeExecutorTool( + tool_config={"conversation_id": CONV, "attachments": []}, user_id=USER + ) + mgr = _Manager() + out = tool._materialize_inputs(mgr, "sess", [real_id]) + assert out["loaded"] == ["inputs/real.txt"] + assert calls == [] # existing artifact -> no bridge + + +# --------------------------------------------------------------------------- +# read_document wiring +# --------------------------------------------------------------------------- +def _patch_read_document_repo(monkeypatch): + import application.agents.tools.read_document as rd + + monkeypatch.setattr(rd, "db_readonly", _fake_db) + monkeypatch.setattr(rd, "ArtifactsRepository", _FakeArtifactsRepo) + + +@pytest.mark.unit +def test_read_document_bridges_attachment_then_enqueues(monkeypatch): + import application.agents.tools.read_document as rd + from application.agents.tools.read_document import ReadDocumentTool + + _storage, calls = _patch_bridge(monkeypatch, persisted_id="art-RD") + _patch_read_document_repo(monkeypatch) + att = _attachment(filename="invoice.pdf", path="up/invoice.pdf") + _FakeAttachmentsRepo.rows = {att["id"]: att} + + captured: Dict[str, Any] = {} + + class _AR: + def get(self, timeout=None): + return {"status": "ok", "content": "parsed", "truncated": False} + + import application.api.user.tasks as tasks + + def _apply_async(args=None, queue=None, **kw): + captured["args"] = args + return _AR() + + monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async) + + tool = ReadDocumentTool( + tool_config={"conversation_id": CONV, "attachments": [att]}, user_id=USER + ) + out = tool.execute_action("read_document", input="invoice.pdf", persist=False) + assert out["status"] == "ok" + # The bridged artifact id (not the raw name) was enqueued, with the conversation parent. + assert captured["args"][0] == "art-RD" + assert captured["args"][1] == {"conversation_id": CONV} + assert len(calls) == 1 + _ = rd # keep import referenced + + +@pytest.mark.unit +def test_read_document_workflow_scope_does_not_bridge(monkeypatch): + from application.agents.tools.read_document import ReadDocumentTool + + _storage, calls = _patch_bridge(monkeypatch) + _patch_read_document_repo(monkeypatch) + att = _attachment(filename="wf.pdf") + _FakeAttachmentsRepo.rows = {att["id"]: att} + + import application.api.user.tasks as tasks + monkeypatch.setattr( + tasks.parse_document, "apply_async", + lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not enqueue")), + ) + + tool = ReadDocumentTool( + tool_config={"workflow_run_id": CONV, "attachments": [att]}, user_id=USER + ) + out = tool.execute_action("read_document", input="wf.pdf", persist=False) + assert out["status"] == "error" and "not found in this conversation/run" in out["error"] + assert calls == []