diff --git a/application/agents/tools/artifact_ref.py b/application/agents/tools/artifact_ref.py index 830c622d..24751624 100644 --- a/application/agents/tools/artifact_ref.py +++ b/application/agents/tools/artifact_ref.py @@ -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, diff --git a/application/agents/tools/code_executor.py b/application/agents/tools/code_executor.py index 62e4fe60..e71c18d7 100644 --- a/application/agents/tools/code_executor.py +++ b/application/agents/tools/code_executor.py @@ -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: diff --git a/application/agents/tools/read_document.py b/application/agents/tools/read_document.py index 6f031fa0..7acdf60b 100644 --- a/application/agents/tools/read_document.py +++ b/application/agents/tools/read_document.py @@ -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."} diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index 713a2fa8..adc20b83 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -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( diff --git a/application/api/user/artifacts/routes.py b/application/api/user/artifacts/routes.py index 03e8c318..96911794 100644 --- a/application/api/user/artifacts/routes.py +++ b/application/api/user/artifacts/routes.py @@ -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 diff --git a/application/parser/document_reader.py b/application/parser/document_reader.py index 1d18ac5b..4524394d 100644 --- a/application/parser/document_reader.py +++ b/application/parser/document_reader.py @@ -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 diff --git a/application/sandbox/artifacts_capture.py b/application/sandbox/artifacts_capture.py index 581202c8..07270a84 100644 --- a/application/sandbox/artifacts_capture.py +++ b/application/sandbox/artifacts_capture.py @@ -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 diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py index c3a515be..ebe26aa6 100644 --- a/application/storage/db/repositories/artifacts.py +++ b/application/storage/db/repositories/artifacts.py @@ -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 = "" diff --git a/application/worker.py b/application/worker.py index 6cca0728..d7a05a67 100755 --- a/application/worker.py +++ b/application/worker.py @@ -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 diff --git a/deployment/k8s/deployments/sandbox-deploy.yaml b/deployment/k8s/deployments/sandbox-deploy.yaml index 941c99ef..24f92e81 100644 --- a/deployment/k8s/deployments/sandbox-deploy.yaml +++ b/deployment/k8s/deployments/sandbox-deploy.yaml @@ -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: diff --git a/deployment/k8s/network-policies/sandbox-egress-policy.yaml b/deployment/k8s/network-policies/sandbox-egress-policy.yaml index 489cb6eb..bf9ca8d2 100644 --- a/deployment/k8s/network-policies/sandbox-egress-policy.yaml +++ b/deployment/k8s/network-policies/sandbox-egress-policy.yaml @@ -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 diff --git a/deployment/optional/docker-compose.optional.sandbox-egress.yaml b/deployment/optional/docker-compose.optional.sandbox-egress.yaml index 4af85004..eca41f23 100644 --- a/deployment/optional/docker-compose.optional.sandbox-egress.yaml +++ b/deployment/optional/docker-compose.optional.sandbox-egress.yaml @@ -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 +# ` / ``, 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 diff --git a/deployment/sandbox/Dockerfile b/deployment/sandbox/Dockerfile index 39f25748..b1ac1816 100644 --- a/deployment/sandbox/Dockerfile +++ b/deployment/sandbox/Dockerfile @@ -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 diff --git a/deployment/sandbox/README.md b/deployment/sandbox/README.md index bb63d8e9..e68a80e8 100644 --- a/deployment/sandbox/README.md +++ b/deployment/sandbox/README.md @@ -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) diff --git a/frontend/src/agents/workflow/WorkflowBuilder.tsx b/frontend/src/agents/workflow/WorkflowBuilder.tsx index b3c81fce..753586da 100644 --- a/frontend/src/agents/workflow/WorkflowBuilder.tsx +++ b/frontend/src/agents/workflow/WorkflowBuilder.tsx @@ -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); } diff --git a/frontend/src/agents/workflow/WorkflowPreview.tsx b/frontend/src/agents/workflow/WorkflowPreview.tsx index 3b8151c5..fbc2e92b 100644 --- a/frontend/src/agents/workflow/WorkflowPreview.tsx +++ b/frontend/src/agents/workflow/WorkflowPreview.tsx @@ -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( + null, + ); const [openDetailsIndex, setOpenDetailsIndex] = useState(null); const [openArtifactsIndex, setOpenArtifactsIndex] = useState( 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({ )}
+ {sendBlockedMessage && ( +

+ {sendBlockedMessage} +

+ )} handleQuestionSubmission(text)} loading={status === 'loading'} diff --git a/frontend/src/agents/workflow/WorkflowRunArtifacts.tsx b/frontend/src/agents/workflow/WorkflowRunArtifacts.tsx index 7a17259b..53703468 100644 --- a/frontend/src/agents/workflow/WorkflowRunArtifacts.tsx +++ b/frontend/src/agents/workflow/WorkflowRunArtifacts.tsx @@ -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; } diff --git a/frontend/src/agents/workflow/components/PromptTextArea.tsx b/frontend/src/agents/workflow/components/PromptTextArea.tsx index c85a7221..ce133990 100644 --- a/frontend/src/agents/workflow/components/PromptTextArea.tsx +++ b/frontend/src/agents/workflow/components/PromptTextArea.tsx @@ -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__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); } } diff --git a/frontend/src/agents/workflow/documentConfig.test.ts b/frontend/src/agents/workflow/documentConfig.test.ts index a4ce2ce8..7066f873 100644 --- a/frontend/src/agents/workflow/documentConfig.test.ts +++ b/frontend/src/agents/workflow/documentConfig.test.ts @@ -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' }]); }); diff --git a/frontend/src/agents/workflow/documentConfig.ts b/frontend/src/agents/workflow/documentConfig.ts index c524c07e..191d76fe 100644 --- a/frontend/src/agents/workflow/documentConfig.ts +++ b/frontend/src/agents/workflow/documentConfig.ts @@ -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(); 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 }); diff --git a/frontend/src/agents/workflow/workflowPreviewSlice.test.ts b/frontend/src/agents/workflow/workflowPreviewSlice.test.ts index ebadefc3..c329652c 100644 --- a/frontend/src/agents/workflow/workflowPreviewSlice.test.ts +++ b/frontend/src/agents/workflow/workflowPreviewSlice.test.ts @@ -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(); diff --git a/frontend/src/agents/workflow/workflowPreviewSlice.ts b/frontend/src/agents/workflow/workflowPreviewSlice.ts index 061a4178..b7888885 100644 --- a/frontend/src/agents/workflow/workflowPreviewSlice.ts +++ b/frontend/src/agents/workflow/workflowPreviewSlice.ts @@ -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()); diff --git a/frontend/src/api/endpoints.ts b/frontend/src/api/endpoints.ts index 9b634cdd..8f23f896 100644 --- a/frontend/src/api/endpoints.ts +++ b/frontend/src/api/endpoints.ts @@ -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', diff --git a/frontend/src/api/services/userService.ts b/frontend/src/api/services/userService.ts index 5e283034..bd65491c 100644 --- a/frontend/src/api/services/userService.ts +++ b/frontend/src/api/services/userService.ts @@ -339,8 +339,16 @@ const userService = { artifactId: string, token: string | null, version?: number, + disposition?: 'url', ): Promise => - 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, diff --git a/frontend/src/components/DocumentArtifactView.tsx b/frontend/src/components/DocumentArtifactView.tsx index d2edc6fe..61db9692 100644 --- a/frontend/src/components/DocumentArtifactView.tsx +++ b/frontend/src/components/DocumentArtifactView.tsx @@ -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'); diff --git a/frontend/src/components/artifactViewUtils.test.ts b/frontend/src/components/artifactViewUtils.test.ts index 33f6794a..5f08d253 100644 --- a/frontend/src/components/artifactViewUtils.test.ts +++ b/frontend/src/components/artifactViewUtils.test.ts @@ -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('

