fix: more artefact hardening

This commit is contained in:
Alex committed 2026-07-04 11:42:27 +02:00
1 parent 47cc0314af
commit 94a845aa82
37 files changed
+1233 -165

No files matched your search

+19 -3
View File
@@ -1,8 +1,11 @@
"""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.
A ref is NOT a stored column: ``A{n}`` is the artifact's STABLE per-parent ``ref_seq``, assigned at
creation and kept in the artifact's ``metadata``, so deleting an earlier artifact no longer
re-points a later ref the model already holds. Artifacts created before ``ref_seq`` existed have
none, so resolution falls back to the legacy positional (n-th by created_at) lookup. Refs resolve
only inside the caller's parent (``conversation_id`` or ``workflow_run_id``), never cross-tenant;
resolution still goes through the parent-scoped authz gate.
"""
from __future__ import annotations
@@ -41,6 +44,19 @@ def resolve_artifact_id(
"""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:
# A ref is the artifact's stable per-parent ``ref_seq``: resolve by it first so a
# deletion of an earlier artifact never re-points this ref. Legacy rows (created
# before ref_seq) and repos without the newer method fall back to the positional
# (n-th by created_at) lookup.
by_seq = getattr(repo, "resolve_id_by_ref_seq", None)
if callable(by_seq):
resolved = by_seq(
position,
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
)
if resolved is not None:
return resolved
return repo.artifact_id_at_position(
position,
conversation_id=conversation_id,
+7 -2
View File
@@ -20,6 +20,7 @@ from application.sandbox.artifacts_capture import (
infer_mime as _infer_mime,
kind_for_mime as _kind_for_mime,
snapshot_signatures,
unique_input_path,
)
from application.sandbox.base import ExecResult
from application.sandbox.sandbox_creator import SandboxCreator
@@ -224,6 +225,9 @@ class CodeExecutorTool(Tool):
if not inputs:
return {"loaded": loaded}
storage = StorageCreator.get_storage()
# Two inputs whose current versions share a filename would clobber each other at
# the same ``inputs/{name}`` path; track used paths and disambiguate deterministically.
used_paths: set = set()
for raw_id in inputs:
raw = str(raw_id).strip()
if not raw:
@@ -295,12 +299,13 @@ class CodeExecutorTool(Tool):
return {"error": f"failed to read input artifact {artifact_id}."}
if max_bytes and len(data) > max_bytes:
return {"error": f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit."}
rel_path = unique_input_path(f"inputs/{filename}", used_paths)
try:
manager.put_file(session_id, f"inputs/{filename}", data)
manager.put_file(session_id, rel_path, data)
except Exception:
logger.exception("code_executor: put_file failed for input artifact")
return {"error": f"failed to stage input artifact {artifact_id} into the workspace."}
loaded.append(f"inputs/{filename}")
loaded.append(rel_path)
return {"loaded": loaded}
def _bridge_chat_attachment(self, raw: str) -> Any:
+28 -13
View File
@@ -14,6 +14,8 @@ from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional
from celery import current_task
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.agents.tools.attachment_bridge import (
AttachmentBridgeError,
@@ -188,31 +190,44 @@ class ReadDocumentTool(Tool):
return result
def _dispatch(self, artifact_id: str, options: Dict[str, Any]) -> Dict[str, Any]:
"""Enqueue ``parse_document`` on the dedicated parsing queue and await with a timeout.
"""Parse INLINE inside a Celery worker, else dispatch to the parsing queue and await.
Deadlock note: this tool may run in the WEB process (/stream) OR inside a
Celery worker (headless/scheduled agents). Routing to the dedicated
``parsing`` queue + separate parsing workers avoids worker-self-deadlock;
the await degrades (returns an error) on timeout/failure rather than hanging.
This tool runs in the WEB process (/stream) OR inside a Celery worker
(headless/scheduled/workflow agents). When it already runs inside a worker that also
serves the ``parsing`` queue (the shipped default ``-Q docsgpt,parsing``), dispatching
and blocking on ``get()`` would self-deadlock: concurrent agent tasks each hold a pool
slot blocked in ``get()`` so ``parse_document`` never gets a free slot. So parse INLINE
in-process inside a worker; only dispatch+await (degrading on timeout/failure) from web.
"""
parent = self._parent()
# ``current_task`` is a Celery proxy: truthy only while this runs inside a worker task,
# falsy in the web process (the bare proxy is NOT identity-None, so test truthiness).
if current_task:
from application.worker import run_parse_document
try:
result = run_parse_document(artifact_id, parent, self.user_id, options)
except Exception as exc:
logger.exception("read_document: inline parse failed")
return {"status": "error", "error": f"document parsing failed: {type(exc).__name__}: {exc}"}
if not isinstance(result, dict):
return {"status": "error", "error": "document parsing produced an unexpected result."}
return result
from celery.exceptions import TimeoutError as CeleryTimeoutError
from application.api.user.tasks import parse_document
parent = self._parent()
timeout = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120))
queue = getattr(settings, "DOCUMENT_PARSE_QUEUE", "parsing")
try:
async_result = parse_document.apply_async(
args=[artifact_id, parent, self.user_id, options], queue=queue
)
# ``read_document`` runs both from the web request path AND from inside a
# Celery task (headless/scheduled agents). In a prefork worker
# ``task_join_will_block()`` is process-wide, so the default
# ``disable_sync_subtasks=True`` makes ``get()`` raise RuntimeError
# ("Never call result.get() within a task!"). The dedicated ``parsing``
# queue + separate workers already avoid the real self-deadlock, so opt
# out of the blanket guard explicitly.
# The web process (not a worker) awaits here; ``disable_sync_subtasks=False`` keeps
# the call correct if invoked from a non-prefork (eventlet/gevent) worker where the
# inline branch above still ran but the blanket guard would otherwise raise.
result = async_result.get(timeout=timeout, disable_sync_subtasks=False)
except (CeleryTimeoutError, TimeoutError):
return {"status": "error", "error": f"document parsing timed out after {int(timeout)}s."}
@@ -345,6 +345,16 @@ class WorkflowEngine:
has_structured_response = False
first_chunk = True
for event in node_agent.gen(formatted_prompt):
# A tool that pauses for approval makes the LLM handler yield
# ``tool_calls_pending`` and end. An ephemeral node agent has no resume path,
# so silently continuing would leave the node with empty output (or a confusing
# "Structured output was expected" when it has a json_schema). Fail visibly.
if event.get("type") == "tool_calls_pending":
raise ValueError(
f'Node "{node.title}" uses a tool that requires approval, which is not '
"supported inside a workflow. Disable require_approval for tools used in "
"workflow nodes."
)
if "answer" in event:
chunk = str(event["answer"])
full_response_parts.append(chunk)
@@ -496,6 +506,7 @@ class WorkflowEngine:
"""Stage referenced input artifacts (run-scoped, never cross-tenant) into the workspace."""
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.core.settings import settings
from application.sandbox.artifacts_capture import unique_input_path
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
from application.storage.storage_creator import StorageCreator
@@ -507,6 +518,9 @@ class WorkflowEngine:
return loaded
max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 0) or 0)
storage = StorageCreator.get_storage()
# Two inputs whose current versions share a filename would clobber each other at the
# same ``inputs/{name}`` path; track used paths and disambiguate deterministically.
used_paths: set = set()
for raw in raw_ids:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
@@ -541,8 +555,9 @@ class WorkflowEngine:
raise ValueError(
f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit."
)
manager.put_file(session_id, f"inputs/{filename}", data)
loaded.append(f"inputs/{filename}")
rel_path = unique_input_path(f"inputs/{filename}", used_paths)
manager.put_file(session_id, rel_path, data)
loaded.append(rel_path)
return loaded
def _materialize_node_attachments(
+9
View File
@@ -384,6 +384,15 @@ class DownloadArtifact(Resource):
),
500,
)
# A 302 to a cross-origin S3 URL can't be read by the app's authed
# fetch (the bucket has no CORS grant for the app origin). When the
# client opts in via ?disposition=url (or Accept: application/json)
# hand the presigned URL back as JSON so it can navigate to it
# top-level (no CORS). Default stays a 302 so nothing else breaks.
if request.args.get("disposition") == "url" or (
"application/json" in request.headers.get("Accept", "")
):
return make_response(jsonify({"success": True, "url": url}), 200)
return redirect(url, code=302)
# Stream the bytes in chunks instead of buffering the whole object in
+29 -21
View File
@@ -22,10 +22,10 @@ from application.utils import safe_filename
logger = logging.getLogger(__name__)
# Cap the text returned to the LLM so a huge document can't flood context; the
# full result is still persisted as a ``data`` artifact. When the text exceeds
# the cap a head+tail window keeps both the document's beginning AND end (e.g.
# totals/signatures) within the byte budget.
# Default cap for the LLM-facing VIEW of the text (applied in ``bound_parse_payload``,
# NOT during parsing) so a huge document can't flood context; the full result is still
# persisted as a ``data`` artifact. When the text exceeds the cap a head+tail window keeps
# both the document's beginning AND end (e.g. totals/signatures) within the byte budget.
_TEXT_MAX_BYTES = 8000
_MAX_TABLES_RETURNED = 20
_MAX_TABLE_ROWS = 50
@@ -55,18 +55,26 @@ def truncate_text_head_tail(text: str, max_bytes: Optional[int] = None) -> str:
return f"{head_text}\n\n...[truncated {dropped} bytes]...\n\n{tail_text}"
def bound_parse_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
def bound_parse_payload(payload: Dict[str, Any], max_chars: Optional[int] = None) -> Dict[str, Any]:
"""Bound every shape of a parse payload so the Redis-backed result stays small.
``content`` is re-windowed and ``chunks`` is capped in count and per-chunk
length. ``structured`` is left as-is: it rides back so json_schema validation
in the tool can run against it, and it is already bounded by the input byte
cap plus the table caps (``_compact_table`` / ``summary``); the full result is
also persisted as a ``data`` artifact. The dict is mutated in place.
This is where ALL view-bounding happens: parsing now returns the FULL content and the
persisted ``data`` artifact keeps it, while the view ridden back through the Redis result
backend is bounded here. ``content`` is capped to ``max_chars`` when given, else re-windowed
to a head+tail byte window; ``chunks`` is capped in count and per-chunk length. ``structured``
is left as-is: it rides back so json_schema validation in the tool can run against it, and it
is already bounded by the input byte cap plus the table caps (``_compact_table`` /
``summary``). ``payload['truncated']`` is set when the content view actually cut. The dict is
mutated in place.
"""
content = payload.get("content")
if isinstance(content, str):
payload["content"] = truncate_text_head_tail(content)
if max_chars and int(max_chars) > 0:
capped = content[: int(max_chars)]
else:
capped = truncate_text_head_tail(content)
payload["content"] = capped
payload["truncated"] = capped != content
chunks = payload.get("chunks")
if isinstance(chunks, list):
@@ -336,7 +344,7 @@ def _shape(
extracted = _docling_structured(path, ocr_enabled=ocr_enabled, include_tables=include_tables)
except Exception as exc:
return {"error": f"structured parsing requires Docling: {type(exc).__name__}: {exc}"}
bounded, truncated = _bounded(extracted["markdown"], max_chars)
bounded, truncated = _bounded(extracted["markdown"])
return {
"output": "structured",
"content": bounded,
@@ -364,7 +372,7 @@ def _shape(
except Exception:
text, tables = _parse_to_text(parser, path), []
text = _apply_pages(text, pages)
bounded, truncated = _bounded(text, max_chars)
bounded, truncated = _bounded(text)
payload: Dict[str, Any] = {"output": output, "content": bounded, "truncated": truncated}
if tables:
payload["tables"] = tables
@@ -388,17 +396,17 @@ def _shape(
path, ocr_enabled=ocr_enabled, include_tables=True)["tables"]]
except Exception:
tables = []
bounded, truncated = _bounded(text, max_chars)
bounded, truncated = _bounded(text)
payload: Dict[str, Any] = {"output": output, "content": bounded, "truncated": truncated}
if tables:
payload["tables"] = tables
return payload
def _bounded(text: str, max_chars: Optional[int]) -> tuple[str, bool]:
"""Bound text to ``max_chars`` (chars) or the default byte window; flag truncation."""
if max_chars and int(max_chars) > 0:
capped = text[: int(max_chars)]
return capped, len(capped) < len(text)
bounded = truncate_text_head_tail(text)
return bounded, bounded != text
def _bounded(text: str) -> tuple[str, bool]:
"""Return the FULL extracted text (never truncated here); the view is bounded in ``bound_parse_payload``.
Parsing keeps the complete text so the persisted ``data`` artifact is the full parse;
``max_chars`` and the default head+tail window now bound only the LLM-facing view.
"""
return text, False
+56 -4
View File
@@ -13,7 +13,8 @@ import hashlib
import io
import logging
import mimetypes
from typing import Any, Dict, List, Optional, Tuple
import os
from typing import Any, Dict, List, Optional, Set, Tuple
from sqlalchemy import text
@@ -63,6 +64,25 @@ _KIND_BY_MIME_PREFIX: Dict[str, str] = {
}
def unique_input_path(rel_path: str, used: Set[str]) -> str:
"""Return ``rel_path`` (reserving it) or a ``-2``/``-3``... suffixed variant when already staged.
Two inputs whose current versions share a filename would otherwise clobber each other at the
same ``inputs/{name}`` path; the numeric suffix is inserted before the extension. Mutates
``used`` to record the returned path.
"""
if rel_path not in used:
used.add(rel_path)
return rel_path
base, ext = os.path.splitext(rel_path)
n = 2
while f"{base}-{n}{ext}" in used:
n += 1
unique = f"{base}-{n}{ext}"
used.add(unique)
return unique
def infer_mime(filename: str) -> str:
"""Infer a mime type from a filename, falling back to a generic binary type."""
mime, _ = mimetypes.guess_type(filename)
@@ -271,16 +291,36 @@ 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_from_metadata(metadata: Any) -> Optional[str]:
"""Return the short ref (``A{ref_seq}``) from a stored metadata dict, or None when absent/non-numeric."""
if not isinstance(metadata, dict):
return None
try:
seq = int(metadata.get("ref_seq"))
except (TypeError, ValueError):
return None
return make_ref(seq) if seq >= 1 else None
def _ref_for(
repo: ArtifactsRepository,
artifact_id: str,
*,
conversation_id: Optional[str],
workflow_run_id: Optional[str],
metadata: Any = None,
) -> Optional[str]:
"""Compute the short ref (``A{n}``) for an artifact from its position in its parent; None on failure."""
"""Compute the short ref for an artifact: its STABLE stored ``ref_seq``, else positional fallback.
The ref_seq is assigned at creation and kept in ``metadata`` so it survives deletions of
earlier artifacts. Legacy rows created before ref_seq existed have none, so they fall back to
the (mutable) position-in-parent; ``resolve_artifact_id`` mirrors this fallback.
"""
if conversation_id is None and workflow_run_id is None:
return None
seq_ref = _ref_from_metadata(metadata)
if seq_ref is not None:
return seq_ref
try:
position = repo.position_in_parent(
artifact_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
@@ -339,8 +379,13 @@ 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)
# The freshly-created row carries the stable ``ref_seq`` in its metadata.
ref = _ref_for(
repo, artifact_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
repo,
artifact_id,
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
metadata=artifact.get("metadata"),
)
storage.save_file(io.BytesIO(data), storage_path)
saved_key = storage_path
@@ -402,8 +447,15 @@ 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)
# The identity's stable ``ref_seq`` was assigned at creation; read it so an
# edit hands back the SAME ref the model already holds.
existing = repo.get_artifact(str(artifact_id))
ref = _ref_for(
repo, str(artifact_id), conversation_id=conversation_id, workflow_run_id=workflow_run_id
repo,
str(artifact_id),
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
metadata=(existing or {}).get("metadata"),
)
storage.save_file(io.BytesIO(data), storage_path)
saved_key = storage_path
@@ -57,7 +57,21 @@ class ArtifactsRepository:
preview_text: Optional[str] = None,
produced_by: Any = None,
) -> dict:
"""Create the identity row and its version 1 atomically; return the artifact dict."""
"""Create the identity row and its version 1 atomically; return the artifact dict.
When a parent is given, a STABLE per-parent ``ref_seq`` (max existing + 1) is
computed BEFORE the insert and stored into ``metadata`` so the short ``A{n}`` ref
the model receives survives later deletions. See ``next_ref_seq`` for the (benign)
concurrent-create dup caveat.
"""
metadata_payload = metadata
if conversation_id is not None or workflow_run_id is not None:
ref_seq = self.next_ref_seq(
conversation_id=conversation_id, workflow_run_id=workflow_run_id
)
merged = dict(metadata) if isinstance(metadata, dict) else {}
merged["ref_seq"] = ref_seq
metadata_payload = merged
artifact = self._conn.execute(
text(
"""
@@ -84,7 +98,7 @@ class ArtifactsRepository:
"message_id": message_id,
"kind": kind,
"title": title,
"metadata": json.dumps(metadata) if metadata is not None else None,
"metadata": json.dumps(metadata_payload) if metadata_payload is not None else None,
},
).fetchone()
artifact_dict = _artifact_to_dict(artifact)
@@ -273,6 +287,61 @@ class ArtifactsRepository:
).fetchone()
return str(row[0]) if row is not None else None
def next_ref_seq(
self,
*,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
) -> int:
"""Return the next stable ``ref_seq`` for a parent (max stored + 1; 1 when empty).
Caveat: two concurrent creates in one parent can read the same MAX and get a
duplicate ``ref_seq`` (creates within a parent are normally serial from one agent
turn); a dup resolves to the earliest — acceptable, and strictly better than the
re-point bug. A DB unique index is a possible future hardening (no migration here).
"""
if conversation_id is None and workflow_run_id is None:
raise ValueError("next_ref_seq requires conversation_id or workflow_run_id")
clauses, params = self._parent_clauses(conversation_id, workflow_run_id)
# Guard the cast: legacy rows have no ``ref_seq`` (NULL, ignored), and the regex
# match keeps any non-numeric value from erroring the cast.
row = self._conn.execute(
text(
"SELECT COALESCE(MAX(CASE WHEN metadata ->> 'ref_seq' ~ '^[0-9]+$' "
"THEN (metadata ->> 'ref_seq')::int END), 0) + 1 "
f"FROM artifacts WHERE {' AND '.join(clauses)}"
),
params,
).fetchone()
return int(row[0]) if row is not None else 1
def resolve_id_by_ref_seq(
self,
seq: int,
*,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
) -> Optional[str]:
"""Return the id of the artifact whose stored ``ref_seq`` == ``seq`` in a parent, or None.
Earliest ``created_at`` wins on the rare concurrent-create tie (see ``next_ref_seq``).
"""
if conversation_id is None and workflow_run_id is None:
raise ValueError("resolve_id_by_ref_seq requires conversation_id or workflow_run_id")
if not isinstance(seq, int) or seq < 1:
return None
clauses, params = self._parent_clauses(conversation_id, workflow_run_id)
params["seq"] = str(seq)
row = self._conn.execute(
text(
f"SELECT id FROM artifacts WHERE {' AND '.join(clauses)} "
"AND metadata ->> 'ref_seq' = :seq "
"ORDER BY created_at ASC, id ASC 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 = ""
+22 -14
View File
@@ -1611,11 +1611,18 @@ def attachment_worker(self, file_info, user):
def parse_document_worker(self, artifact_id, parent, user_id, options):
"""Parse an input artifact's bytes to a shaped result on the dedicated parsing queue.
"""Thin Celery-task wrapper; delegates to the process-agnostic ``run_parse_document``."""
return run_parse_document(artifact_id, parent, user_id, options)
Security: the artifact is re-resolved through the run-scoped gate IN THE WORKER
(never trusting a raw storage path) so authz is enforced independently here, in
addition to the pre-enqueue check in the tool.
def run_parse_document(artifact_id, parent, user_id, options):
"""Parse an input artifact's bytes to a shaped result; runnable inline OR on the parsing queue.
Security: the artifact is re-resolved through the run-scoped gate here (never trusting a
raw storage path) so authz is enforced independently, in addition to the pre-enqueue check
in the tool. ``read_document`` calls this directly (in-process) when it already runs inside a
Celery worker; the web process dispatches ``parse_document`` to the parsing queue, which lands
here via ``parse_document_worker``.
"""
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.parser.document_reader import bound_parse_payload, parse_document_bytes
@@ -1628,7 +1635,7 @@ def parse_document_worker(self, artifact_id, parent, user_id, options):
return {"status": "error", "error": "parse_document requires a conversation_id or workflow_run_id."}
# Re-resolve through the parent-scoped gate so a forged/cross-run id is rejected
# in the worker too; resolve a short ref to an id within this parent only.
# here too; resolve a short ref to an id within this parent only.
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
@@ -1646,7 +1653,7 @@ def parse_document_worker(self, artifact_id, parent, user_id, options):
return {"status": "error", "error": f"input artifact {artifact_id} not found in this conversation/run."}
version = repo.get_version(resolved_id, artifact["current_version"])
except Exception:
logging.error("parse_document_worker: failed to resolve input artifact", exc_info=True)
logging.error("run_parse_document: failed to resolve input artifact", exc_info=True)
return {"status": "error", "error": f"failed to load input artifact {artifact_id}."}
if not version or not version.get("storage_path"):
@@ -1657,9 +1664,11 @@ def parse_document_worker(self, artifact_id, parent, user_id, options):
try:
data = StorageCreator.get_storage().get_file(version["storage_path"]).read()
except Exception:
logging.error("parse_document_worker: failed to read input artifact bytes", exc_info=True)
logging.error("run_parse_document: failed to read input artifact bytes", exc_info=True)
return {"status": "error", "error": f"failed to read input artifact {artifact_id}."}
# Parse returns the FULL content (no max_chars/window here) so the persisted artifact is the
# complete parse; all view-bounding is applied by ``bound_parse_payload`` below.
result = parse_document_bytes(
data,
filename,
@@ -1667,7 +1676,6 @@ def parse_document_worker(self, artifact_id, parent, user_id, options):
ocr=options.get("ocr", "auto"),
pages=options.get("pages"),
engine=options.get("engine", "auto"),
max_chars=options.get("max_chars"),
include_tables=bool(options.get("include_tables", True)),
)
if result.get("error"):
@@ -1675,17 +1683,17 @@ def parse_document_worker(self, artifact_id, parent, user_id, options):
payload = {"status": "ok", **result}
if options.get("persist"):
# The full shaped result is persisted by reference; only a bounded view rides
# back through the Redis result backend (the bytes live in the artifact).
# Persist the FULL shaped result by reference (bytes live in the artifact); only the
# bounded view computed below rides back through the Redis result backend.
artifact_ref = _persist_parse_result(result, display_name, user_id, parent, options)
if isinstance(artifact_ref, dict) and artifact_ref.get("error"):
payload["artifact_error"] = artifact_ref["error"]
elif artifact_ref is not None:
payload["artifact"] = artifact_ref
# Bound the Redis-backed view across all shapes: content is re-windowed, chunks are
# capped, and structured (needed for json_schema validation) is bounded by the input
# cap + table caps. The FULL result already lives in the persisted artifact above.
payload = bound_parse_payload(payload)
# Bound the Redis-backed VIEW across all shapes: content is capped by max_chars (else the
# default head+tail window), chunks are count/length-capped, and structured (needed for
# json_schema validation) rides back as-is. The FULL result already lives in the artifact.
payload = bound_parse_payload(payload, max_chars=options.get("max_chars"))
return payload
@@ -34,6 +34,10 @@ spec:
# runtimeClassName: gvisor
securityContext:
runAsNonRoot: true
# Numeric UID the kubelet can verify as non-root; MUST match the sandbox
# image's `USER` (deployment/sandbox/Dockerfile). Without it the pod
# fails to start with CreateContainerConfigError.
runAsUser: 10001
seccompProfile:
type: RuntimeDefault
containers:
@@ -40,6 +40,7 @@ spec:
egress:
# DNS to cluster resolvers (kube-dns). Restricted to UDP/TCP 53 so the
# broad-egress rule below does not need to whitelist the resolver IP.
# Kept for clusters WITHOUT NodeLocal DNSCache, where the resolver is a pod.
- to:
- namespaceSelector: {}
ports:
@@ -47,6 +48,22 @@ spec:
port: 53
- protocol: TCP
port: 53
# NodeLocal DNSCache: on clusters running it, the stub resolver is a
# link-local IP (commonly 169.254.20.10), NOT a pod, and it sits inside the
# 169.254.0.0/16 range carved out of the broad-egress rule below -- so
# without this explicit allow ALL DNS is dropped and every hostname lookup
# (pip install, public APIs) times out. NetworkPolicy rules are
# additive-allow, so this narrowly re-permits port 53 to that one IP.
# Operators whose NodeLocal DNSCache uses a different link-local IP must
# adjust this cidr.
- to:
- ipBlock:
cidr: 169.254.20.10/32 # NodeLocal DNSCache link-local resolver
ports:
- protocol: UDP
port: 53
- protocol: TCP
port: 53
# Broad public-internet egress, with every private / link-local / ULA /
# carrier-grade-NAT / documentation range carved out via `except`. The
# cloud metadata IP 169.254.169.254 is inside the 169.254.0.0/16 hole, so
@@ -11,31 +11,68 @@
# (control plane to backend/worker) and `sandbox-egress` (its internet route).
# This overlay flips `sandbox-egress` to `internal: true`, which removes the
# runner's direct route to the internet / host / RFC1918 / metadata entirely --
# so even raw sockets in arbitrary sandbox code have nowhere to go except the
# proxy on `sandbox-net`. The proxy is the ONLY container with an internet route
# so raw sockets in arbitrary sandbox code have no route OFF the host except the
# deny-private proxy. The proxy is the ONLY container with an internet route
# (on its own `sandbox-egress-out` network) and its ACL denies private
# destinations. We flip a network's scalar property by KEY (a well-defined
# Compose merge); we do NOT try to remove an item from the service's `networks`
# LIST, which Compose cannot express (it unions lists) -- the reason a naive
# `networks: [sandbox-net]` override silently leaves the runner on its old net.
#
# Two viable approaches; pick one:
# WHAT THIS OVERLAY DOES NOT CONTAIN (read before enabling the sandbox):
# The runner STILL shares `sandbox-net` with `backend` and `worker` -- that is
# its control path and it cannot be removed without breaking code execution. A
# shared Docker network is bidirectional and Compose cannot sever it
# one-directionally, so arbitrary sandbox code can still open sockets to
# `backend:7091` and reach the worker. The internal flip contains raw sockets to
# the internet / host / RFC1918 / metadata, NOT to backend/worker on sandbox-net.
# (The Kubernetes NetworkPolicy DOES block this, via its RFC1918 egress
# carve-out; compose has no equivalent.)
#
# (1) Host / cloud firewall (simplest, belt-and-suspenders for single-host):
# Also drop egress to RFC1918 (10/8, 172.16/12, 192.168/16), link-local
# (169.254/16, incl. the 169.254.169.254 metadata IP), and ULA on the
# docsgpt-sandbox container's interface using iptables/nftables on the
# Docker host. Example (host root):
# SBX=$(docker inspect -f '{{.NetworkSettings.Networks.docsgpt-oss_sandbox-egress.IPAddress}}' docsgpt-sandbox)
# iptables -I DOCKER-USER -s "$SBX" -d 169.254.0.0/16 -j DROP
# iptables -I DOCKER-USER -s "$SBX" -d 10.0.0.0/8 -j DROP
# iptables -I DOCKER-USER -s "$SBX" -d 172.16.0.0/12 -j DROP
# iptables -I DOCKER-USER -s "$SBX" -d 192.168.0.0/16 -j DROP
# (Allow established/return traffic and DNS as needed for your setup.)
# MITIGATION -- required when enabling the sandbox (do the first, ideally both):
# - Run the backend with real authentication (`AUTH_TYPE` != none / a real auth
# provider) so a reachable API rejects unauthenticated requests. Without it,
# runner->backend reach is a free control-plane bypass.
# - Add a host-firewall DROP for runner->backend/worker on sandbox-net
# (approach (1) below).
#
# Two layers below: (2) is what this overlay wires up (internet egress via the
# deny-private proxy); (1) is the host-firewall DROP that closes the
# runner->backend/worker gap. Apply both for the strongest posture.
#
# (1) Host-firewall DROP for runner->backend/worker (closes the gap above). Run
# as root on the Docker host AFTER `... up`. The project name is pinned to
# `docsgpt-oss` (see `name:` in docker-compose.yaml), so the shared control
# net is `docsgpt-oss_sandbox-net`. `docker compose ... ps -q` resolves each
# container id regardless of its generated name (there is no fixed
# `container_name`), and the `index` function is REQUIRED because a
# Go-template dotted key cannot contain hyphens -- the old
# `{{.NetworkSettings.Networks.docsgpt-oss_sandbox-egress.IPAddress}}` form
# never parsed. Both IPs are read on sandbox-net so source+dest match the
# shared path:
# CF="-f deployment/docker-compose.yaml -f deployment/optional/docker-compose.optional.sandbox-egress.yaml --profile sandbox"
# NET=docsgpt-oss_sandbox-net
# SBX=$(docker inspect -f "{{(index .NetworkSettings.Networks \"$NET\").IPAddress}}" "$(docker compose $CF ps -q docsgpt-sandbox)")
# BE=$( docker inspect -f "{{(index .NetworkSettings.Networks \"$NET\").IPAddress}}" "$(docker compose $CF ps -q backend)")
# WK=$( docker inspect -f "{{(index .NetworkSettings.Networks \"$NET\").IPAddress}}" "$(docker compose $CF ps -q worker)")
# iptables -I DOCKER-USER -s "$SBX" -d "$BE" -j DROP # runner -> backend (incl. :7091)
# iptables -I DOCKER-USER -s "$SBX" -d "$WK" -j DROP # runner -> worker
# (These container-to-container DROPs require the host's
# `bridge-nf-call-iptables=1`, Docker's default, so bridged traffic
# traverses DOCKER-USER.) If a lookup is empty (renamed project, or the
# container is not up yet), resolve by hand instead of running a broken
# command: `docker compose $CF ps` to list ids, then `docker inspect
# <runner-id>` / `<backend-id>`, read the sandbox-net
# `.NetworkSettings.Networks.*.IPAddress`, and add the same DROP rules.
# Re-run after any recreate -- container IPs are not stable across `up`.
# (Internet / RFC1918 / metadata egress is already gone via the internal
# flip, so no separate SSRF drop is needed here.)
#
# (2) Egress-gateway sidecar (this overlay): the runner has NO direct internet
# route (see above); its only path out is the deny-private forward proxy.
# Point the runner's HTTP(S) client at the proxy via the env vars below.
# This does NOT stop runner->backend/worker on sandbox-net -- pair it with
# backend auth and/or approach (1).
#
# Apply alongside the base stack:
# docker compose -f deployment/docker-compose.yaml \
@@ -43,11 +80,14 @@
# (also works layered on docker-compose-hub.yaml / docker-compose-azure.yaml,
# which share the same service/network names).
#
# NOTE: because `sandbox-egress` is internal here, non-HTTP raw egress is blocked
# at L3 regardless of what the code does (there is simply no route). HTTP(S) that
# honors the proxy env is filtered by the proxy ACL. For belt-and-suspenders on a
# multi-tenant or untrusted deployment, also apply approach (1) or the Kubernetes
# NetworkPolicy.
# NOTE: because `sandbox-egress` is internal here, non-HTTP raw egress TO THE
# INTERNET / HOST / RFC1918 / METADATA is blocked at L3 regardless of what the
# code does (there is simply no route). HTTP(S) that honors the proxy env is
# filtered by the proxy ACL. This does NOT block raw egress to `backend:7091` /
# the worker on the shared `sandbox-net` (see WHAT THIS OVERLAY DOES NOT CONTAIN
# / MITIGATION above): enable backend auth and/or apply approach (1). On a
# multi-tenant or untrusted deployment, prefer the Kubernetes NetworkPolicy,
# which blocks the internal path too.
services:
# Forward proxy that denies private/link-local/metadata destinations and
+4 -1
View File
@@ -58,7 +58,10 @@ COPY kernel-launch.sh /opt/docsgpt/kernel-launch.sh
RUN chmod 0555 /opt/docsgpt/kernel-launch.sh
COPY kernels/docsgpt-python/kernel.json /usr/local/share/jupyter/kernels/docsgpt-python/kernel.json
USER sandbox
# Numeric UID (not the name) so a kubelet with `runAsNonRoot: true` can verify
# the user is non-root without resolving /etc/passwd. This uid MUST match
# `runAsUser` in deployment/k8s/deployments/sandbox-deploy.yaml.
USER 10001
WORKDIR /home/sandbox
EXPOSE 8888
+14 -5
View File
@@ -228,11 +228,20 @@ The hardened container runs **without `NET_ADMIN`**, so it cannot self-apply
[`deployment/optional/docker-compose.optional.sandbox-egress.yaml`](../optional/docker-compose.optional.sandbox-egress.yaml),
which flips `sandbox-egress` to `internal: true` (removing the runner's direct
internet/RFC1918/metadata route entirely) and forces egress through a
deny-private **egress-gateway proxy** sidecar; for belt-and-suspenders on
untrusted/multi-tenant hosts also add a **host/cloud firewall rule** (drop the
four private ranges on the sandbox container's interface), since a forward
proxy only filters code that honors `HTTP(S)_PROXY` — the `internal` flip is
what contains raw sockets. Note the broker/DB published ports are bound to
deny-private **egress-gateway proxy** sidecar. That `internal` flip is what
contains **raw sockets to the internet/host/RFC1918/metadata** (a forward proxy
only filters code that honors `HTTP(S)_PROXY`).
**What compose canNOT contain:** the runner stays on `sandbox-net` with the
backend and worker — that is its control path, and a shared Docker network is
bidirectional, so Compose cannot sever it one-directionally. Sandbox code can
therefore still open sockets to `backend:7091` and the worker. This is a real
gap the Kubernetes NetworkPolicy closes (via its RFC1918 egress carve-out) but
compose cannot. **Mitigate it** when enabling the sandbox: run the backend with
real authentication (`AUTH_TYPE` != none / a real auth provider) so a reachable
API rejects unauthenticated requests — **required** — and/or add a host-firewall
`DROP` for runner→backend/worker on `sandbox-net` (see approach (1) in the
overlay file's header comment). Note the broker/DB published ports are bound to
`127.0.0.1` so the runner cannot reach them via the host gateway either.
## Other hardening (deployment-level)
@@ -723,25 +723,34 @@ function WorkflowBuilderInner() {
}, [isPublishing]);
useEffect(() => {
// A keystroke targeting a text field / editable popover (the Code / JSON
// schema textareas, the Output Variable input) must edit the field, never
// fire a canvas shortcut.
const isEditableTarget = (target: EventTarget | null): boolean => {
const el = target as HTMLElement | null;
return (
el?.tagName === 'INPUT' ||
el?.tagName === 'TEXTAREA' ||
el?.isContentEditable === true
);
};
// Shared guard for both canvas shortcuts (Delete/Backspace to remove the
// selected node, Escape to close the config panel): ignore the keystroke
// while typing in a field, while the Preview Sheet is open (its own inputs
// own the keys — a stray Delete there must not delete the node behind it),
// or once another handler has already consumed the event. Kept in one place
// so the two branches can't drift apart.
const shouldIgnoreShortcut = (e: KeyboardEvent): boolean =>
e.defaultPrevented || showPreview || isEditableTarget(e.target);
const handleKeyDown = (e: KeyboardEvent) => {
if (e.key === 'Delete' && selectedNode) {
if (shouldIgnoreShortcut(e)) return;
if ((e.key === 'Delete' || e.key === 'Backspace') && selectedNode) {
handleDeleteNode();
return;
}
if (e.key === 'Escape') {
// While the Preview Sheet is open, leave Escape to its own dismissal
// (popovers, the Sheet itself) so a reflexive Escape can't tear down
// the node-config panel and lose the run/attachments/typed prompt.
if (showPreview) return;
// Ignore Escape originating from text fields / editable popovers so it
// dismisses the field rather than the whole config panel.
const target = e.target as HTMLElement | null;
if (
target?.tagName === 'INPUT' ||
target?.tagName === 'TEXTAREA' ||
target?.isContentEditable
) {
return;
}
setShowNodeConfig(false);
setSelectedNode(null);
}
@@ -32,6 +32,7 @@ import {
addQuery,
fetchWorkflowPreviewAnswer,
handleWorkflowPreviewAbort,
previewSendBlockReason,
resendQuery,
resetWorkflowPreview,
selectActiveNodeId,
@@ -443,6 +444,9 @@ export default function WorkflowPreview({
const hasCompletedAttachment = completedAttachments.length > 0;
const [lastQueryReturnedErr, setLastQueryReturnedErr] = useState(false);
const [sendBlockedMessage, setSendBlockedMessage] = useState<string | null>(
null,
);
const [openDetailsIndex, setOpenDetailsIndex] = useState<number | null>(null);
const [openArtifactsIndex, setOpenArtifactsIndex] = useState<number | null>(
null,
@@ -493,6 +497,20 @@ export default function WorkflowPreview({
isRetry?: boolean;
index?: number;
}) => {
// An unsaved draft can't bridge uploaded documents into the run (no
// persisted workflow_id), so block the send rather than run with the docs
// silently dropped. The uploads are kept (not cleared) for a retry after
// the workflow is saved.
const blockReason = previewSendBlockReason(
workflowId,
hasCompletedAttachment,
);
if (blockReason) {
setSendBlockedMessage(blockReason);
return;
}
setSendBlockedMessage(null);
const trimmedQuestion = question.trim();
// Doc-driven nodes read ``input_documents`` rather than the query, so an
// attachment-only run is allowed to proceed with an empty question.
@@ -509,9 +527,15 @@ export default function WorkflowPreview({
handleFetchAnswer({ question: trimmedQuestion, index: undefined });
}
},
[dispatch, handleFetchAnswer, hasCompletedAttachment],
[dispatch, handleFetchAnswer, hasCompletedAttachment, workflowId],
);
// Clear the block message once it no longer applies (the workflow was saved,
// or the attachments were removed) so a stale warning doesn't linger.
useEffect(() => {
if (workflowId || !hasCompletedAttachment) setSendBlockedMessage(null);
}, [workflowId, hasCompletedAttachment]);
const handleQuestionSubmission = (
question?: string,
updated?: boolean,
@@ -710,6 +734,11 @@ export default function WorkflowPreview({
)}
</div>
<div className="bg-card absolute right-0 bottom-0 left-0 flex w-full flex-col gap-2 px-4 pt-2 pb-4">
{sendBlockedMessage && (
<p className="text-xs text-red-500" role="alert">
{sendBlockedMessage}
</p>
)}
<MessageInput
onSubmit={(text) => handleQuestionSubmission(text)}
loading={status === 'loading'}
@@ -47,9 +47,15 @@ export default function WorkflowRunArtifacts({
.then(async (res: Response) => {
if (cancelled) return;
if (!res.ok) {
// The run row is missing/unauthorized (e.g. an unsaved-draft preview):
// show an informational empty state rather than a hard error.
setArtifacts([]);
// A 403 is the expected unsaved-draft / unauthorized case (no
// persisted run row): show the informational empty state. Any other
// non-OK status (500, an expired 401, ...) is a real failure — surface
// an error + Retry instead of masquerading as "no artifacts".
if (res.status === 403) {
setArtifacts([]);
} else {
setError('Failed to load artifacts');
}
setLoading(false);
return;
}
@@ -14,6 +14,12 @@ export interface WorkflowVariable {
label: string;
templatePath: string;
section: string;
// True when the variable resolves to artifact reference(s) at run time
// (uploaded input_documents, code-node outputs) rather than plain LLM/state
// TEXT. Only artifact-bearing variables may be picked as node Documents;
// selecting a text output there makes the engine append the literal variable
// name and the node hard-fails. Heuristic keyed on the producing node type.
producesArtifact?: boolean;
}
const GLOBAL_CONTEXT_VARIABLES: WorkflowVariable[] = [
@@ -118,41 +124,54 @@ export function extractUpstreamVariables(
label: 'agent.input_documents',
templatePath: 'agent.input_documents',
section: 'Workflow input',
// Uploaded documents are artifact references.
producesArtifact: true,
},
...GLOBAL_CONTEXT_VARIABLES,
];
const seen = new Set(variables.map((variable) => variable.templatePath));
const upstreamIds = getUpstreamNodeIds(selectedNodeId, edges);
const pushNodeOutput = (
node: Node,
outputName: string,
sectionFallback: string,
producesArtifact: boolean,
) => {
const templatePath = toAgentTemplatePath(outputName);
if (seen.has(templatePath)) return;
seen.add(templatePath);
variables.push({
label: templatePath,
templatePath,
section: node.data.title || node.data.label || sectionFallback,
producesArtifact,
});
};
for (const node of nodes) {
if (!upstreamIds.has(node.id)) continue;
if (node.type === 'agent') {
const defaultOutputTemplatePath = toAgentTemplatePath(
if (node.type === 'agent' || node.type === 'code') {
// Agent and code nodes both expose `node_<id>_output` and an optional
// `output_variable`, but only code-node outputs resolve to artifact
// references in the engine — agent outputs are LLM TEXT. So both feed the
// prompt-variable popover, while only code outputs are offered as
// Documents (see toDocumentVariableOptions / producesArtifact).
const producesArtifact = node.type === 'code';
const sectionFallback = node.type === 'code' ? 'Code' : 'Agent';
pushNodeOutput(
node,
`node_${node.id}_output`,
sectionFallback,
producesArtifact,
);
if (!seen.has(defaultOutputTemplatePath)) {
seen.add(defaultOutputTemplatePath);
variables.push({
label: defaultOutputTemplatePath,
templatePath: defaultOutputTemplatePath,
section: node.data.title || node.data.label || 'Agent',
});
}
const outputVariable = String(
node.data?.config?.output_variable || '',
).trim();
if (outputVariable) {
const templatePath = toAgentTemplatePath(outputVariable);
if (!seen.has(templatePath)) {
seen.add(templatePath);
variables.push({
label: templatePath,
templatePath,
section: node.data.title || node.data.label || 'Agent',
});
}
pushNodeOutput(node, outputVariable, sectionFallback, producesArtifact);
}
}
@@ -88,27 +88,33 @@ describe('stripAgentPrefix', () => {
});
describe('toDocumentVariableOptions', () => {
const make = (templatePath: string): WorkflowVariable => ({
const make = (
templatePath: string,
producesArtifact = false,
): WorkflowVariable => ({
label: templatePath,
templatePath,
section: 'x',
producesArtifact,
});
it('keeps input_documents and upstream outputs as bare names', () => {
it('offers only artifact-producing variables (input_documents, code outputs)', () => {
const options = toDocumentVariableOptions([
make('agent.input_documents'),
make('agent.wire_doc'),
make("agent['weird name']"),
make('agent.input_documents', true),
make('agent.node_code1_output', true),
make("agent['weird name']", true),
]);
expect(options).toEqual([
{ value: 'input_documents', label: 'input_documents' },
{ value: 'wire_doc', label: 'wire_doc' },
{ value: 'node_code1_output', label: 'node_code1_output' },
{ value: 'weird name', label: 'weird name' },
]);
});
it('excludes query, chat_history and global context variables', () => {
it('excludes plain agent/state TEXT outputs that are not artifact refs', () => {
const options = toDocumentVariableOptions([
make('agent.node_agent1_output'),
make('agent.some_state_var'),
make('agent.query'),
make('agent.chat_history'),
make('source.content'),
@@ -120,8 +126,8 @@ describe('toDocumentVariableOptions', () => {
it('deduplicates by bare name', () => {
const options = toDocumentVariableOptions([
make('agent.wire_doc'),
make('agent.wire_doc'),
make('agent.wire_doc', true),
make('agent.wire_doc', true),
]);
expect(options).toEqual([{ value: 'wire_doc', label: 'wire_doc' }]);
});
+12 -10
View File
@@ -59,23 +59,25 @@ export function stripAgentPrefix(templatePath: string): string {
return trimmed;
}
/** Multiselect options of upstream document variables, stored as bare names. */
/**
* Multiselect options of upstream document variables, stored as bare names.
*
* Only variables that carry artifact references at run time are offered
* (`producesArtifact`): the uploaded `input_documents` and code-node outputs.
* Plain agent/state TEXT outputs are excluded — the engine can't resolve them
* to artifacts, so picking one appends the literal variable name and the node
* hard-fails. The distinction is assigned upstream in `extractUpstreamVariables`
* keyed on the producing node type.
*/
export function toDocumentVariableOptions(
variables: WorkflowVariable[],
): { value: string; label: string }[] {
const options: { value: string; label: string }[] = [];
const seen = new Set<string>();
for (const variable of variables) {
const path = variable.templatePath;
const isInputDocuments = path === 'agent.input_documents';
const isAgentOutput =
path.startsWith('agent.') &&
path !== 'agent.query' &&
path !== 'agent.chat_history';
const isAgentBracket = path.startsWith('agent[');
if (!isInputDocuments && !isAgentOutput && !isAgentBracket) continue;
if (!variable.producesArtifact) continue;
const bareName = stripAgentPrefix(path);
const bareName = stripAgentPrefix(variable.templatePath);
if (!bareName || seen.has(bareName)) continue;
seen.add(bareName);
options.push({ value: bareName, label: bareName });
@@ -4,9 +4,12 @@ import { Attachment } from '../../upload/uploadSlice';
import reducer, {
addQuery,
collectCompletedAttachmentIds,
collectRunAttachmentIds,
previewSendBlockReason,
resetWorkflowPreview,
setPreviewOpen,
setWorkflowRunId,
UNSAVED_DRAFT_ATTACHMENTS_MESSAGE,
} from './workflowPreviewSlice';
const seedState = () => reducer(undefined, { type: '@@INIT' });
@@ -39,6 +42,40 @@ describe('collectCompletedAttachmentIds', () => {
});
});
describe('collectRunAttachmentIds', () => {
const completed = [att({ id: 'done', status: 'completed' })];
it('returns completed ids for a saved workflow', () => {
expect(collectRunAttachmentIds(completed, 'wf-1')).toEqual(['done']);
});
it('returns [] for an unsaved draft so uploads are neither sent nor cleared', () => {
expect(collectRunAttachmentIds(completed, null)).toEqual([]);
expect(collectRunAttachmentIds(completed, undefined)).toEqual([]);
expect(collectRunAttachmentIds(completed, '')).toEqual([]);
});
});
describe('previewSendBlockReason', () => {
it('blocks an unsaved draft that has completed attachments', () => {
expect(previewSendBlockReason(null, true)).toBe(
UNSAVED_DRAFT_ATTACHMENTS_MESSAGE,
);
expect(previewSendBlockReason(undefined, true)).toBe(
UNSAVED_DRAFT_ATTACHMENTS_MESSAGE,
);
expect(previewSendBlockReason('', true)).toBe(
UNSAVED_DRAFT_ATTACHMENTS_MESSAGE,
);
});
it('allows a saved workflow, or an unsaved draft with no attachments', () => {
expect(previewSendBlockReason('wf-1', true)).toBeNull();
expect(previewSendBlockReason('wf-1', false)).toBeNull();
expect(previewSendBlockReason(null, false)).toBeNull();
});
});
describe('setWorkflowRunId', () => {
it('stores the run id on the addressed query', () => {
let state = seedState();
@@ -76,6 +76,40 @@ export function collectCompletedAttachmentIds(
.map((att) => att.id);
}
/** User-facing message shown when a Preview send is blocked. */
export const UNSAVED_DRAFT_ATTACHMENTS_MESSAGE =
'Save the workflow before attaching documents in Preview.';
/**
* Reason a Preview send must be blocked, or null when it may proceed. An
* unsaved draft has no persisted ``workflow_id``, so the backend can't bridge
* uploaded documents into the run — it would execute with the docs silently
* dropped. Block instead, keeping the uploads for a retry after the save.
*/
export function previewSendBlockReason(
workflowId: string | null | undefined,
hasCompletedAttachment: boolean,
): string | null {
if (!workflowId && hasCompletedAttachment) {
return UNSAVED_DRAFT_ATTACHMENTS_MESSAGE;
}
return null;
}
/**
* Completed attachment ids eligible for a run. Attachments only survive as
* run-scoped artifacts when the run persists a ``workflow_runs`` row, which
* needs a saved ``workflow_id``; on an unsaved draft none are eligible (the UI
* blocks that send upstream), so they are neither sent nor cleared.
*/
export function collectRunAttachmentIds(
attachments: Attachment[],
workflowId: string | null | undefined,
): string[] {
if (!workflowId) return [];
return collectCompletedAttachmentIds(attachments);
}
export const fetchWorkflowPreviewAnswer = createAsyncThunk<
void,
{
@@ -100,8 +134,12 @@ export const fetchWorkflowPreviewAnswer = createAsyncThunk<
if (state.preference) {
// Doc-driven workflows read uploaded files as run-scoped artifacts; pass
// the completed attachment ids so the backend bridges them (user-scoped).
const attachmentIds = collectCompletedAttachmentIds(
// Gated on a saved workflow_id: an unsaved draft can't bridge them, so we
// neither send nor clear the uploads (the Preview UI blocks that send and
// keeps them for a retry after saving).
const attachmentIds = collectRunAttachmentIds(
state.upload.attachments,
workflowId,
);
if (attachmentIds.length > 0) dispatch(clearAttachments());
+11 -4
View File
@@ -114,10 +114,17 @@ const endpoints = {
`/api/artifacts/${artifactId}`,
LIST_WORKFLOW_RUN_ARTIFACTS: (workflowRunId: string) =>
`/api/artifacts?workflow_run_id=${encodeURIComponent(workflowRunId)}`,
DOWNLOAD_ARTIFACT: (artifactId: string, version?: number) =>
version != null
? `/api/artifacts/${artifactId}/download?version=${version}`
: `/api/artifacts/${artifactId}/download`,
DOWNLOAD_ARTIFACT: (
artifactId: string,
version?: number,
disposition?: string,
) => {
const params = new URLSearchParams();
if (version != null) params.set('version', String(version));
if (disposition) params.set('disposition', disposition);
const qs = params.toString();
return `/api/artifacts/${artifactId}/download${qs ? `?${qs}` : ''}`;
},
RESTORE_ARTIFACT: (artifactId: string) =>
`/api/artifacts/${artifactId}/restore`,
WORKFLOWS: '/api/workflows',
+9 -1
View File
@@ -339,8 +339,16 @@ const userService = {
artifactId: string,
token: string | null,
version?: number,
disposition?: 'url',
): Promise<Response> =>
apiClient.get(endpoints.USER.DOWNLOAD_ARTIFACT(artifactId, version), token),
apiClient.get(
endpoints.USER.DOWNLOAD_ARTIFACT(artifactId, version, disposition),
token,
// ?disposition=url asks the s3 strategy to return the presigned URL as
// JSON (for a top-level navigation) instead of a CORS-blocked 302; the
// Accept header is the same opt-in via content negotiation.
disposition === 'url' ? { Accept: 'application/json' } : {},
),
restoreArtifactVersion: (
artifactId: string,
version: number,
@@ -279,6 +279,7 @@ export default function DocumentArtifactView({
artifact.id,
token,
selectedVersion,
'url',
);
const ok = await triggerResponseDownload(response, filename);
if (!ok && isMountedRef.current) setActionError('Download failed');
@@ -1,4 +1,4 @@
import { describe, expect, it } from 'vitest';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import {
buildPreviewDocument,
@@ -12,7 +12,9 @@ import {
PREVIEW_CSP,
PREVIEW_CSP_MERMAID,
previewModeForKind,
readPresignedUrlEnvelope,
sortVersionsDesc,
triggerResponseDownload,
type ArtifactVersion,
} from './artifactViewUtils';
@@ -234,6 +236,104 @@ describe('buildPreviewDocument', () => {
});
});
describe('readPresignedUrlEnvelope', () => {
const json = (body: unknown, contentType = 'application/json') =>
new Response(JSON.stringify(body), {
headers: { 'Content-Type': contentType },
});
it('extracts the presigned url from the s3 JSON envelope', async () => {
const url = await readPresignedUrlEnvelope(
json({ success: true, url: 'https://signed.example/x?sig=1' }),
);
expect(url).toBe('https://signed.example/x?sig=1');
});
it('returns null for a byte stream and leaves the body readable', async () => {
const res = new Response('<p>hi</p>', {
headers: { 'Content-Type': 'text/html' },
});
expect(await readPresignedUrlEnvelope(res)).toBeNull();
expect(await res.text()).toBe('<p>hi</p>');
});
it('returns null for a JSON artifact whose bytes are not the envelope', async () => {
const res = json({ foo: 1 });
expect(await readPresignedUrlEnvelope(res)).toBeNull();
// clone() was used to peek, so the original body is still consumable.
expect(await res.json()).toEqual({ foo: 1 });
});
it('rejects a non-http(s) url in the envelope', async () => {
expect(
await readPresignedUrlEnvelope(
json({ success: true, url: 'javascript:alert(1)' }),
),
).toBeNull();
expect(
await readPresignedUrlEnvelope(json({ success: true, url: '' })),
).toBeNull();
});
});
describe('triggerResponseDownload', () => {
let anchors: HTMLAnchorElement[];
beforeEach(() => {
anchors = [];
vi.spyOn(URL, 'createObjectURL').mockReturnValue('blob:mock');
vi.spyOn(URL, 'revokeObjectURL').mockImplementation(() => undefined);
const orig = document.createElement.bind(document);
vi.spyOn(document, 'createElement').mockImplementation((tag: string) => {
const el = orig(tag);
if (String(tag).toLowerCase() === 'a') {
vi.spyOn(el as HTMLAnchorElement, 'click').mockImplementation(
() => undefined,
);
anchors.push(el as HTMLAnchorElement);
}
return el;
});
});
afterEach(() => {
vi.restoreAllMocks();
});
it('returns false for a non-OK response', async () => {
const res = new Response('', { status: 500 });
expect(await triggerResponseDownload(res, 'f.bin')).toBe(false);
});
it('navigates top-level to the presigned url for the s3 envelope', async () => {
const res = new Response(
JSON.stringify({ success: true, url: 'https://signed.example/x' }),
{ headers: { 'Content-Type': 'application/json' } },
);
expect(await triggerResponseDownload(res, 'f.bin')).toBe(true);
// No blob path for the presigned envelope.
expect(URL.createObjectURL).not.toHaveBeenCalled();
const anchor = anchors.at(-1)!;
expect(anchor.href).toBe('https://signed.example/x');
expect(anchor.target).toBe('_blank');
expect(anchor.click).toHaveBeenCalled();
});
it('saves streamed bytes via a blob url for the backend strategy', async () => {
const res = new Response('BINARY', {
headers: {
'Content-Type': 'application/octet-stream',
'Content-Disposition': 'attachment; filename="deck.pptx"',
},
});
expect(await triggerResponseDownload(res, 'fallback.bin')).toBe(true);
expect(URL.createObjectURL).toHaveBeenCalled();
const anchor = anchors.at(-1)!;
expect(anchor.download).toBe('deck.pptx');
expect(URL.revokeObjectURL).toHaveBeenCalledWith('blob:mock');
});
});
describe('filenameFromContentDisposition', () => {
it('reads the plain quoted filename', () => {
expect(
+59 -2
View File
@@ -258,14 +258,71 @@ export function filenameFromContentDisposition(
}
/**
* Turn a fetched download Response into a browser file-save, honoring
* `Content-Disposition` when present. Returns false if the response is not OK.
* Detect the `URL_STRATEGY=s3` presigned-URL envelope and return its `url`, or
* null when the response is a byte stream (`URL_STRATEGY=backend`). The download
* endpoint returns `{ success: true, url }` as JSON only when asked via
* `?disposition=url`; a plain byte stream (even a JSON artifact's own bytes)
* carries the file's own content-type and never this exact shape. The response
* is `clone()`d before reading so the caller can still consume the original body
* as bytes when this is not the envelope. The `url` must be an absolute http(s)
* URL to further guard against a JSON artifact coincidentally matching.
*/
export async function readPresignedUrlEnvelope(
response: Response,
): Promise<string | null> {
const contentType = response.headers.get('Content-Type') ?? '';
if (!contentType.toLowerCase().includes('application/json')) return null;
const data = await response
.clone()
.json()
.catch(() => null);
if (
data &&
typeof data === 'object' &&
(data as { success?: unknown }).success === true &&
typeof (data as { url?: unknown }).url === 'string' &&
/^https?:\/\//i.test((data as { url: string }).url)
) {
return (data as { url: string }).url;
}
return null;
}
/** Trigger a top-level browser navigation to `url` (opens a new tab). */
function triggerUrlNavigation(url: string, fallbackName: string): void {
const link = document.createElement('a');
link.href = url;
// `download` is ignored for cross-origin URLs, so the object's own
// Content-Disposition drives the save; open a new tab so the app isn't
// navigated away if the object renders inline.
link.download = fallbackName;
link.target = '_blank';
link.rel = 'noopener noreferrer';
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
}
/**
* Turn a fetched download Response into a browser file-save. Supports both URL
* strategies from a single call: the s3 strategy returns a presigned-URL JSON
* envelope (fetching its bytes cross-origin would be CORS-blocked, so navigate
* to it top-level), while the backend strategy streams the bytes (saved via a
* blob URL, honoring `Content-Disposition`). Returns false if the response is
* not OK, or is a malformed JSON envelope with no usable url.
*/
export async function triggerResponseDownload(
response: Response,
fallbackName: string,
): Promise<boolean> {
if (!response.ok) return false;
const presignedUrl = await readPresignedUrlEnvelope(response);
if (presignedUrl) {
triggerUrlNavigation(presignedUrl, fallbackName);
return true;
}
const name =
filenameFromContentDisposition(
response.headers.get('Content-Disposition'),
+14 -3
View File
@@ -4,6 +4,7 @@ import userService from '../api/services/userService';
import {
type BytesPreviewMode,
MAX_INLINE_TEXT_BYTES,
readPresignedUrlEnvelope,
} from './artifactViewUtils';
type BytesState =
@@ -49,12 +50,22 @@ export function useArtifactBytes(
};
userService
.downloadArtifact(artifactId, token, version)
.downloadArtifact(artifactId, token, version, 'url')
.then(async (response: Response) => {
if (!response.ok) throw new Error('download failed');
// Under URL_STRATEGY=s3 the endpoint returns a presigned URL; fetch the
// bytes directly from it (a bucket without a CORS grant for the app
// origin blocks this read → caught below → download-card fallback). The
// backend strategy streams the bytes here, so this is a no-op for it.
const presignedUrl = await readPresignedUrlEnvelope(response);
const bytesResponse = presignedUrl
? await fetch(presignedUrl)
: response;
if (!bytesResponse.ok) throw new Error('bytes fetch failed');
if (mode === 'image') {
const blob = await response.blob();
const blob = await bytesResponse.blob();
if (cancelled) return;
revoke();
const url = URL.createObjectURL(blob);
@@ -64,7 +75,7 @@ export function useArtifactBytes(
}
// html / svg / text families render as a (bounded) string.
const blob = await response.blob();
const blob = await bytesResponse.blob();
if (cancelled) return;
if (blob.size > MAX_INLINE_TEXT_BYTES) throw new Error('too large');
const text = await blob.text();
+48
View File
@@ -7,8 +7,11 @@ against the ephemeral ``pg_conn`` fixture. Agent construction is bypassed via
from contextlib import contextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
def _make_agent(*, workflow_id=None, workflow=None, workflow_owner=None,
@@ -334,3 +337,48 @@ class TestGen:
assert results == [
{"type": "error", "error": "Failed to load workflow configuration."}
]
class TestAgentNodeApprovalPause:
"""A node agent whose tool pauses for approval must fail the node visibly, not emit empty output."""
def test_tool_calls_pending_raises_clear_error(self, monkeypatch):
from application.agents.workflows import workflow_engine as we
from application.agents.workflows.schemas import (
NodeType,
Workflow,
WorkflowGraph,
WorkflowNode,
)
from application.agents.workflows.workflow_engine import WorkflowEngine
# An ephemeral node agent whose LLM handler yields the pause signal and ends,
# emitting no "answer". Previously the engine dropped it and the node completed
# with empty output (or raised a confusing "Structured output was expected").
class _PendingAgent:
attachments = None
def gen(self, prompt):
yield {"type": "tool_calls_pending", "data": {"pending_tool_calls": [{"id": "x"}]}}
monkeypatch.setattr(
we.WorkflowNodeAgentFactory, "create", lambda **kw: _PendingAgent()
)
graph = WorkflowGraph(workflow=Workflow(name="Approval Pause"), nodes=[], edges=[])
agent = SimpleNamespace(
endpoint="stream", llm_name="openai", model_id="gpt-4o-mini", api_key="k",
chat_history=[], user="u", decoded_token={"sub": "u"},
)
engine = WorkflowEngine(
graph, agent, workflow_run_id="22222222-2222-2222-2222-222222222222"
)
node = WorkflowNode(
id="agent_1", workflow_id="wf-1", type=NodeType.AGENT, title="Extractor",
position={"x": 0, "y": 0},
config={"llm_name": "openai", "system_prompt": "s", "tools": ["t"]},
)
with pytest.raises(ValueError, match="requires approval"):
list(engine._execute_agent_node(node))
+51
View File
@@ -396,6 +396,57 @@ def test_materialize_code_inputs_rejects_oversize(monkeypatch):
assert manager.put_files == [] # nothing staged
def test_materialize_code_inputs_dedupes_same_filename(monkeypatch):
"""Two code-node inputs sharing a filename stage to DISTINCT inputs/ paths (no clobber)."""
from contextlib import contextmanager
monkeypatch.setattr(
"application.agents.tools.artifact_ref.resolve_artifact_id",
lambda repo, raw, **k: str(raw),
)
class _Repo:
def __init__(self, conn):
pass
def get_artifact_in_parent(self, artifact_id, *, workflow_run_id=None, conversation_id=None):
return {"id": artifact_id, "current_version": 1}
def get_version(self, artifact_id, version):
# Same filename, distinct stored bytes per input.
return {"filename": "data.csv", "storage_path": f"p/{artifact_id}.csv"}
@contextmanager
def _readonly():
yield object()
class _Storage:
def get_file(self, path):
import io
return io.BytesIO(path.encode())
monkeypatch.setattr(
"application.storage.db.repositories.artifacts.ArtifactsRepository", _Repo
)
monkeypatch.setattr("application.storage.db.session.db_readonly", _readonly)
monkeypatch.setattr(
"application.storage.storage_creator.StorageCreator.get_storage",
staticmethod(lambda: _Storage()),
)
engine = _engine()
manager = _FakeManager(_Result(ok=True, stdout="ok"))
loaded = engine._materialize_code_inputs(
manager, engine._session_id(), ["art-a", "art-b"], "user-code"
)
assert loaded == ["inputs/data.csv", "inputs/data-2.csv"]
staged = dict(manager.put_files)
assert set(staged) == {"inputs/data.csv", "inputs/data-2.csv"}
assert staged["inputs/data.csv"] != staged["inputs/data-2.csv"] # no clobber
# ---------------------------------------------------------------------------
# Pass-by-reference: survives serialization + CEL branches on the metadata.
# ---------------------------------------------------------------------------
+59
View File
@@ -102,3 +102,62 @@ def test_resolve_ref_under_workflow_run_parent():
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")]
# ---------------------------------------------------------------------------
# Stable per-parent ref_seq: a ref no longer re-points after an earlier delete
# ---------------------------------------------------------------------------
class _SeqRepo:
"""Fake repo backing resolve_id_by_ref_seq (stable) with the legacy positional fallback."""
def __init__(self, *, seq_map=None, pos_map=None, conv=None, run=None):
self.seq_map = seq_map or {}
self.pos_map = pos_map or {}
self.conv = conv
self.run = run
self.seq_calls = []
self.pos_calls = []
def _in_scope(self, conversation_id, workflow_run_id) -> bool:
if self.conv is not None and conversation_id != self.conv:
return False
if self.run is not None and workflow_run_id != self.run:
return False
return True
def resolve_id_by_ref_seq(self, seq, *, conversation_id=None, workflow_run_id=None):
self.seq_calls.append((seq, conversation_id, workflow_run_id))
if not self._in_scope(conversation_id, workflow_run_id):
return None
return self.seq_map.get(seq)
def artifact_id_at_position(self, n, *, conversation_id=None, workflow_run_id=None):
self.pos_calls.append((n, conversation_id, workflow_run_id))
if not self._in_scope(conversation_id, workflow_run_id):
return None
return self.pos_map.get(n)
def test_ref_resolves_by_stable_seq_not_shifted_position():
b, c = str(uuid.uuid4()), str(uuid.uuid4())
# After deleting the earlier A(seq 1): ref_seq keeps B=2/C=3; positions shifted to B=1/C=2.
repo = _SeqRepo(seq_map={2: b, 3: c}, pos_map={1: b, 2: c}, conv="conv-1")
assert resolve_artifact_id(repo, "A2", conversation_id="conv-1") == b
assert resolve_artifact_id(repo, "A3", conversation_id="conv-1") == c
# The stable ref_seq hit; the (shifted) positional fallback was never consulted.
assert repo.pos_calls == []
def test_ref_falls_back_to_position_for_legacy_rows():
legacy = str(uuid.uuid4())
# Legacy row has no ref_seq -> resolve_id_by_ref_seq misses -> positional fallback resolves it.
repo = _SeqRepo(seq_map={}, pos_map={1: legacy}, conv="conv-1")
assert resolve_artifact_id(repo, "A1", conversation_id="conv-1") == legacy
assert repo.seq_calls == [(1, "conv-1", None)]
assert repo.pos_calls == [(1, "conv-1", None)]
def test_ref_seq_does_not_cross_parents():
b = str(uuid.uuid4())
repo = _SeqRepo(seq_map={2: b}, pos_map={2: b}, conv="conv-1")
assert resolve_artifact_id(repo, "A2", conversation_id="conv-OTHER") is None
@@ -302,3 +302,55 @@ def test_malformed_json_schema_rejected_before_enqueue(monkeypatch):
)
out = _tool().execute_action("read_document", input=_ART_ID, json_schema={"properties": {}}, persist=False)
assert out["status"] == "error" and "invalid json_schema" in out["error"]
# ---------------------------------------------------------------------------
# Worker-context: parse INLINE inside a Celery worker; dispatch from the web
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_dispatch_inline_when_in_worker(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
# Inside a worker current_task is truthy -> parse inline, never enqueue (else the
# parsing queue self-deadlocks the worker that also serves it).
monkeypatch.setattr(rd, "current_task", object())
import application.api.user.tasks as tasks
monkeypatch.setattr(
tasks.parse_document, "apply_async",
lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not enqueue inside a worker")),
)
import application.worker as worker
called: Dict[str, Any] = {}
def _fake_run(artifact_id, parent, user_id, options):
called["args"] = (artifact_id, parent, user_id)
called["options"] = options
return {"status": "ok", "content": "inline", "truncated": False}
monkeypatch.setattr(worker, "run_parse_document", _fake_run)
out = _tool().execute_action("read_document", input=_ART_ID, persist=False)
assert out["status"] == "ok" and out["content"] == "inline"
# Same auth re-resolution + parent shape as the dispatch path.
assert called["args"] == (_ART_ID, {"conversation_id": "conv-1"}, "u-1")
assert called["options"]["persist"] is False
@pytest.mark.unit
def test_dispatch_enqueues_when_not_in_worker(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
# Web process: current_task falsy -> dispatch to the parsing queue, never inline.
monkeypatch.setattr(rd, "current_task", None)
captured = _patch_task(monkeypatch, payload={"status": "ok", "content": "queued", "truncated": False})
import application.worker as worker
monkeypatch.setattr(
worker, "run_parse_document",
lambda *a, **k: (_ for _ in ()).throw(AssertionError("web path must dispatch, not inline")),
)
out = _tool().execute_action("read_document", input=_ART_ID, persist=False)
assert out["status"] == "ok" and out["content"] == "queued"
assert captured["args"][0] == _ART_ID
assert captured["args"][1] == {"conversation_id": "conv-1"}
+30
View File
@@ -580,6 +580,36 @@ class TestDownloadArtifact:
assert resp.status_code == 302
assert resp.headers["Location"].startswith("https://signed.example/")
def test_s3_strategy_disposition_url_returns_json(
self, _patch_db, flask_app, token_owner, monkeypatch
):
# ?disposition=url opts into a JSON envelope (for a top-level browser
# navigation) instead of the CORS-blocked cross-origin 302.
from application.api.user.artifacts.routes import DownloadArtifact
art = self._seed(_patch_db)
storage = _FakeStorage(b"unused")
monkeypatch.setattr(
"application.api.user.artifacts.routes.StorageCreator.get_storage",
lambda: storage,
)
monkeypatch.setattr(
"application.api.user.artifacts.routes.settings.URL_STRATEGY",
"s3",
raising=False,
)
resp = _call(
flask_app,
DownloadArtifact,
art["id"],
token=token_owner,
query={"disposition": "url"},
)
assert resp.status_code == 200
assert resp.json["success"] is True
assert resp.json["url"].startswith("https://signed.example/")
def test_stranger_denied(self, _patch_db, flask_app, monkeypatch):
from application.api.user.artifacts.routes import DownloadArtifact
+36 -9
View File
@@ -61,19 +61,21 @@ def test_markdown_output_default():
@pytest.mark.unit
def test_max_chars_truncates_and_flags():
def test_max_chars_does_not_truncate_parse_output():
# max_chars now bounds only the VIEW (bound_parse_payload); parse returns the FULL text
# so the persisted artifact is the complete parse.
out = parse_document_bytes(("A" * 100).encode(), "note.txt", output="text", max_chars=10, include_tables=False)
assert out["truncated"] is True
assert len(out["content"]) == 10
assert out["content"] == "A" * 100
assert out["truncated"] is False
@pytest.mark.unit
def test_default_window_truncates_large_text():
big = ("A" * (dr._TEXT_MAX_BYTES * 3)).encode()
out = parse_document_bytes(big, "note.txt", output="text", include_tables=False)
assert out["truncated"] is True
assert "...[truncated" in out["content"]
assert len(out["content"].encode("utf-8")) <= dr._TEXT_MAX_BYTES + 64
def test_large_text_is_full_in_parse_output():
# The default head+tail window moved to bound_parse_payload; parse keeps the full text.
big = "A" * (dr._TEXT_MAX_BYTES * 3)
out = parse_document_bytes(big.encode(), "note.txt", output="text", include_tables=False)
assert out["content"] == big
assert out["truncated"] is False
# ---------------------------------------------------------------------------
@@ -348,3 +350,28 @@ def test_bound_parse_payload_keeps_structured_for_validation():
out = bound_parse_payload({"output": "structured", "content": "# ok", "structured": structured})
# structured must survive so the tool's json_schema validation can run on it.
assert out["structured"] == structured
@pytest.mark.unit
def test_bound_parse_payload_max_chars_bounds_view_only():
# max_chars caps the returned view; the parse itself is unaffected (see parse tests above).
content = "A" * 100
out = bound_parse_payload({"output": "text", "content": content}, max_chars=10)
assert out["content"] == "A" * 10
assert out["truncated"] is True
@pytest.mark.unit
def test_bound_parse_payload_default_window_when_no_max_chars():
big = "A" * (dr._TEXT_MAX_BYTES * 3)
out = bound_parse_payload({"output": "text", "content": big})
assert "...[truncated" in out["content"]
assert len(out["content"].encode("utf-8")) <= dr._TEXT_MAX_BYTES + 64
assert out["truncated"] is True
@pytest.mark.unit
def test_bound_parse_payload_small_content_not_flagged():
out = bound_parse_payload({"output": "text", "content": "hi"}, max_chars=10)
assert out["content"] == "hi"
assert out["truncated"] is False
@@ -41,7 +41,8 @@ class TestCreateArtifact:
assert artifact["user_id"] == "user-1"
assert artifact["kind"] == "presentation"
assert artifact["title"] == "Q3 deck"
assert artifact["metadata"] == {"source": "chat"}
# A stable per-parent ref_seq is merged into the caller's metadata at creation.
assert artifact["metadata"] == {"source": "chat", "ref_seq": 1}
assert artifact["current_version"] == 1
assert artifact["id"] is not None
assert artifact["_id"] == artifact["id"]
@@ -460,6 +461,97 @@ class TestVirtualRefPositions:
repo.artifact_id_at_position(1)
class TestStableRefSeq:
"""A stable per-parent ``ref_seq`` (stored in metadata) backs the ``A{n}`` refs and survives deletes."""
def test_create_assigns_incrementing_ref_seq(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)
assert a["metadata"]["ref_seq"] == 1
assert b["metadata"]["ref_seq"] == 2
assert c["metadata"]["ref_seq"] == 3
def test_ref_seq_is_parent_scoped(self, pg_conn):
repo = _repo(pg_conn)
a1 = repo.create_artifact("u", "document", conversation_id=_conversation_id())
b1 = repo.create_artifact("u", "document", workflow_run_id=_conversation_id())
# Each parent numbers ref_seq from 1 independently.
assert a1["metadata"]["ref_seq"] == 1
assert b1["metadata"]["ref_seq"] == 1
def test_ref_seq_merges_without_clobbering_caller_metadata(self, pg_conn):
repo = _repo(pg_conn)
a = repo.create_artifact(
"u", "document", conversation_id=_conversation_id(), metadata={"k": "v"}
)
assert a["metadata"] == {"k": "v", "ref_seq": 1}
def test_next_ref_seq_is_monotonic_across_delete(self, pg_conn):
repo = _repo(pg_conn)
conv = _conversation_id()
a = repo.create_artifact("u", "document", conversation_id=conv)
repo.create_artifact("u", "document", conversation_id=conv) # ref_seq 2
repo.create_artifact("u", "document", conversation_id=conv) # ref_seq 3
assert repo.next_ref_seq(conversation_id=conv) == 4
# Deleting an earlier artifact must NOT lower the next seq (no reuse -> no re-point).
repo.delete_artifact(a["id"])
assert repo.next_ref_seq(conversation_id=conv) == 4
def test_resolve_id_by_ref_seq_is_stable_after_delete(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)
assert repo.resolve_id_by_ref_seq(2, conversation_id=conv) == b["id"]
assert repo.resolve_id_by_ref_seq(3, conversation_id=conv) == c["id"]
# Delete the FIRST artifact; B/C refs must NOT shift.
repo.delete_artifact(a["id"])
assert repo.resolve_id_by_ref_seq(2, conversation_id=conv) == b["id"]
assert repo.resolve_id_by_ref_seq(3, conversation_id=conv) == c["id"]
assert repo.resolve_id_by_ref_seq(1, conversation_id=conv) is None # A is gone
def test_resolve_id_by_ref_seq_missing_and_invalid(self, pg_conn):
repo = _repo(pg_conn)
conv = _conversation_id()
repo.create_artifact("u", "document", conversation_id=conv)
assert repo.resolve_id_by_ref_seq(99, conversation_id=conv) is None
assert repo.resolve_id_by_ref_seq(0, conversation_id=conv) is None
with pytest.raises(ValueError):
repo.next_ref_seq()
with pytest.raises(ValueError):
repo.resolve_id_by_ref_seq(1)
def test_resolver_resolves_a_n_by_stable_seq_after_delete(self, pg_conn):
from application.agents.tools.artifact_ref import resolve_artifact_id
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)
repo.delete_artifact(a["id"])
# The tool path resolves A2 to B via ref_seq even though B is now positionally first.
assert resolve_artifact_id(repo, "A2", conversation_id=conv) == b["id"]
def test_legacy_row_without_ref_seq_falls_back_to_position(self, pg_conn):
from application.agents.tools.artifact_ref import resolve_artifact_id
repo = _repo(pg_conn)
conv = _conversation_id()
art = repo.create_artifact("u", "document", conversation_id=conv)
# Simulate a pre-migration row that never got a ref_seq.
pg_conn.execute(
text("UPDATE artifacts SET metadata = NULL WHERE id = CAST(:id AS uuid)"),
{"id": art["id"]},
)
assert repo.resolve_id_by_ref_seq(1, conversation_id=conv) is None
# No ref_seq -> positional fallback still resolves A1.
assert resolve_artifact_id(repo, "A1", conversation_id=conv) == art["id"]
class TestCascadeDelete:
def test_deleting_artifact_removes_versions(self, pg_conn):
repo = _repo(pg_conn)
+51
View File
@@ -428,3 +428,54 @@ def test_materialize_inputs_rejects_oversize_by_declared_size(monkeypatch):
out = _tool()._materialize_inputs(manager, "conv-1", [_ART_ID])
assert "exceeds" in out["error"] and "sandbox input limit" in out["error"]
assert manager.put_files == {} # nothing staged
def test_materialize_inputs_dedupes_same_filename(monkeypatch):
"""Two inputs whose current versions share a filename stage to DISTINCT inputs/ paths."""
from application.agents.tools import code_executor as ce
id_a = str(uuid.uuid4())
id_b = str(uuid.uuid4())
class _Repo:
def __init__(self, conn):
pass
def artifact_id_at_position(self, n, *, conversation_id=None, workflow_run_id=None):
return None
def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None):
if conversation_id != "conv-1":
return None
return {"id": artifact_id, "current_version": 1, "title": "seed.csv"}
def get_version(self, artifact_id, version):
# Both inputs carry the SAME current filename but distinct stored bytes.
return {"filename": "seed.csv", "storage_path": f"p/{artifact_id}.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(path.encode()) # distinct bytes per artifact path
monkeypatch.setattr(ce, "db_readonly", lambda: _Conn())
monkeypatch.setattr(ce, "ArtifactsRepository", _Repo)
monkeypatch.setattr(ce.StorageCreator, "get_storage", staticmethod(lambda: _Storage()))
manager = _InputManager()
out = _tool()._materialize_inputs(manager, "conv-1", [id_a, id_b])
assert "error" not in out
# The colliding second input is suffixed before the extension; both are reported.
assert out["loaded"] == ["inputs/seed.csv", "inputs/seed-2.csv"]
assert set(manager.put_files) == {"inputs/seed.csv", "inputs/seed-2.csv"}
# Each path holds its own artifact's bytes (no clobber).
assert manager.put_files["inputs/seed.csv"] != manager.put_files["inputs/seed-2.csv"]
@@ -206,3 +206,61 @@ def test_result_payload_chunks_are_bounded(monkeypatch):
assert out["total_chunks"] == dr._MAX_CHUNKS_RETURNED * 3
assert all("...[truncated" in c for c in out["chunks"])
assert all(len(c) < len(huge_chunk) for c in out["chunks"])
def _capture_persist(monkeypatch):
"""Patch persist_new_artifact to record what would be stored; return the capture dict."""
import application.sandbox.artifacts_capture as ac
captured: Dict[str, Any] = {}
def _fake_persist(**kwargs):
captured.update(kwargs)
return {"artifact_id": "new-art", "version": 1, "filename": "x.json",
"mime_type": "application/json", "size": len(kwargs.get("data", b""))}
monkeypatch.setattr(ac, "persist_new_artifact", _fake_persist)
return captured
@pytest.mark.unit
def test_persist_keeps_full_content_while_view_is_bounded(monkeypatch):
# A >8000-byte doc: the persisted artifact must keep the FULL text, while the returned
# (Redis/LLM) view is head+tail windowed.
import json
_patch_repo(monkeypatch, found=True, run="run-1")
big = "Z" * 12000
_patch_parse(monkeypatch, {"output": "markdown", "content": big, "truncated": False})
captured = _capture_persist(monkeypatch)
out = worker.parse_document_worker(
None, _ART_ID, {"workflow_run_id": "run-1"}, "u-1", {"output": "markdown", "persist": True}
)
persisted = json.loads(captured["data"].decode("utf-8"))
assert persisted["content"] == big # FULL parse persisted
assert len(persisted["content"]) == 12000
assert len(out["content"]) < len(big) # bounded view
assert "...[truncated" in out["content"]
assert out["truncated"] is True
@pytest.mark.unit
def test_persist_full_but_view_respects_max_chars(monkeypatch):
# max_chars bounds only the returned view; the persisted artifact stays full.
import json
_patch_repo(monkeypatch, found=True, run="run-1")
big = "Z" * 5000
_patch_parse(monkeypatch, {"output": "markdown", "content": big, "truncated": False})
captured = _capture_persist(monkeypatch)
out = worker.parse_document_worker(
None, _ART_ID, {"workflow_run_id": "run-1"}, "u-1",
{"output": "markdown", "persist": True, "max_chars": 100},
)
assert json.loads(captured["data"].decode("utf-8"))["content"] == big # full persisted
assert len(out["content"]) == 100 # view capped by max_chars
assert out["truncated"] is True