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