diff --git a/application/agents/tools/artifact_generator.py b/application/agents/tools/artifact_generator.py index 451f6f37..18f5dc18 100644 --- a/application/agents/tools/artifact_generator.py +++ b/application/agents/tools/artifact_generator.py @@ -17,6 +17,7 @@ import logging import uuid from typing import Any, Dict, List, Optional +from application.agents.tools.artifact_ref import resolve_artifact_id from application.agents.tools.base import Tool from application.core.settings import settings from application.sandbox.artifacts_capture import ( @@ -269,7 +270,8 @@ class ArtifactGeneratorTool(Tool): "name": "create_artifact", "description": ( "Render a new editable document from a JSON spec and store it as version 1. " - "The spec is the source of truth; the rendered file is derived." + "The spec is the source of truth; the rendered file is derived. The response " + "carries a short ref (like `A1`) you can pass to edit_artifact/rewrite_artifact." ), "active": True, "parameters": { @@ -292,7 +294,11 @@ class ArtifactGeneratorTool(Tool): "parameters": { "type": "object", "properties": { - "id": {"type": "string", "description": "Artifact id to edit."}, + "id": { + "type": "string", + "description": "Artifact to edit; accepts the short ref like `A1` " + "returned by a previous artifact action, or the full artifact id.", + }, "spec_patch": { "type": "object", "description": "RFC 7386 merge-patch; null values delete keys.", @@ -308,7 +314,11 @@ class ArtifactGeneratorTool(Tool): "parameters": { "type": "object", "properties": { - "id": {"type": "string", "description": "Artifact id to rewrite."}, + "id": { + "type": "string", + "description": "Artifact to rewrite; accepts the short ref like `A1` " + "returned by a previous artifact action, or the full artifact id.", + }, "spec": {"type": "object", "description": "Replacement spec matching the kind's schema."}, }, "required": ["id", "spec"], @@ -386,25 +396,22 @@ class ArtifactGeneratorTool(Tool): def _edit(self, **kwargs: Any) -> Dict[str, Any]: """Merge-patch the current spec, re-render, and append a version.""" - artifact_id = kwargs.get("id") spec_patch = kwargs.get("spec_patch") if not isinstance(spec_patch, dict): return {"status": "error", "error": "spec_patch must be a JSON object (merge-patch)."} - loaded = self._load_current(artifact_id) + loaded = self._load_current(kwargs.get("id")) if loaded.get("error"): return {"status": "error", "error": loaded["error"]} - kind = loaded["kind"] new_spec = merge_patch(loaded["spec"], spec_patch) - return self._reversion(artifact_id, kind, new_spec, "edit_artifact") + return self._reversion(loaded["artifact_id"], loaded["kind"], new_spec, "edit_artifact") def _rewrite(self, **kwargs: Any) -> Dict[str, Any]: """Replace the spec wholesale, re-render, and append a version.""" - artifact_id = kwargs.get("id") spec = kwargs.get("spec") - loaded = self._load_current(artifact_id) + loaded = self._load_current(kwargs.get("id")) if loaded.get("error"): return {"status": "error", "error": loaded["error"]} - return self._reversion(artifact_id, loaded["kind"], spec, "rewrite_artifact") + return self._reversion(loaded["artifact_id"], loaded["kind"], spec, "rewrite_artifact") def _reversion(self, artifact_id: str, kind: str, spec: Any, action: str) -> Dict[str, Any]: """Validate the new spec, re-render, and append the next version of an existing artifact.""" @@ -425,6 +432,8 @@ class ArtifactGeneratorTool(Tool): mime_type=info["mime"], spec=spec, produced_by=self._produced_by(action, kind), + conversation_id=self.conversation_id, + workflow_run_id=self.workflow_run_id, ) except QuotaExceeded as exc: return {"status": "error", "error": str(exc)} @@ -446,30 +455,41 @@ class ArtifactGeneratorTool(Tool): return {"status": "error", "error": f"invalid {kind} spec: {exc.message}"} return None - def _load_current(self, artifact_id: Any) -> Dict[str, Any]: - """Fetch the parent-scoped artifact and its current-version spec for edit/rewrite.""" - if not isinstance(artifact_id, str) or not artifact_id.strip(): + def _load_current(self, raw_id: Any) -> Dict[str, Any]: + """Resolve a short ref/uuid to its parent-scoped artifact and current-version spec for edit/rewrite.""" + if not isinstance(raw_id, str) or not raw_id.strip(): return {"error": "id is required."} try: with db_readonly() as conn: repo = ArtifactsRepository(conn) + # A ref (A1/A2/...) resolves to an id within this parent only; the + # resolved id is then re-checked through the parent-scoped gate so a + # ref can never reach another tenant. + artifact_id = resolve_artifact_id( + repo, + raw_id.strip(), + conversation_id=self.conversation_id, + workflow_run_id=self.workflow_run_id, + ) + if artifact_id is None: + return {"error": f"artifact {raw_id} not found in this conversation/run."} artifact = repo.get_artifact_in_parent( artifact_id, conversation_id=self.conversation_id, workflow_run_id=self.workflow_run_id, ) if artifact is None: - return {"error": f"artifact {artifact_id} not found in this conversation/run."} + return {"error": f"artifact {raw_id} not found in this conversation/run."} version = repo.get_version(artifact_id, artifact["current_version"]) except Exception: logger.exception("artifact_generator: failed to load artifact") - return {"error": f"failed to load artifact {artifact_id}."} + return {"error": f"failed to load artifact {raw_id}."} if not version or version.get("spec") is None: - return {"error": f"artifact {artifact_id} has no editable spec."} + return {"error": f"artifact {raw_id} has no editable spec."} kind = self._kind_for(artifact, version) if kind is None: - return {"error": f"artifact {artifact_id} is not a spec-rendered document."} - return {"kind": kind, "spec": version["spec"]} + return {"error": f"artifact {raw_id} is not a spec-rendered document."} + return {"artifact_id": artifact_id, "kind": kind, "spec": version["spec"]} @staticmethod def _kind_for(artifact: Dict[str, Any], version: Dict[str, Any]) -> Optional[str]: diff --git a/application/agents/tools/artifact_ref.py b/application/agents/tools/artifact_ref.py new file mode 100644 index 00000000..830c622d --- /dev/null +++ b/application/agents/tools/artifact_ref.py @@ -0,0 +1,51 @@ +"""Virtual short artifact handles (``A1``, ``A2``, ...) the model can type to reference an artifact. + +A ref is NOT persisted: ``A{n}`` is the n-th artifact (1-based, created_at asc) within the +caller's parent (``conversation_id`` or ``workflow_run_id``). Refs resolve only inside that +parent, never cross-tenant; resolution still goes through the parent-scoped authz gate. +""" + +from __future__ import annotations + +import re +from typing import Any, Optional + +from application.storage.db.base_repository import looks_like_uuid + +_REF_RE = re.compile(r"^[Aa](\d+)$") + + +def make_ref(position: int) -> str: + """Build the short ref string for a 1-based position (``1`` -> ``"A1"``).""" + return f"A{position}" + + +def parse_ref(value: Any) -> Optional[int]: + """Parse a short ref like ``A1``/``a2`` into its 1-based position, or None when it is not a ref.""" + if not isinstance(value, str): + return None + match = _REF_RE.match(value.strip()) + if match is None: + return None + position = int(match.group(1)) + return position if position >= 1 else None + + +def resolve_artifact_id( + repo: Any, + raw: Any, + *, + conversation_id: Optional[str] = None, + workflow_run_id: Optional[str] = None, +) -> Optional[str]: + """Resolve a short ref or a uuid to an artifact id, scoped to the caller's parent; None otherwise.""" + position = parse_ref(raw) + if position is not None: + return repo.artifact_id_at_position( + position, + conversation_id=conversation_id, + workflow_run_id=workflow_run_id, + ) + if looks_like_uuid(raw): + return str(raw).strip() + return None diff --git a/application/agents/tools/code_executor.py b/application/agents/tools/code_executor.py index 095505a2..1cadb86f 100644 --- a/application/agents/tools/code_executor.py +++ b/application/agents/tools/code_executor.py @@ -6,6 +6,7 @@ import logging import re from typing import Any, Dict, List, Optional, Tuple +from application.agents.tools.artifact_ref import resolve_artifact_id from application.agents.tools.base import Tool from application.core.settings import settings from application.sandbox.artifacts_capture import ( @@ -96,7 +97,9 @@ class CodeExecutorTool(Tool): "inputs": { "type": "array", "items": {"type": "string"}, - "description": "Artifact ids (from this conversation/run) to materialize into the workspace.", + "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.", }, "timeout": { "type": "integer", @@ -228,19 +231,33 @@ class CodeExecutorTool(Tool): return {"loaded": loaded} storage = StorageCreator.get_storage() for raw_id in inputs: - artifact_id = str(raw_id).strip() - if not artifact_id: + raw = str(raw_id).strip() + if not raw: continue + artifact_id: Optional[str] = raw try: with db_readonly() as conn: repo = ArtifactsRepository(conn) - artifact = repo.get_artifact_in_parent( - artifact_id, + # A short ref (A1/A2/...) resolves to an id within this parent + # only; the resolved id still passes through the parent-scoped + # gate so a ref can never reach another tenant. + artifact_id = resolve_artifact_id( + repo, + raw, conversation_id=self.conversation_id, workflow_run_id=self.workflow_run_id, ) + artifact = ( + repo.get_artifact_in_parent( + artifact_id, + conversation_id=self.conversation_id, + workflow_run_id=self.workflow_run_id, + ) + if artifact_id is not None + else None + ) if artifact is None: - return {"error": f"input artifact {artifact_id} not found in this conversation/run."} + 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") diff --git a/application/agents/tools/document_extractor.py b/application/agents/tools/document_extractor.py index 4f489e9f..a93f9104 100644 --- a/application/agents/tools/document_extractor.py +++ b/application/agents/tools/document_extractor.py @@ -18,6 +18,7 @@ import logging import uuid from typing import Any, Dict, List, Optional +from application.agents.tools.artifact_ref import resolve_artifact_id from application.agents.tools.base import Tool from application.core.json_schema_utils import ( JsonSchemaValidationError, @@ -141,7 +142,8 @@ class DocumentExtractorTool(Tool): "properties": { "input": { "type": "string", - "description": "Artifact id (from this conversation/run) of the document to extract.", + "description": "Document (from this conversation/run) to extract; accepts the short " + "ref like `A1` returned by a previous artifact action, or the full artifact id.", }, "json_schema": { "type": "object", @@ -302,24 +304,38 @@ class DocumentExtractorTool(Tool): # ------------------------------------------------------------------ # Input / payload helpers # ------------------------------------------------------------------ - def _load_input(self, artifact_id: str) -> Dict[str, Any]: - """Fetch the parent-scoped input artifact's current-version bytes; never cross-tenant.""" + def _load_input(self, raw_id: str) -> Dict[str, Any]: + """Resolve a short ref/uuid, then fetch the parent-scoped input bytes; never cross-tenant.""" + artifact_id: Optional[str] = raw_id try: with db_readonly() as conn: repo = ArtifactsRepository(conn) - artifact = repo.get_artifact_in_parent( - artifact_id, + # A ref (A1/A2/...) resolves to an id within this parent only; the + # resolved id is re-checked through the parent-scoped gate so a ref + # can never reach another tenant. + artifact_id = resolve_artifact_id( + repo, + raw_id, conversation_id=self.conversation_id, workflow_run_id=self.workflow_run_id, ) + artifact = ( + repo.get_artifact_in_parent( + artifact_id, + conversation_id=self.conversation_id, + workflow_run_id=self.workflow_run_id, + ) + if artifact_id is not None + else None + ) if artifact is None: - return {"error": f"input artifact {artifact_id} not found in this conversation/run."} + return {"error": f"input artifact {raw_id} not found in this conversation/run."} version = repo.get_version(artifact_id, artifact["current_version"]) except Exception: logger.exception("document_extractor: failed to load input artifact") - return {"error": f"failed to load input artifact {artifact_id}."} + return {"error": f"failed to load input artifact {raw_id}."} if not version or not version.get("storage_path"): - return {"error": f"input artifact {artifact_id} has no stored content."} + return {"error": f"input artifact {raw_id} has no stored content."} display_name = version.get("filename") or artifact.get("title") or artifact_id filename = safe_filename(display_name) try: @@ -327,7 +343,7 @@ class DocumentExtractorTool(Tool): data = file_obj.read() except Exception: logger.exception("document_extractor: failed to read input artifact bytes") - return {"error": f"failed to read input artifact {artifact_id}."} + return {"error": f"failed to read input artifact {raw_id}."} return {"data": data, "filename": filename, "title": display_name} def _compact_payload(self, extracted: Dict[str, Any]) -> Dict[str, Any]: diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index dd07c71d..78a0b0b9 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -424,27 +424,31 @@ class WorkflowEngine: self, manager: Any, session_id: str, inputs: List[str], user_id: str ) -> List[str]: """Stage referenced input artifacts (run-scoped, never cross-tenant) into the workspace.""" - from application.storage.db.base_repository import looks_like_uuid + from application.agents.tools.artifact_ref import resolve_artifact_id from application.storage.db.repositories.artifacts import ArtifactsRepository from application.storage.db.session import db_readonly from application.storage.storage_creator import StorageCreator from application.utils import safe_filename loaded: List[str] = [] - artifact_ids = self._resolve_input_artifact_ids(inputs) - if not artifact_ids: + raw_ids = self._resolve_input_artifact_ids(inputs) + if not raw_ids: return loaded storage = StorageCreator.get_storage() - for artifact_id in artifact_ids: - if not looks_like_uuid(artifact_id): - raise ValueError(f"input artifact {artifact_id} not found in this run.") + for raw in raw_ids: with db_readonly() as conn: repo = ArtifactsRepository(conn) - artifact = repo.get_artifact_in_parent( - artifact_id, workflow_run_id=self.workflow_run_id + # A short ref (A1/A2/...) resolves to an id within this run only; + # the resolved id is re-checked through the run-scoped gate so a ref + # can never reach another tenant. + 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"input artifact {artifact_id} not found in this run.") + raise ValueError(f"input artifact {raw} 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 artifact {artifact_id} has no stored content.") diff --git a/application/sandbox/artifacts_capture.py b/application/sandbox/artifacts_capture.py index f3f27006..107b0ee1 100644 --- a/application/sandbox/artifacts_capture.py +++ b/application/sandbox/artifacts_capture.py @@ -16,6 +16,7 @@ from typing import Any, Dict, List, Optional, Tuple from sqlalchemy import text +from application.agents.tools.artifact_ref import make_ref from application.core.settings import settings from application.storage.db.repositories.artifacts import ArtifactsRepository from application.storage.db.session import db_session @@ -220,6 +221,26 @@ def _enforce_user_quota(repo: ArtifactsRepository, user_id: str, added_bytes: in raise QuotaExceeded(f"artifact storage quota reached ({max_total} bytes); delete artifacts to free space") +def _ref_for( + repo: ArtifactsRepository, + artifact_id: str, + *, + conversation_id: Optional[str], + workflow_run_id: Optional[str], +) -> Optional[str]: + """Compute the short ref (``A{n}``) for an artifact from its position in its parent; None on failure.""" + if conversation_id is None and workflow_run_id is None: + return None + try: + position = repo.position_in_parent( + artifact_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id + ) + except Exception: + logger.exception("artifacts_capture: failed to compute artifact ref") + return None + return make_ref(position) if position >= 1 else None + + def persist_new_artifact( *, user_id: str, @@ -245,6 +266,7 @@ def persist_new_artifact( sha256 = hashlib.sha256(data).hexdigest() storage = StorageCreator.get_storage() saved_key: Optional[str] = None + ref: Optional[str] = None try: with db_session() as conn: repo = ArtifactsRepository(conn) @@ -267,6 +289,9 @@ def persist_new_artifact( artifact_id = str(artifact["id"]) storage_path = _storage_key(user_id, artifact_id, 1, safe_name) _set_version_storage_path(conn, artifact_id, 1, storage_path) + ref = _ref_for( + repo, artifact_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id + ) storage.save_file(io.BytesIO(data), storage_path) saved_key = storage_path except QuotaExceeded: @@ -277,13 +302,16 @@ def persist_new_artifact( logger.exception("artifacts_capture: failed to persist new artifact") _cleanup_orphan(storage, saved_key) return None - return { + payload: Dict[str, Any] = { "artifact_id": artifact_id, "version": 1, "filename": safe_name, "mime_type": mime_type, "size": size, } + if ref is not None: + payload["ref"] = ref + return payload def append_artifact_version( @@ -296,6 +324,8 @@ def append_artifact_version( spec: Any = None, preview_text: Optional[str] = None, produced_by: Any = None, + conversation_id: Optional[str] = None, + workflow_run_id: Optional[str] = None, ) -> Optional[Dict[str, Any]]: """Append a new version (new spec + new bytes) to an existing artifact; return its reference.""" safe_name = safe_filename(filename) @@ -303,6 +333,7 @@ def append_artifact_version( sha256 = hashlib.sha256(data).hexdigest() storage = StorageCreator.get_storage() saved_key: Optional[str] = None + ref: Optional[str] = None try: with db_session() as conn: repo = ArtifactsRepository(conn) @@ -321,6 +352,9 @@ def append_artifact_version( version_number = int(version["version"]) storage_path = _storage_key(user_id, artifact_id, version_number, safe_name) _set_version_storage_path(conn, artifact_id, version_number, storage_path) + ref = _ref_for( + repo, str(artifact_id), conversation_id=conversation_id, workflow_run_id=workflow_run_id + ) storage.save_file(io.BytesIO(data), storage_path) saved_key = storage_path except QuotaExceeded: @@ -329,10 +363,13 @@ def append_artifact_version( logger.exception("artifacts_capture: failed to append artifact version") _cleanup_orphan(storage, saved_key) return None - return { + payload: Dict[str, Any] = { "artifact_id": str(artifact_id), "version": version_number, "filename": safe_name, "mime_type": mime_type, "size": size, } + if ref is not None: + payload["ref"] = ref + return payload diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py index 49ec3392..4cc06da7 100644 --- a/application/storage/db/repositories/artifacts.py +++ b/application/storage/db/repositories/artifacts.py @@ -163,6 +163,72 @@ class ArtifactsRepository: ) return [_artifact_to_dict(r) for r in result.fetchall()] + def position_in_parent( + self, + artifact_id: str, + *, + conversation_id: Optional[str] = None, + workflow_run_id: Optional[str] = None, + ) -> int: + """Return the 1-based position of an artifact within its parent (by created_at, id tie-break); 0 if absent.""" + if conversation_id is None and workflow_run_id is None: + raise ValueError("position_in_parent requires conversation_id or workflow_run_id") + outer, params = self._parent_clauses(conversation_id, workflow_run_id, alias="a") + inner, _ = self._parent_clauses(conversation_id, workflow_run_id, alias="t") + params["id"] = artifact_id + # The inner SELECT applies the same parent scope, so an artifact in a + # different parent yields no anchor row -> count() is 0 (not in this parent). + row = self._conn.execute( + text( + f"SELECT count(*) FROM artifacts a " + f"WHERE {' AND '.join(outer)} AND (a.created_at, a.id) <= (" + f" SELECT t.created_at, t.id FROM artifacts t " + f" WHERE {' AND '.join(inner)} AND t.id = CAST(:id AS uuid)" + f")" + ), + params, + ).fetchone() + return int(row[0]) if row is not None else 0 + + def artifact_id_at_position( + self, + n: int, + *, + conversation_id: Optional[str] = None, + workflow_run_id: Optional[str] = None, + ) -> Optional[str]: + """Return the id of the n-th artifact (1-based, created_at asc, id tie-break) in a parent, or None.""" + if conversation_id is None and workflow_run_id is None: + raise ValueError("artifact_id_at_position requires conversation_id or workflow_run_id") + if not isinstance(n, int) or n < 1: + return None + clauses, params = self._parent_clauses(conversation_id, workflow_run_id) + params["offset"] = n - 1 + row = self._conn.execute( + text( + f"SELECT id FROM artifacts WHERE {' AND '.join(clauses)} " + f"ORDER BY created_at ASC, id ASC OFFSET :offset LIMIT 1" + ), + params, + ).fetchone() + return str(row[0]) if row is not None else None + + @staticmethod + def _parent_clauses( + conversation_id: Optional[str], workflow_run_id: Optional[str], alias: str = "" + ) -> tuple[list[str], dict[str, Any]]: + """Build the parent-scope WHERE clauses + params (optionally column-aliased) for the position helpers.""" + prefix = f"{alias}." if alias else "" + clauses: list[str] = [] + params: dict[str, Any] = {} + if conversation_id is not None: + clauses.append(f"{prefix}conversation_id = CAST(:conversation_id AS uuid)") + params["conversation_id"] = conversation_id + if workflow_run_id is not None: + clauses.append(f"{prefix}workflow_run_id = CAST(:workflow_run_id AS uuid)") + params["workflow_run_id"] = workflow_run_id + return clauses, params + def count_for_user(self, user_id: str) -> int: """Return how many artifacts ``user_id`` currently owns (quota accounting).""" row = self._conn.execute( diff --git a/tests/agents/tools/test_artifact_ref.py b/tests/agents/tools/test_artifact_ref.py new file mode 100644 index 00000000..a7d2b505 --- /dev/null +++ b/tests/agents/tools/test_artifact_ref.py @@ -0,0 +1,104 @@ +"""Unit tests for the virtual short artifact-ref helpers (make_ref / parse_ref / resolve_artifact_id). + +No DB: ``resolve_artifact_id`` is exercised against a tiny fake repo that records the +parent scope it was asked for, proving a ref resolves only within the caller's parent. +""" + +from __future__ import annotations + +import uuid + +from application.agents.tools.artifact_ref import make_ref, parse_ref, resolve_artifact_id + + +class _FakeRepo: + """Records the position lookup and returns a canned id only for the scoped parent.""" + + def __init__(self, *, mapping=None, conv=None, run=None): + self.mapping = mapping or {} + self.conv = conv + self.run = run + self.calls = [] + + def artifact_id_at_position(self, n, *, conversation_id=None, workflow_run_id=None): + self.calls.append((n, conversation_id, workflow_run_id)) + if self.conv is not None and conversation_id != self.conv: + return None + if self.run is not None and workflow_run_id != self.run: + return None + return self.mapping.get(n) + + +# --------------------------------------------------------------------------- +# make_ref / parse_ref +# --------------------------------------------------------------------------- +def test_make_ref_formats_position(): + assert make_ref(1) == "A1" + assert make_ref(2) == "A2" + assert make_ref(42) == "A42" + + +def test_parse_ref_round_trips_make_ref(): + for n in (1, 2, 9, 10, 137): + assert parse_ref(make_ref(n)) == n + + +def test_parse_ref_is_case_insensitive_and_trims(): + assert parse_ref("a1") == 1 + assert parse_ref("A3") == 3 + assert parse_ref(" A7 ") == 7 + + +def test_parse_ref_rejects_non_refs(): + assert parse_ref("A0") is None # 1-based; position 0 is not a ref + assert parse_ref("A") is None + assert parse_ref("AA1") is None + assert parse_ref("1") is None + assert parse_ref("B1") is None + assert parse_ref(str(uuid.uuid4())) is None # a uuid is not a ref + assert parse_ref(None) is None + assert parse_ref(7) is None + + +# --------------------------------------------------------------------------- +# resolve_artifact_id +# --------------------------------------------------------------------------- +def test_resolve_ref_uses_parent_scoped_position(): + target = str(uuid.uuid4()) + repo = _FakeRepo(mapping={1: target}, conv="conv-1") + out = resolve_artifact_id(repo, "A1", conversation_id="conv-1") + assert out == target + assert repo.calls == [(1, "conv-1", None)] + + +def test_resolve_ref_does_not_cross_parents(): + target = str(uuid.uuid4()) + # The repo only yields the id for conv-1; asking under conv-OTHER yields nothing. + repo = _FakeRepo(mapping={1: target}, conv="conv-1") + assert resolve_artifact_id(repo, "A1", conversation_id="conv-OTHER") is None + + +def test_resolve_out_of_range_ref_returns_none(): + repo = _FakeRepo(mapping={1: str(uuid.uuid4())}, conv="conv-1") + assert resolve_artifact_id(repo, "A9", conversation_id="conv-1") is None + + +def test_resolve_uuid_passthrough_without_touching_repo(): + raw = str(uuid.uuid4()) + repo = _FakeRepo(conv="conv-1") + assert resolve_artifact_id(repo, raw, conversation_id="conv-1") == raw + # A uuid never triggers a position lookup. + assert repo.calls == [] + + +def test_resolve_garbage_returns_none(): + repo = _FakeRepo(conv="conv-1") + assert resolve_artifact_id(repo, "not-a-ref-or-uuid", conversation_id="conv-1") is None + assert resolve_artifact_id(repo, "", conversation_id="conv-1") is None + + +def test_resolve_ref_under_workflow_run_parent(): + target = str(uuid.uuid4()) + repo = _FakeRepo(mapping={2: target}, run="run-9") + assert resolve_artifact_id(repo, "A2", workflow_run_id="run-9") == target + assert repo.calls == [(2, None, "run-9")] diff --git a/tests/agents/tools/test_document_extractor_unit.py b/tests/agents/tools/test_document_extractor_unit.py index e8524c1c..0a5ef211 100644 --- a/tests/agents/tools/test_document_extractor_unit.py +++ b/tests/agents/tools/test_document_extractor_unit.py @@ -16,6 +16,7 @@ import os import sys import tempfile import types +import uuid import application.agents.tools.document_extractor as de from application.agents.tools.document_extractor import ( @@ -27,6 +28,10 @@ from application.agents.tools.document_extractor import ( ) from application.sandbox.base import ExecResult +# Production artifact ids are uuids; the resolver gates non-uuid/non-ref inputs to +# "not found", so the stubbed happy-path input uses a real uuid. +_ART_ID = str(uuid.uuid4()) + # A canned Docling-style extraction result the mocked program "writes" to result.json. _CANNED_EXTRACT = { "ok": True, @@ -87,6 +92,16 @@ def _stub_repo(monkeypatch, *, found: bool, conv: Optional[str], run: Optional[s def __init__(self, conn): pass + def artifact_id_at_position(self, n, *, conversation_id=None, workflow_run_id=None): + # The lone artifact lives at position 1, scoped to the right parent only. + if not found or n != 1: + return None + if conv is not None and conversation_id != conv: + return None + if run is not None and workflow_run_id != run: + return None + return _ART_ID + def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None): # Mirror the real scope gate: the artifact resolves only for the right parent. if not found: @@ -188,7 +203,7 @@ def test_extract_stages_input_and_shapes_compact_payload(monkeypatch): manager = _manager_with_extract(_CANNED_EXTRACT) _patch_manager(monkeypatch, manager) - out = _tool().execute_action("extract_document", input="art-1", persist=False) + out = _tool().execute_action("extract_document", input=_ART_ID, persist=False) assert out["status"] == "ok" structured = out["structured"] @@ -214,7 +229,7 @@ def test_extract_removes_its_scratch_token_dir(monkeypatch): manager = _manager_with_extract(_CANNED_EXTRACT) _patch_manager(monkeypatch, manager) - _tool().execute_action("extract_document", input="art-1", persist=False) + _tool().execute_action("extract_document", input=_ART_ID, persist=False) # The per-render token dir (extract/) is removed before/at teardown so # a reused session doesn't accumulate staged inputs + results on disk. @@ -241,7 +256,7 @@ def test_extract_persists_data_artifact_by_reference(monkeypatch): monkeypatch.setattr(de, "persist_new_artifact", _fake_persist) tool = _tool() - out = tool.execute_action("extract_document", input="art-1") # persist defaults to true + out = tool.execute_action("extract_document", input=_ART_ID) # persist defaults to true assert out["status"] == "ok" assert out["artifact"]["artifact_id"] == "new-art" assert tool.get_artifact_id("extract_document") == "new-art" @@ -264,7 +279,7 @@ def test_quota_exceeded_on_persist_surfaces_cleanly(monkeypatch): monkeypatch.setattr(de, "persist_new_artifact", _quota_blocked) - out = _tool().execute_action("extract_document", input="art-1") + out = _tool().execute_action("extract_document", input=_ART_ID) # Extraction still succeeds; the quota failure surfaces as a non-fatal note. assert out["status"] == "ok" assert "artifact" not in out @@ -283,13 +298,42 @@ def test_cross_tenant_input_is_denied(monkeypatch): manager = _manager_with_extract(_CANNED_EXTRACT) _patch_manager(monkeypatch, manager) - out = _tool().execute_action("extract_document", input="art-1", persist=False) + out = _tool().execute_action("extract_document", input=_ART_ID, persist=False) assert out["status"] == "error" assert "not found in this conversation/run" in out["error"] # Nothing was staged or executed for a denied input. assert manager.exec_programs == [] +def test_extract_accepts_short_ref_input(monkeypatch): + # A short ref (A1) resolves to the parent-scoped artifact at position 1 and extracts. + _stub_repo(monkeypatch, found=True, conv="conv-1", run=None) + _patch_storage(monkeypatch) + _patch_no_persist(monkeypatch) + manager = _manager_with_extract(_CANNED_EXTRACT) + _patch_manager(monkeypatch, manager) + + out = _tool().execute_action("extract_document", input="A1", persist=False) + assert out["status"] == "ok" + assert out["structured"]["markdown"].startswith("# Statement") + # The resolved input was staged and the program executed. + assert manager.exec_programs + + +def test_extract_short_ref_out_of_range_is_clean_error(monkeypatch): + # A2 has no artifact behind it (stub only fills position 1) -> clean not-found. + _stub_repo(monkeypatch, found=True, conv="conv-1", run=None) + _patch_storage(monkeypatch) + _patch_no_persist(monkeypatch) + manager = _manager_with_extract(_CANNED_EXTRACT) + _patch_manager(monkeypatch, manager) + + out = _tool().execute_action("extract_document", input="A2", persist=False) + assert out["status"] == "error" + assert "not found in this conversation/run" in out["error"] + assert manager.exec_programs == [] + + def test_missing_input_artifact_is_denied(monkeypatch): _stub_repo(monkeypatch, found=False, conv="conv-1", run=None) _patch_storage(monkeypatch) @@ -320,7 +364,7 @@ def test_json_schema_validation_passes(monkeypatch): _patch_manager(monkeypatch, _manager_with_extract(_CANNED_EXTRACT)) out = _tool().execute_action( - "extract_document", input="art-1", json_schema=_schema_requiring_texts(), persist=False + "extract_document", input=_ART_ID, json_schema=_schema_requiring_texts(), persist=False ) assert out["status"] == "ok" @@ -333,7 +377,7 @@ def test_json_schema_validation_fails_with_clean_error(monkeypatch): schema = {"type": "object", "required": ["amount"], "properties": {"amount": {"type": "number"}}} _patch_manager(monkeypatch, _manager_with_extract(_CANNED_EXTRACT)) - out = _tool().execute_action("extract_document", input="art-1", json_schema=schema, persist=False) + out = _tool().execute_action("extract_document", input=_ART_ID, json_schema=schema, persist=False) assert out["status"] == "error" assert "did not match json_schema" in out["error"] @@ -346,7 +390,7 @@ def test_malformed_json_schema_rejected_before_run(monkeypatch): _patch_manager(monkeypatch, manager) # No "type"/"schema" key -> normalize_json_schema_payload rejects it up front. - out = _tool().execute_action("extract_document", input="art-1", json_schema={"properties": {}}, persist=False) + out = _tool().execute_action("extract_document", input=_ART_ID, json_schema={"properties": {}}, persist=False) assert out["status"] == "error" and "invalid json_schema" in out["error"] assert manager.exec_programs == [] @@ -364,7 +408,7 @@ def test_docling_unavailable_surfaces_clean_error(monkeypatch): } _patch_manager(monkeypatch, _manager_with_extract(err_result)) - out = _tool().execute_action("extract_document", input="art-1", persist=False) + out = _tool().execute_action("extract_document", input=_ART_ID, persist=False) assert out["status"] == "error" assert "docling is not available" in out["error"] @@ -378,7 +422,7 @@ def test_exec_error_surfaces_clean_error(monkeypatch): ) _patch_manager(monkeypatch, manager) - out = _tool().execute_action("extract_document", input="art-1", persist=False) + out = _tool().execute_action("extract_document", input=_ART_ID, persist=False) assert out["status"] == "error" assert "TimeoutError: exceeded 60s" in out["error"] @@ -393,14 +437,14 @@ def test_extraction_program_is_fixed_and_params_are_data(monkeypatch): manager = _manager_with_extract(_CANNED_EXTRACT) _patch_manager(monkeypatch, manager) - _tool().execute_action("extract_document", input="art-1", persist=False) + _tool().execute_action("extract_document", input=_ART_ID, persist=False) program = manager.exec_programs[0] # The executed program is the fixed template with only server-controlled path # literals substituted — no input id, filename, or document content appears. assert "json.load(open(" in program assert "DocumentConverter" in program - assert "art-1" not in program + assert _ART_ID not in program assert "statement.pdf" not in program # Params reached the program as a JSON DATA file, not via interpolation. params_files = [p for p in manager.put_files if p.endswith("params.json")] @@ -439,7 +483,7 @@ def test_malicious_param_value_is_not_executed(monkeypatch): manager = _manager_with_extract(_CANNED_EXTRACT) _patch_manager(monkeypatch, manager) - out = _tool().execute_action("extract_document", input="art-1", persist=False) + out = _tool().execute_action("extract_document", input=_ART_ID, persist=False) assert out["status"] == "ok" program = manager.exec_programs[0] @@ -549,7 +593,7 @@ def test_compact_payload_bounds_table_rows_and_cell_bytes(monkeypatch): extract["tables"] = [huge_table] _patch_manager(monkeypatch, _manager_with_extract(extract)) - out = _tool().execute_action("extract_document", input="art-1", persist=False) + out = _tool().execute_action("extract_document", input=_ART_ID, persist=False) assert out["status"] == "ok" table = out["structured"]["tables"][0] # Rows are capped and flagged with the original total. @@ -572,7 +616,7 @@ def test_oversized_input_is_rejected_before_exec(monkeypatch): manager = _manager_with_extract(_CANNED_EXTRACT) _patch_manager(monkeypatch, manager) - out = _tool().execute_action("extract_document", input="art-1", persist=False) + out = _tool().execute_action("extract_document", input=_ART_ID, persist=False) assert out["status"] == "error" assert "too large" in out["error"] # Nothing was staged or executed for an oversized input. diff --git a/tests/integration/test_artifact_generator_e2e.py b/tests/integration/test_artifact_generator_e2e.py index 20f8b54c..3bcc12bc 100644 --- a/tests/integration/test_artifact_generator_e2e.py +++ b/tests/integration/test_artifact_generator_e2e.py @@ -240,6 +240,71 @@ def test_edit_cross_tenant_denied(wired_tool): assert artifact["current_version"] == 1 +def test_short_ref_numbering_and_edit_by_ref(wired_tool): + tool, conversation_id, pg_engine, storage = wired_tool + from pptx import Presentation + + # Successive creates in one conversation get A1, A2 (1-based, by created order). + first = tool.execute_action( + "create_artifact", kind="presentation", title="One", spec={"slides": [{"title": "a"}]} + ) + assert first["status"] == "ok", first + assert first["ref"] == "A1" + second = tool.execute_action( + "create_artifact", kind="document", spec={"sections": [{"heading": "h", "paragraphs": ["p"]}]} + ) + assert second["status"] == "ok", second + assert second["ref"] == "A2" + + # edit_artifact(id="A1") resolves the ref -> appends v2 to the FIRST artifact only. + edited = tool.execute_action("edit_artifact", id="A1", spec_patch={"slides": [{"title": "a"}, {"title": "b"}]}) + assert edited["status"] == "ok", edited + assert edited["version"] == 2 + assert edited["artifact_id"] == first["artifact_id"] + # The append keeps the artifact's position, so its ref is stable. + assert edited["ref"] == "A1" + + with pg_engine.connect() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact_in_parent(first["artifact_id"], conversation_id=conversation_id) + v1 = repo.get_version(first["artifact_id"], 1) + v2 = repo.get_version(first["artifact_id"], 2) + assert artifact["current_version"] == 2 + # Append-only: v1 preserved with one slide, v2 has two. + assert len(v1["spec"]["slides"]) == 1 + assert len(v2["spec"]["slides"]) == 2 + assert len(Presentation(storage.get_file(v2["storage_path"])).slides) == 2 + + # The uuid path still works alongside refs. + by_uuid = tool.execute_action("edit_artifact", id=first["artifact_id"], spec_patch={"title": "X"}) + assert by_uuid["status"] == "ok" + assert by_uuid["version"] == 3 + + +def test_short_ref_is_conversation_scoped(wired_tool): + tool, conversation_id, pg_engine, storage = wired_tool + + # Seed an artifact in a DIFFERENT conversation; its A1 must not resolve here. + other_conversation = str(uuid.uuid4()) + _seed_presentation(pg_engine, storage, other_conversation) + + # This conversation has no artifacts yet, so A1 resolves to nothing (no cross-parent leak). + out = tool.execute_action("edit_artifact", id="A1", spec_patch={"title": "x"}) + assert out["status"] == "error" + assert "not found in this conversation/run" in out["error"] + + +def test_short_ref_out_of_range_is_clean_error(wired_tool): + tool, conversation_id, pg_engine, storage = wired_tool + + created = tool.execute_action("create_artifact", kind="presentation", spec={"slides": [{"title": "a"}]}) + assert created["ref"] == "A1" + # A2 has no artifact behind it -> clean not-found, no crash. + out = tool.execute_action("edit_artifact", id="A2", spec_patch={"title": "x"}) + assert out["status"] == "error" + assert "not found in this conversation/run" in out["error"] + + def _seed_presentation(pg_engine, storage, conversation_id) -> str: """Create a presentation artifact (row + version + bytes) in a given conversation; return its id.""" spec = {"slides": [{"title": "seed"}]} diff --git a/tests/integration/test_workflow_code_node_e2e.py b/tests/integration/test_workflow_code_node_e2e.py index ae7075b2..888c8124 100644 --- a/tests/integration/test_workflow_code_node_e2e.py +++ b/tests/integration/test_workflow_code_node_e2e.py @@ -153,11 +153,13 @@ def test_code_node_persists_artifact_reference_in_state(wired): list(engine._execute_code_node(node)) ref = engine.state["report"] - # The state holds an artifact REFERENCE (JSON primitives), never bytes. - assert set(ref) == {"artifact_id", "version", "filename", "mime_type", "size"} + # The state holds an artifact REFERENCE (JSON primitives), never bytes. The + # short handle ``ref`` (A1) lets a later node address the artifact by position. + assert set(ref) == {"artifact_id", "version", "filename", "mime_type", "size", "ref"} assert ref["filename"] == "report.txt" assert ref["mime_type"] == "text/plain" assert ref["size"] == len(b"compliance ok") + assert ref["ref"] == "A1" assert all(not isinstance(v, (bytes, bytearray)) for v in ref.values()) assert engine.state["node_code_1_output"] == ref diff --git a/tests/storage/db/repositories/test_artifacts.py b/tests/storage/db/repositories/test_artifacts.py index 2d2f15a4..75ee3905 100644 --- a/tests/storage/db/repositories/test_artifacts.py +++ b/tests/storage/db/repositories/test_artifacts.py @@ -322,6 +322,91 @@ class TestQuotaEnforcement: ac._enforce_user_quota(_repo(pg_conn), "alice", 10_000_000, new_artifact=True) # disabled -> no raise +class TestVirtualRefPositions: + """Parent-scoped 1-based positions that back the short ``A{n}`` refs.""" + + @staticmethod + def _stamp(pg_conn, artifact_id: str, seconds: int) -> None: + """Give an artifact a distinct created_at so ordering is deterministic in one txn. + + In production each persist is its own transaction (distinct ``now()``); the + repo test seeds several rows in a single transaction where ``now()`` collides, + so explicit timestamps make the created_at ordering testable. + """ + pg_conn.execute( + text( + "UPDATE artifacts SET created_at = TIMESTAMPTZ '2026-01-01 00:00:00+00' " + "+ (:s || ' seconds')::interval WHERE id = CAST(:id AS uuid)" + ), + {"s": seconds, "id": artifact_id}, + ) + + def test_position_and_id_round_trip_in_order(self, pg_conn): + repo = _repo(pg_conn) + conv = _conversation_id() + a = repo.create_artifact("u", "document", conversation_id=conv) + b = repo.create_artifact("u", "document", conversation_id=conv) + c = repo.create_artifact("u", "document", conversation_id=conv) + for idx, art in enumerate((a, b, c)): + self._stamp(pg_conn, art["id"], idx) + + assert repo.position_in_parent(a["id"], conversation_id=conv) == 1 + assert repo.position_in_parent(b["id"], conversation_id=conv) == 2 + assert repo.position_in_parent(c["id"], conversation_id=conv) == 3 + + assert repo.artifact_id_at_position(1, conversation_id=conv) == a["id"] + assert repo.artifact_id_at_position(2, conversation_id=conv) == b["id"] + assert repo.artifact_id_at_position(3, conversation_id=conv) == c["id"] + + def test_position_is_parent_scoped(self, pg_conn): + repo = _repo(pg_conn) + conv_a = _conversation_id() + conv_b = _conversation_id() + a1 = repo.create_artifact("u", "document", conversation_id=conv_a) + a2 = repo.create_artifact("u", "document", conversation_id=conv_a) + b1 = repo.create_artifact("u", "document", conversation_id=conv_b) + for idx, art in enumerate((a1, a2, b1)): + self._stamp(pg_conn, art["id"], idx) + + # Each parent numbers from 1 independently. + assert repo.position_in_parent(a2["id"], conversation_id=conv_a) == 2 + assert repo.position_in_parent(b1["id"], conversation_id=conv_b) == 1 + # An artifact is invisible (position 0) under the wrong parent. + assert repo.position_in_parent(a2["id"], conversation_id=conv_b) == 0 + + def test_id_at_position_does_not_cross_parents(self, pg_conn): + repo = _repo(pg_conn) + conv_a = _conversation_id() + conv_b = _conversation_id() + a1 = repo.create_artifact("u", "document", conversation_id=conv_a) + # conv_b has no artifacts, so its A1 resolves to nothing (no leak from conv_a). + assert repo.artifact_id_at_position(1, conversation_id=conv_a) == a1["id"] + assert repo.artifact_id_at_position(1, conversation_id=conv_b) is None + + def test_out_of_range_and_invalid_position_return_none(self, pg_conn): + repo = _repo(pg_conn) + conv = _conversation_id() + repo.create_artifact("u", "document", conversation_id=conv) + assert repo.artifact_id_at_position(2, conversation_id=conv) is None + assert repo.artifact_id_at_position(0, conversation_id=conv) is None + assert repo.artifact_id_at_position(-1, conversation_id=conv) is None + + def test_workflow_run_parent_supported(self, pg_conn): + repo = _repo(pg_conn) + run = _conversation_id() + a = repo.create_artifact("u", "document", workflow_run_id=run) + assert repo.position_in_parent(a["id"], workflow_run_id=run) == 1 + assert repo.artifact_id_at_position(1, workflow_run_id=run) == a["id"] + + def test_position_helpers_require_a_parent(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create_artifact("u", "document", conversation_id=_conversation_id()) + with pytest.raises(ValueError): + repo.position_in_parent(created["id"]) + with pytest.raises(ValueError): + repo.artifact_id_at_position(1) + + class TestCascadeDelete: def test_deleting_artifact_removes_versions(self, pg_conn): repo = _repo(pg_conn) diff --git a/tests/test_code_executor_tool.py b/tests/test_code_executor_tool.py index 624bce59..2101a31f 100644 --- a/tests/test_code_executor_tool.py +++ b/tests/test_code_executor_tool.py @@ -7,6 +7,8 @@ run in the fast unit suite. The end-to-end persistence path is covered by from __future__ import annotations +import uuid + from application.agents.tools.code_executor import ( CodeExecutorTool, _infer_mime, @@ -266,3 +268,84 @@ def test_session_kept_alive_on_positive_ttl(monkeypatch): monkeypatch, manager, code="print(1)", ttl=30, capture_artifacts=False ) assert manager.closed == [] + + +# --------------------------------------------------------------------------- +# Input materialization: short-ref + uuid resolution (no live sandbox/DB) +# --------------------------------------------------------------------------- +_ART_ID = str(uuid.uuid4()) + + +class _InputManager: + """Records files staged into the workspace by _materialize_inputs.""" + + def __init__(self) -> None: + self.put_files: dict = {} + + def put_file(self, session_id, dest_path, data): + self.put_files[dest_path] = data + + +def _patch_input_repo(monkeypatch, *, found_position: bool, conv: str): + """Patch db_readonly + ArtifactsRepository so a ref/uuid resolves only within ``conv``.""" + from application.agents.tools import code_executor as ce + + class _Repo: + def __init__(self, conn): + pass + + def artifact_id_at_position(self, n, *, conversation_id=None, workflow_run_id=None): + if not found_position or n != 1 or conversation_id != conv: + return None + return _ART_ID + + def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None): + if conversation_id != conv: + return None + return {"id": artifact_id, "current_version": 1, "title": "seed.csv"} + + def get_version(self, artifact_id, version): + return {"filename": "seed.csv", "storage_path": f"inputs/u/artifacts/{artifact_id}/v1/seed.csv"} + + class _Conn: + def __enter__(self): + return object() + + def __exit__(self, *exc): + return False + + class _Storage: + def get_file(self, path): + import io + + return io.BytesIO(b"col\n1\n") + + monkeypatch.setattr(ce, "db_readonly", lambda: _Conn()) + monkeypatch.setattr(ce, "ArtifactsRepository", _Repo) + monkeypatch.setattr(ce.StorageCreator, "get_storage", staticmethod(lambda: _Storage())) + + +def test_materialize_inputs_accepts_short_ref(monkeypatch): + _patch_input_repo(monkeypatch, found_position=True, conv="conv-1") + manager = _InputManager() + out = _tool()._materialize_inputs(manager, "conv-1", ["A1"]) + assert "error" not in out + assert out["loaded"] == ["inputs/seed.csv"] + assert manager.put_files["inputs/seed.csv"] == b"col\n1\n" + + +def test_materialize_inputs_accepts_uuid(monkeypatch): + _patch_input_repo(monkeypatch, found_position=False, conv="conv-1") + manager = _InputManager() + out = _tool()._materialize_inputs(manager, "conv-1", [_ART_ID]) + assert "error" not in out + assert out["loaded"] == ["inputs/seed.csv"] + + +def test_materialize_inputs_out_of_range_ref_is_clean_error(monkeypatch): + _patch_input_repo(monkeypatch, found_position=True, conv="conv-1") + manager = _InputManager() + out = _tool()._materialize_inputs(manager, "conv-1", ["A2"]) + assert "A2" in out["error"] + assert "not found in this conversation/run" in out["error"] + assert manager.put_files == {}