Add short artifact handle (A1) so the model can edit by reference

Give each produced artifact a short virtual handle - A1, A2, ... - the n-th
artifact in the conversation or workflow run (case-insensitive, not stored). The
tools return it in their results, and the edit, rewrite, and input parameters
accept either the handle or the full id. A handle resolves only within the
caller's own conversation or run, and the resolved id is still checked against
that parent before any read, so it cannot reach another tenant's artifact. This
fixes the model creating a duplicate instead of a new version when asked to
edit an artifact.
This commit is contained in:
Alex committed 2026-06-24 15:35:06 +01:00
1 parent a06b348c34
commit e7f765b628
13 files changed
+655 -61

No files matched your search

+38 -18
View File
@@ -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]:
+51
View File
@@ -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
+23 -6
View File
@@ -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")
+25 -9
View File
@@ -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]:
@@ -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.")
+39 -2
View File
@@ -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
@@ -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(
+104
View File
@@ -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")]
@@ -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/<token>) 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.
@@ -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"}]}
@@ -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
@@ -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)
+83
View File
@@ -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 == {}