hi

', { + headers: { 'Content-Type': 'text/html' }, + }); + expect(await readPresignedUrlEnvelope(res)).toBeNull(); + expect(await res.text()).toBe('

hi

'); + }); + + 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( diff --git a/frontend/src/components/artifactViewUtils.ts b/frontend/src/components/artifactViewUtils.ts index b5b12ee6..75b7420b 100644 --- a/frontend/src/components/artifactViewUtils.ts +++ b/frontend/src/components/artifactViewUtils.ts @@ -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 { + 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 { 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'), diff --git a/frontend/src/components/useArtifactBytes.ts b/frontend/src/components/useArtifactBytes.ts index afa62f06..d58b77e7 100644 --- a/frontend/src/components/useArtifactBytes.ts +++ b/frontend/src/components/useArtifactBytes.ts @@ -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(); diff --git a/tests/agents/test_workflow_agent_graph.py b/tests/agents/test_workflow_agent_graph.py index f5f17c3a..f2726a3c 100644 --- a/tests/agents/test_workflow_agent_graph.py +++ b/tests/agents/test_workflow_agent_graph.py @@ -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)) diff --git a/tests/agents/test_workflow_code_node.py b/tests/agents/test_workflow_code_node.py index 5f693961..a1961ba7 100644 --- a/tests/agents/test_workflow_code_node.py +++ b/tests/agents/test_workflow_code_node.py @@ -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. # --------------------------------------------------------------------------- diff --git a/tests/agents/tools/test_artifact_ref.py b/tests/agents/tools/test_artifact_ref.py index a7d2b505..429ffb4c 100644 --- a/tests/agents/tools/test_artifact_ref.py +++ b/tests/agents/tools/test_artifact_ref.py @@ -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 diff --git a/tests/agents/tools/test_read_document_tool.py b/tests/agents/tools/test_read_document_tool.py index f4160b58..31192687 100644 --- a/tests/agents/tools/test_read_document_tool.py +++ b/tests/agents/tools/test_read_document_tool.py @@ -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"} diff --git a/tests/api/user/test_artifacts_routes.py b/tests/api/user/test_artifacts_routes.py index 23efd0b1..a796b0d2 100644 --- a/tests/api/user/test_artifacts_routes.py +++ b/tests/api/user/test_artifacts_routes.py @@ -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 diff --git a/tests/parser/test_document_reader.py b/tests/parser/test_document_reader.py index 06019a7b..248a6bd8 100644 --- a/tests/parser/test_document_reader.py +++ b/tests/parser/test_document_reader.py @@ -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 diff --git a/tests/storage/db/repositories/test_artifacts.py b/tests/storage/db/repositories/test_artifacts.py index 31302d15..3673ca01 100644 --- a/tests/storage/db/repositories/test_artifacts.py +++ b/tests/storage/db/repositories/test_artifacts.py @@ -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) diff --git a/tests/test_code_executor_tool.py b/tests/test_code_executor_tool.py index 333afc80..f2a33ecf 100644 --- a/tests/test_code_executor_tool.py +++ b/tests/test_code_executor_tool.py @@ -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"] diff --git a/tests/worker/test_parse_document_worker.py b/tests/worker/test_parse_document_worker.py index 62e36829..1a762587 100644 --- a/tests/worker/test_parse_document_worker.py +++ b/tests/worker/test_parse_document_worker.py @@ -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