diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index 57a1d380..03aaf1a4 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -48,6 +48,10 @@ _SESSION_ID_RE = re.compile(r"[^A-Za-z0-9_-]+") # egress-open sandbox code (see ``_json_safe_state``). _CODE_STATE_EXCLUDED_KEYS = frozenset({"chat_history"}) +# Sentinel id stamped on the synthetic text attachment that flags documents a +# node skipped past the per-node blocking-extract cap (callers/tests match on it). +_EXTRACT_TRUNCATION_ID = "workflow-extract-truncated" + class WorkflowEngine: MAX_EXECUTION_STEPS = 50 @@ -511,6 +515,7 @@ class WorkflowEngine: supported = set(supported_types) supports_images = any(t.startswith("image/") for t in supported) max_files = int(getattr(settings, "WORKFLOW_NODE_NATIVE_MAX_FILES", 5)) + extract_max = int(getattr(settings, "WORKFLOW_NODE_EXTRACT_MAX_FILES", 5)) max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024)) # One read-only connection for the whole batch; the resolved-version @@ -539,6 +544,8 @@ class WorkflowEngine: attachments: List[Dict[str, Any]] = [] native_count = 0 + extract_count = 0 + dropped_for_cap = 0 for artifact_id, artifact, version in resolved: storage_path = version["storage_path"] mime_type = version.get("mime_type") or "application/octet-stream" @@ -571,6 +578,22 @@ class WorkflowEngine: attachments.append({"id": artifact_id, "mime_type": mime_type, "path": storage_path}) native_count += 1 else: + # Inline-text mimes are read directly (cheap); other mimes route + # through the parsing worker -- a blocking, ~120s per-document call. + # Cap how many documents a single node sends down that path so a + # node referencing many non-native documents (e.g. the ``*`` token) + # can't serialize dozens of parses. Inline text is not capped. + needs_parse = not self._is_inline_text_mime(mime_type) + if needs_parse: + # Count (and gate on) the parse ATTEMPT, not the success: a + # timed-out/failed parse is the ~120s worst case we must bound, + # so it has to consume cap budget too. Otherwise a degraded + # parsing backend (every parse fails) never advances the count + # and the node keeps issuing blocking parses without limit. + if extract_count >= extract_max: + dropped_for_cap += 1 + continue + extract_count += 1 content = self._extract_attachment_text( artifact_id, storage_path, mime_type, filename, max_bytes ) @@ -585,8 +608,26 @@ class WorkflowEngine: attachments.append( {"id": artifact_id, "mime_type": "text/plain", "content": content} ) + if dropped_for_cap: + logger.warning( + "Workflow node %s: blocking-extract cap (%d) reached; %d document(s) omitted", + node_title, extract_max, dropped_for_cap, + ) + attachments.append( + self._extract_truncation_note(node_title, extract_max, dropped_for_cap) + ) return attachments + @staticmethod + def _extract_truncation_note(node_title: str, cap: int, dropped: int) -> Dict[str, Any]: + """Build a non-fatal text attachment flagging documents skipped past the per-node extract cap.""" + content = ( + f'[Notice] Only the first {cap} document(s) for node "{node_title}" were extracted ' + f"to text; {dropped} additional document(s) were omitted to bound execution time. " + "Reference fewer documents, or use a model that reads them natively." + ) + return {"id": _EXTRACT_TRUNCATION_ID, "mime_type": "text/plain", "content": content} + @staticmethod def _agent_supported_attachment_types(node_agent: "BaseAgent") -> List[str]: """Return the provider's authoritative supported attachment mime types (handler's source).""" @@ -654,7 +695,13 @@ class WorkflowEngine: args=[artifact_id, {"workflow_run_id": self.workflow_run_id}, user_id, options], queue=queue, ) - result = async_result.get(timeout=timeout) + # A workflow can run inside a Celery worker (scheduled runs / webhooks). + # 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 + # explicitly (mirrors application/agents/tools/read_document.py). + result = async_result.get(timeout=timeout, disable_sync_subtasks=False) except (CeleryTimeoutError, TimeoutError): logger.warning("Workflow node: document parse timed out for %s", artifact_id) return None diff --git a/application/api/user/artifacts/routes.py b/application/api/user/artifacts/routes.py index 19824e10..daf2d693 100644 --- a/application/api/user/artifacts/routes.py +++ b/application/api/user/artifacts/routes.py @@ -5,7 +5,15 @@ from __future__ import annotations import re from typing import Optional -from flask import current_app, jsonify, make_response, redirect, request +from flask import ( + Response, + current_app, + jsonify, + make_response, + redirect, + request, + stream_with_context, +) from flask_restx import Namespace, Resource from application.api import api @@ -183,6 +191,53 @@ class GetArtifact(Resource): current_app.logger.error(f"Error retrieving artifact: {err}", exc_info=True) return make_response(jsonify({"success": False}), 400) + @api.doc(description="Delete an artifact and all its versions (owner only)") + def delete(self, artifact_id: str): + if not looks_like_uuid(artifact_id): + return make_response( + jsonify({"success": False, "message": "Artifact not found"}), 404 + ) + user_id = resolve_authenticated_user() + try: + with db_session() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact(artifact_id) + if artifact is None: + return make_response( + jsonify({"success": False, "message": "Artifact not found"}), 404 + ) + # Delete is a WRITE: only the parent owner may delete; share + # links / read-only collaborators are denied (read access only). + if not authorize_artifact_write(conn, artifact, user_id): + return make_response( + jsonify({"success": False, "message": "Forbidden"}), 403 + ) + storage_paths = repo.delete_artifact(artifact_id) + # Reap the bytes best-effort AFTER the row delete commits. + _reap_storage(storage_paths) + return make_response(jsonify({"success": True}), 200) + except Exception as err: + current_app.logger.error(f"Error deleting artifact: {err}", exc_info=True) + return make_response(jsonify({"success": False}), 500) + + +def _reap_storage(paths: list) -> None: + """Best-effort delete artifact bytes after the DB rows are gone; never raises.""" + if not paths: + return + try: + storage = StorageCreator.get_storage() + except Exception: + current_app.logger.warning("artifact delete: storage unavailable", exc_info=True) + return + for path in paths: + try: + storage.delete_file(path) + except Exception: + current_app.logger.warning( + "artifact delete: failed to delete bytes %s", path, exc_info=True + ) + @artifacts_ns.route("/artifacts//versions/") class GetArtifactVersion(Resource): @@ -294,13 +349,26 @@ class DownloadArtifact(Resource): ) return redirect(url, code=302) + # Stream the bytes in chunks instead of buffering the whole object in + # worker memory (artifacts can be many MB); close the handle when done. file_obj = storage.get_file(storage_path) - response = make_response(file_obj.read()) - response.headers.set("Content-Type", mime_type) - response.headers.set( - "Content-Disposition", f'attachment; filename="{filename}"' + + def _stream(): + try: + for chunk in iter(lambda: file_obj.read(65536), b""): + yield chunk + finally: + close = getattr(file_obj, "close", None) + if callable(close): + close() + + return Response( + stream_with_context(_stream()), + mimetype=mime_type, + headers={ + "Content-Disposition": f'attachment; filename="{filename}"' + }, ) - return response except FileNotFoundError: return make_response( jsonify({"success": False, "message": "File not found"}), 404 diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index eb271870..bd3d644a 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -254,7 +254,29 @@ def process_agent_webhook(self, agent_id, payload, idempotency_key=None): # separate parsing worker and never self-deadlocks the awaiting worker. @celery.task(bind=True, acks_late=False, autoretry_for=()) def parse_document(self, artifact_id, parent, user_id, options=None): - return parse_document_worker(self, artifact_id, parent, user_id, options or {}) + """Parse an input artifact on the parsing queue; self-terminate at the soft time limit.""" + from celery.exceptions import SoftTimeLimitExceeded + + try: + return parse_document_worker(self, artifact_id, parent, user_id, options or {}) + except SoftTimeLimitExceeded: + # A pathological/malicious document must not pin a parsing-worker slot past the + # window the caller already abandoned. Return the worker's clean error shape so + # the slot frees and the Redis result backend still gets a terminal result. + limit = getattr(self, "soft_time_limit", None) + suffix = f" after {int(limit)}s" if limit else "" + return {"status": "error", "error": f"document parsing timed out{suffix}."} + + +# Bind the soft limit to DOCUMENT_PARSE_TIMEOUT (the same window read_document awaits) so +# the prefork worker self-terminates a runaway parse instead of pinning the slot; the hard +# limit is the SIGKILL backstop if the soft handler can't unwind in time. +try: + from application.core.settings import settings as _parse_settings + parse_document.soft_time_limit = int(_parse_settings.DOCUMENT_PARSE_TIMEOUT) + parse_document.time_limit = parse_document.soft_time_limit + 30 +except Exception: + pass @celery.task(**DURABLE_TASK) diff --git a/application/core/settings.py b/application/core/settings.py index 58288c25..16723e41 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -362,6 +362,12 @@ class Settings(BaseSettings): # inputs). Files past the cap are extracted to text or dropped, not attached # natively, to bound context/cost. Re-uses SANDBOX_MAX_INPUT_BYTES per file. WORKFLOW_NODE_NATIVE_MAX_FILES: int = 5 + # Per-agent-node cap on documents extracted to text via the parsing worker. + # Each non-native, non-text document issues a separate blocking parse, so a + # node referencing many documents (e.g. the ``*`` token) is bounded here to + # avoid serializing dozens of parses; documents past the cap are skipped with + # a truncation note instead of extracted. + WORKFLOW_NODE_EXTRACT_MAX_FILES: int = 5 # Runner container resource caps — consumed by the docsgpt-sandbox compose # service (deployment/sandbox), not by the app client. cgroup CPU/mem caps # are part of the untrusted-code security boundary. diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index 05d17706..ca9f3d8b 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -1000,6 +1000,10 @@ class LLMHandler(ABC): # They will be added on resume to keep call/result pairs together. continue + # One assistant(tool_calls) message per call: track whether the + # success path already appended it so the except below doesn't + # add a second one when create_tool_message fails post-append. + assistant_appended = False try: self.tool_calls.append(call) tool_executor_gen = agent._execute_tool_action(tools_dict, call) @@ -1041,41 +1045,53 @@ class LLMHandler(ABC): if reasoning_content: assistant_msg["reasoning_content"] = reasoning_content updated_messages.append(assistant_msg) + assistant_appended = True updated_messages.append(self.create_tool_message(call, tool_response)) except Exception as e: logger.error(f"Error executing tool: {str(e)}", exc_info=True) + # The error tool message's tool_call_id must match the assistant + # tool_call id that precedes it. When the success path already + # appended one (create_tool_message failed afterwards) that id is + # the executor-returned ``call_id``; otherwise the except builds + # its own assistant message below from ``call.id``. + error_id = call_id if assistant_appended else call.id error_call = ToolCall( - id=call.id, name=call.name, arguments=call.arguments + id=error_id, name=call.name, arguments=call.arguments ) error_response = f"Error executing tool: {str(e)}" # Mirror the success path: a role:"tool" message must follow # an assistant message carrying its tool_calls, or the next # provider completion 400s ("'tool' must be a response to a - # preceding message with 'tool_calls'"). - args_str = ( - json.dumps(call.arguments) - if isinstance(call.arguments, dict) - else call.arguments - ) - tool_call_obj = { - "id": call.id, - "type": "function", - "function": { - "name": call.name, - "arguments": args_str, - }, - } - if call.thought_signature: - tool_call_obj["thought_signature"] = call.thought_signature - assistant_msg: Dict[str, Any] = { - "role": "assistant", - "content": None, - "tool_calls": [tool_call_obj], - } - if reasoning_content: - assistant_msg["reasoning_content"] = reasoning_content - updated_messages.append(assistant_msg) + # preceding message with 'tool_calls'"). Skip it when the + # success path already appended one for this call — a + # create_tool_message failure after that append would + # otherwise duplicate the assistant message and 400 the + # same way an orphan tool message does. + if not assistant_appended: + args_str = ( + json.dumps(call.arguments) + if isinstance(call.arguments, dict) + else call.arguments + ) + tool_call_obj = { + "id": call.id, + "type": "function", + "function": { + "name": call.name, + "arguments": args_str, + }, + } + if call.thought_signature: + tool_call_obj["thought_signature"] = call.thought_signature + assistant_msg: Dict[str, Any] = { + "role": "assistant", + "content": None, + "tool_calls": [tool_call_obj], + } + if reasoning_content: + assistant_msg["reasoning_content"] = reasoning_content + updated_messages.append(assistant_msg) error_message = self.create_tool_message(error_call, error_response) updated_messages.append(error_message) diff --git a/application/parser/document_reader.py b/application/parser/document_reader.py index ec884840..1d18ac5b 100644 --- a/application/parser/document_reader.py +++ b/application/parser/document_reader.py @@ -139,6 +139,17 @@ def _parse_to_text(parser: Any, path: Path) -> str: return str(parsed) +def _is_docling_parser(parser: Any) -> bool: + """True when ``parser`` is Docling-backed (collecting tables would otherwise re-convert).""" + if parser is None: + return False + try: + from application.parser.file.docling_parser import DoclingParser + except Exception: + return False + return isinstance(parser, DoclingParser) + + def _compact_table(table: Dict[str, Any]) -> Dict[str, Any]: """Bound a single table's rows and cell sizes so one giant table can't bloat context.""" @@ -159,13 +170,27 @@ def _compact_table(table: Dict[str, Any]) -> Dict[str, Any]: return compact -def _docling_structured(path: Path, *, ocr_enabled: bool, include_tables: bool) -> Dict[str, Any]: - """Convert a document with Docling and return markdown + structured dict + bounded tables.""" - from docling.document_converter import DocumentConverter +def _docling_structured( + path: Path, *, ocr_enabled: bool, include_tables: bool, parser: Any = None +) -> Dict[str, Any]: + """Convert a document with Docling and return markdown + structured dict + bounded tables. - converter = DocumentConverter() - doc = converter.convert(str(path)).document - markdown = doc.export_to_markdown() + When ``parser`` is the configured ``DoclingParser`` (the collapse-the-double-convert + path), reuse ITS converter + export so OCR/pipeline options and the export fallback + are honored and the content matches the legacy single-parse output exactly; otherwise + fall back to a vanilla converter. + """ + if _is_docling_parser(parser): + if getattr(parser, "_converter", None) is None: + parser._init_parser() + doc = parser._converter.convert(str(path)).document + markdown = parser._export_content(doc) + else: + from docling.document_converter import DocumentConverter + + converter = DocumentConverter() + doc = converter.convert(str(path)).document + markdown = doc.export_to_markdown() structured = doc.export_to_dict() tables: List[Dict[str, Any]] = [] if include_tables: @@ -323,6 +348,28 @@ def _shape( } parser = _pick_parser(suffix, ocr_enabled=ocr_enabled, engine=engine) + wants_tables = include_tables and engine != "fast" and output != "chunks" + + # A Docling-backed parser already converts the whole document to produce its text. + # When tables are also requested, reuse that single conversion for both the markdown + # content and the tables instead of converting a second time just to collect tables + # (Docling/torch conversion dominates the cost, so a re-convert ~doubles latency). + if wants_tables and _is_docling_parser(parser): + try: + extracted = _docling_structured( + path, ocr_enabled=ocr_enabled, include_tables=True, parser=parser + ) + text = extracted["markdown"] + tables: List[Dict[str, Any]] = [_compact_table(t) for t in extracted["tables"]] + except Exception: + text, tables = _parse_to_text(parser, path), [] + text = _apply_pages(text, pages) + bounded, truncated = _bounded(text, max_chars) + payload: Dict[str, Any] = {"output": output, "content": bounded, "truncated": truncated} + if tables: + payload["tables"] = tables + return payload + if parser is None: # A whitelisted extension with no dedicated parser (e.g. .txt) reads as plain # text, matching SimpleDirectoryReader's standard-read fallback. @@ -335,7 +382,7 @@ def _shape( return {"output": "chunks", "chunks": _to_chunks(text, max_chars), "truncated": False} tables: List[Dict[str, Any]] = [] - if include_tables and engine != "fast": + if wants_tables: try: tables = [_compact_table(t) for t in _docling_structured( path, ocr_enabled=ocr_enabled, include_tables=True)["tables"]] diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py index edbfbb98..b75af252 100644 --- a/application/storage/db/repositories/artifacts.py +++ b/application/storage/db/repositories/artifacts.py @@ -301,6 +301,72 @@ class ArtifactsRepository: ).fetchone() return int(row[0]) if row is not None else 0 + def delete_artifact(self, artifact_id: str) -> list[str]: + """Delete an artifact (+ its versions, via FK cascade); return its stored paths. + + The caller deletes the returned storage objects best-effort AFTER commit; + the DB delete also frees the owner's count/byte quota. + """ + paths = [ + v["storage_path"] + for v in self.list_versions(artifact_id) + if v.get("storage_path") + ] + self._conn.execute( + text("DELETE FROM artifacts WHERE id = CAST(:id AS uuid)"), + {"id": artifact_id}, + ) + return paths + + def storage_paths_for_conversation(self, conversation_id: str) -> list[str]: + """Return every stored version path for a conversation's artifacts (for byte cleanup).""" + result = self._conn.execute( + text( + "SELECT v.storage_path FROM artifact_versions v " + "JOIN artifacts a ON a.id = v.artifact_id " + "WHERE a.conversation_id = CAST(:cid AS uuid) " + "AND v.storage_path IS NOT NULL" + ), + {"cid": conversation_id}, + ) + return [r[0] for r in result.fetchall()] + + def storage_paths_for_user_conversations(self, user_id: str) -> list[str]: + """Return every stored version path for artifacts under a user's conversations.""" + result = self._conn.execute( + text( + "SELECT v.storage_path FROM artifact_versions v " + "JOIN artifacts a ON a.id = v.artifact_id " + "JOIN conversations c ON c.id = a.conversation_id " + "WHERE c.user_id = :user_id AND v.storage_path IS NOT NULL" + ), + {"user_id": user_id}, + ) + return [r[0] for r in result.fetchall()] + + def delete_for_conversation(self, conversation_id: str) -> int: + """Delete all artifacts (+ versions, via cascade) parented to a conversation. + + ``conversation_id`` is a bare uuid column (no FK), so the parent delete + does not cascade — callers invoke this to reclaim the rows + their quota. + """ + result = self._conn.execute( + text("DELETE FROM artifacts WHERE conversation_id = CAST(:cid AS uuid)"), + {"cid": conversation_id}, + ) + return result.rowcount + + def delete_for_user_conversations(self, user_id: str) -> int: + """Delete every artifact parented to one of a user's conversations (+ versions).""" + result = self._conn.execute( + text( + "DELETE FROM artifacts WHERE conversation_id IN " + "(SELECT id FROM conversations WHERE user_id = :user_id)" + ), + {"user_id": user_id}, + ) + return result.rowcount + def get_version(self, artifact_id: str, version: int) -> Optional[dict]: """Fetch a single version row by artifact id and version number.""" result = self._conn.execute( diff --git a/application/storage/db/repositories/conversations.py b/application/storage/db/repositories/conversations.py index b8e3ee25..65787310 100644 --- a/application/storage/db/repositories/conversations.py +++ b/application/storage/db/repositories/conversations.py @@ -15,6 +15,7 @@ Covers every operation the legacy Mongo code performs on from __future__ import annotations import json +import logging from enum import Enum from typing import Optional @@ -25,6 +26,8 @@ from application.storage.db.base_repository import looks_like_uuid, row_to_dict from application.storage.db.models import conversations_table, conversation_messages_table from application.storage.db.serialization import PGNativeJSONEncoder +logger = logging.getLogger(__name__) + class MessageUpdateOutcome(str, Enum): """Discriminated result of ``update_message_by_id``. @@ -535,6 +538,14 @@ class ConversationsRepository: # a non-UUID id reaches this code path. if not looks_like_uuid(conversation_id): return False + # Artifacts are parent-derived: deleting a conversation must take its + # artifacts with it (rows + bytes + the quota they consume). conversation_id + # is a bare uuid (no FK cascade), so delete the rows explicitly and reap the + # stored bytes best-effort. + from application.storage.db.repositories.artifacts import ArtifactsRepository + + artifacts = ArtifactsRepository(self._conn) + paths = artifacts.storage_paths_for_conversation(conversation_id) result = self._conn.execute( text( "DELETE FROM conversations " @@ -542,15 +553,45 @@ class ConversationsRepository: ), {"id": conversation_id, "user_id": user_id}, ) - return result.rowcount > 0 + deleted = result.rowcount > 0 + if deleted: + artifacts.delete_for_conversation(conversation_id) + self._reap_artifact_storage(paths) + return deleted def delete_all_for_user(self, user_id: str) -> int: + from application.storage.db.repositories.artifacts import ArtifactsRepository + + # Reap artifacts BEFORE the conversations: delete_for_user_conversations + # resolves them via a subquery over the still-present conversation rows. + artifacts = ArtifactsRepository(self._conn) + paths = artifacts.storage_paths_for_user_conversations(user_id) + artifacts.delete_for_user_conversations(user_id) result = self._conn.execute( text("DELETE FROM conversations WHERE user_id = :user_id"), {"user_id": user_id}, ) + self._reap_artifact_storage(paths) return result.rowcount + @staticmethod + def _reap_artifact_storage(paths: list) -> None: + """Best-effort delete artifact bytes whose rows were cascade-removed; never raises.""" + if not paths: + return + try: + from application.storage.storage_creator import StorageCreator + + storage = StorageCreator.get_storage() + except Exception: + logger.warning("conversation delete: storage unavailable for artifact cleanup", exc_info=True) + return + for path in paths: + try: + storage.delete_file(path) + except Exception: + logger.warning("conversation delete: failed to delete artifact bytes %s", path, exc_info=True) + # ------------------------------------------------------------------ # Messages # ------------------------------------------------------------------ diff --git a/deployment/docker-compose-azure.yaml b/deployment/docker-compose-azure.yaml index 71f675a8..bbcd9f44 100644 --- a/deployment/docker-compose-azure.yaml +++ b/deployment/docker-compose-azure.yaml @@ -18,8 +18,17 @@ services: - CELERY_BROKER_URL=redis://redis:6379/0 - CELERY_RESULT_BACKEND=redis://redis:6379/1 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt + # Code-execution runner reached over HTTP + WebSocket (no docker socket). + - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 + # Env-scrubbing kernelspec the docsgpt-sandbox image installs. + - SANDBOX_KERNEL_NAME=docsgpt-python ports: - "7091:7091" + networks: + # `default` reaches the sandbox gateway; `data-net` reaches redis/postgres + # (which are NOT on the sandbox's networks). + - default + - data-net volumes: - ../application/indexes:/app/application/indexes - ../application/inputs:/app/application/inputs @@ -42,19 +51,54 @@ services: - CELERY_RESULT_BACKEND=redis://redis:6379/1 - API_URL=http://backend:7091 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt + - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 + - SANDBOX_KERNEL_NAME=docsgpt-python + networks: + - default + - data-net depends_on: redis: condition: service_started postgres: condition: service_healthy + # Always-on code-execution runner (Jupyter Kernel Gateway). In-process kernels, + # no Docker socket, no published host port. On the internal-only sandbox-net + # plus `default` (for internet egress); deliberately NOT on data-net, so + # sandboxed code cannot reach the unauthenticated redis broker (Celery task + # injection -> RCE) or postgres. Do NOT add `env_file: ../.env` here -- the + # runner needs no app secrets. For untrusted multi-tenant workloads use a + # per-session VM via SANDBOX_BACKEND=daytona instead. + docsgpt-sandbox: + build: ./sandbox + mem_limit: ${SANDBOX_MEMORY:-1g} + cpus: ${SANDBOX_CPUS:-1.0} + pids_limit: 256 + read_only: true + environment: + # Keep Jupyter's runtime/connection files on the writable tmpfs. + - JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime + - JUPYTER_DATA_DIR=/tmp/jupyter-data + tmpfs: + - /tmp + networks: + - sandbox-net + - default + redis: image: redis:6-alpine + # Data plane only: kept off the sandbox's networks (a reachable, unauthed + # broker = Celery task injection -> worker RCE). + networks: + - data-net ports: - 6379:6379 postgres: image: postgres:16-alpine + # Data plane only (see redis): kept off the sandbox's networks. + networks: + - data-net environment: - POSTGRES_USER=docsgpt - POSTGRES_PASSWORD=docsgpt @@ -69,5 +113,15 @@ services: timeout: 5s retries: 10 +networks: + # Internal-only runner network (no external egress via this net; the sandbox + # still reaches the internet via the default bridge). + sandbox-net: + internal: true + # Internal data plane: redis/postgres + backend/worker. The sandbox is NOT + # attached, so sandboxed code cannot reach the broker or database. + data-net: + internal: true + volumes: postgres_data: diff --git a/deployment/docker-compose-hub.yaml b/deployment/docker-compose-hub.yaml index b20b1e55..54551023 100644 --- a/deployment/docker-compose-hub.yaml +++ b/deployment/docker-compose-hub.yaml @@ -23,8 +23,17 @@ services: - CELERY_RESULT_BACKEND=redis://redis:6379/1 - CACHE_REDIS_URL=redis://redis:6379/2 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt + # Code-execution runner reached over HTTP + WebSocket (no docker socket). + - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 + # Env-scrubbing kernelspec the docsgpt-sandbox image installs. + - SANDBOX_KERNEL_NAME=docsgpt-python ports: - "7091:7091" + networks: + # `default` reaches the sandbox gateway; `data-net` reaches redis/postgres + # (which are NOT on the sandbox's networks). + - default + - data-net volumes: - ../application/indexes:/app/indexes - ../application/inputs:/app/inputs @@ -49,6 +58,11 @@ services: - API_URL=http://backend:7091 - CACHE_REDIS_URL=redis://redis:6379/2 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt + - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 + - SANDBOX_KERNEL_NAME=docsgpt-python + networks: + - default + - data-net volumes: - ../application/indexes:/app/indexes - ../application/inputs:/app/inputs @@ -59,13 +73,46 @@ services: postgres: condition: service_healthy + # Always-on code-execution runner (Jupyter Kernel Gateway). In-process kernels, + # no Docker socket, no published host port. On the internal-only sandbox-net + # plus `default` (for internet egress); deliberately NOT on data-net, so + # sandboxed code cannot reach the unauthenticated redis broker (Celery task + # injection -> RCE) or postgres. Do NOT add `env_file: ../.env` here -- the + # runner needs no app secrets. For untrusted multi-tenant workloads use a + # per-session VM via SANDBOX_BACKEND=daytona instead. + docsgpt-sandbox: + # Built from the repo (no published arc53/docsgpt-sandbox image exists yet); the + # hub compose ships alongside ./sandbox, so this stays self-contained and the + # stack always starts. Switch to a published image once CI builds/pushes one. + build: ./sandbox + mem_limit: ${SANDBOX_MEMORY:-1g} + cpus: ${SANDBOX_CPUS:-1.0} + pids_limit: 256 + read_only: true + environment: + # Keep Jupyter's runtime/connection files on the writable tmpfs. + - JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime + - JUPYTER_DATA_DIR=/tmp/jupyter-data + tmpfs: + - /tmp + networks: + - sandbox-net + - default + redis: image: redis:6-alpine + # Data plane only: kept off the sandbox's networks (a reachable, unauthed + # broker = Celery task injection -> worker RCE). + networks: + - data-net ports: - 6379:6379 postgres: image: postgres:16-alpine + # Data plane only (see redis): kept off the sandbox's networks. + networks: + - data-net environment: - POSTGRES_USER=docsgpt - POSTGRES_PASSWORD=docsgpt @@ -80,5 +127,15 @@ services: timeout: 5s retries: 10 +networks: + # Internal-only runner network (no external egress via this net; the sandbox + # still reaches the internet via the default bridge). + sandbox-net: + internal: true + # Internal data plane: redis/postgres + backend/worker. The sandbox is NOT + # attached, so sandboxed code cannot reach the broker or database. + data-net: + internal: true + volumes: postgres_data: diff --git a/tests/agents/test_workflow_input_documents.py b/tests/agents/test_workflow_input_documents.py index b91f5326..cb14d872 100644 --- a/tests/agents/test_workflow_input_documents.py +++ b/tests/agents/test_workflow_input_documents.py @@ -16,7 +16,11 @@ import pytest from sqlalchemy import text from application.agents.workflow_agent import WorkflowAgent, _MAX_INPUT_DOCUMENTS -from application.agents.workflows.workflow_engine import WorkflowEngine +from application.agents.workflows.schemas import AgentNodeConfig +from application.agents.workflows.workflow_engine import ( + _EXTRACT_TRUNCATION_ID, + WorkflowEngine, +) from application.storage.db.repositories.artifacts import ArtifactsRepository from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository from application.storage.local import LocalStorage @@ -350,3 +354,142 @@ def test_no_attachments_run_still_works(pg_engine, tmp_path, monkeypatch): ).scalar() assert run is not None assert n == 0 + + +def test_extract_parse_opts_out_of_sync_subtask_guard(monkeypatch): + """_parse_document_text awaits with disable_sync_subtasks=False so it works inside a Celery worker.""" + agent = _agent(str(uuid.uuid4()), []) + engine = WorkflowEngine.__new__(WorkflowEngine) + engine.agent = agent + engine.workflow_run_id = "run-extract" + + import application.api.user.tasks as tasks + + captured: dict = {} + + class _FakeAsyncResult: + def __init__(self): + self.get_kwargs = None + + def get(self, timeout=None, disable_sync_subtasks=True): + self.get_kwargs = {"timeout": timeout, "disable_sync_subtasks": disable_sync_subtasks} + return {"status": "ok", "content": "parsed markdown"} + + def _apply_async(args=None, queue=None, **kw): + result = _FakeAsyncResult() + captured["result"] = result + captured["args"] = args + return result + + monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async) + + out = engine._parse_document_text("artifact-xyz") + + assert out == "parsed markdown" + # A prefork worker's task_join_will_block() is process-wide, so the await must + # opt out of the guard or get() raises RuntimeError("Never call result.get()..."). + assert captured["result"].get_kwargs["disable_sync_subtasks"] is False + # The run-scoped parent + resolved id reached the parsing task. + assert captured["args"][0] == "artifact-xyz" + assert captured["args"][1] == {"workflow_run_id": "run-extract"} + + +def test_node_extract_path_capped_with_truncation_note(pg_engine, tmp_path, monkeypatch): + """A node referencing more docs than the extract cap parses only up to the cap and notes the rest.""" + storage = _wire(pg_engine, tmp_path, monkeypatch) + wf_id = _make_workflow(pg_engine) + + # docx is non-native (no vision) and not inline-text, so each routes through + # the blocking parsing worker -- the path the per-node cap must bound. + docx = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + attachments = [ + _stage_attachment(storage, f"doc-{i}".encode(), f"f{i}.docx", docx) + for i in range(4) + ] + agent = _agent(wf_id, attachments) + _patch_engine(monkeypatch) + list(agent._gen_inner("summarize", log_context=None)) + engine = _RecordingEngine.instances[-1] + + # Cap the blocking-extract path below the doc count so the overflow truncates. + from application.core.settings import settings + + monkeypatch.setattr(settings, "WORKFLOW_NODE_EXTRACT_MAX_FILES", 2, raising=False) + + # Stub the parsing worker so each non-text doc "parses" without a broker, and + # count the blocking calls to prove the overflow docs are never enqueued. + import application.api.user.tasks as tasks + + parse_calls = {"n": 0} + + class _R: + def get(self, timeout=None, disable_sync_subtasks=True): + return {"status": "ok", "content": "PARSED"} + + def _apply_async(args=None, queue=None, **kw): + parse_calls["n"] += 1 + return _R() + + monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async) + + node_config = AgentNodeConfig(input_documents=["*"]) + out = engine._materialize_node_attachments(node_config, "Reviewer", supported_types=[]) + + notes = [a for a in out if a.get("id") == _EXTRACT_TRUNCATION_ID] + extracted = [a for a in out if a.get("id") != _EXTRACT_TRUNCATION_ID] + # Only the cap was extracted; the remaining docs were never sent to the worker. + assert len(extracted) == 2 + assert parse_calls["n"] == 2 + # A single non-fatal truncation note is appended to the node's inlined text. + assert len(notes) == 1 + assert notes[0]["mime_type"] == "text/plain" + assert "omitted" in notes[0]["content"].lower() + + +def test_node_extract_cap_bounds_parse_attempts_even_when_every_parse_times_out( + pg_engine, tmp_path, monkeypatch +): + """The cap must bound parse ATTEMPTS, not successes: a degraded backend where every + parse times out (~120s each) must still issue at most the cap's worth of blocking calls.""" + from celery.exceptions import TimeoutError as CeleryTimeoutError + + storage = _wire(pg_engine, tmp_path, monkeypatch) + wf_id = _make_workflow(pg_engine) + + docx = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + attachments = [ + _stage_attachment(storage, f"doc-{i}".encode(), f"f{i}.docx", docx) + for i in range(4) + ] + agent = _agent(wf_id, attachments) + _patch_engine(monkeypatch) + list(agent._gen_inner("summarize", log_context=None)) + engine = _RecordingEngine.instances[-1] + + from application.core.settings import settings + + monkeypatch.setattr(settings, "WORKFLOW_NODE_EXTRACT_MAX_FILES", 2, raising=False) + + import application.api.user.tasks as tasks + + parse_calls = {"n": 0} + + class _R: + def get(self, timeout=None, disable_sync_subtasks=True): + raise CeleryTimeoutError() # the ~120s worst case the cap must bound + + def _apply_async(args=None, queue=None, **kw): + parse_calls["n"] += 1 + return _R() + + monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async) + + node_config = AgentNodeConfig(input_documents=["*"]) + out = engine._materialize_node_attachments(node_config, "Reviewer", supported_types=[]) + + # Every parse timed out (nothing extracted), but blocking attempts were bounded. + extracted = [a for a in out if a.get("id") != _EXTRACT_TRUNCATION_ID] + assert extracted == [] + assert parse_calls["n"] == 2 # not 4 -- failed parses still consume cap budget + notes = [a for a in out if a.get("id") == _EXTRACT_TRUNCATION_ID] + assert len(notes) == 1 diff --git a/tests/api/user/test_artifacts_routes.py b/tests/api/user/test_artifacts_routes.py index 0322830a..22638dcb 100644 --- a/tests/api/user/test_artifacts_routes.py +++ b/tests/api/user/test_artifacts_routes.py @@ -297,8 +297,9 @@ class TestSharedAccess: # Download # --------------------------------------------------------------------------- class _FakeStorage: - def __init__(self, data: bytes): + def __init__(self, data: bytes = b""): self._data = data + self.deleted: list = [] def get_file(self, path): return io.BytesIO(self._data) @@ -306,6 +307,10 @@ class _FakeStorage: def generate_presigned_url(self, path, expires_in=300): return f"https://signed.example/{path}?exp={expires_in}" + def delete_file(self, path): + self.deleted.append(path) + return True + @pytest.mark.unit class TestDownloadArtifact: @@ -343,6 +348,27 @@ class TestDownloadArtifact: assert resp.headers["Content-Disposition"] == 'attachment; filename="deck.pptx"' assert resp.headers["Content-Type"] == "application/vnd.ms-powerpoint" + def test_local_download_is_streamed_not_buffered( + self, _patch_db, flask_app, token_owner, monkeypatch + ): + # The response body must be a stream (generator), not the whole object + # buffered into memory via make_response(file_obj.read()). + from application.api.user.artifacts.routes import DownloadArtifact + + art = self._seed(_patch_db) + storage = _FakeStorage(b"Z" * 200_000) + monkeypatch.setattr( + "application.api.user.artifacts.routes.StorageCreator.get_storage", + lambda: storage, + ) + monkeypatch.setattr( + "application.api.user.artifacts.routes.settings.URL_STRATEGY", + "backend", raising=False, + ) + resp = _call(flask_app, DownloadArtifact, art["id"], token=token_owner) + assert resp.is_streamed + assert resp.get_data() == b"Z" * 200_000 + def test_s3_strategy_redirects_to_presigned( self, _patch_db, flask_app, token_owner, monkeypatch ): @@ -506,6 +532,91 @@ class TestRestoreArtifact: assert resp.status_code == 403 +@pytest.mark.unit +class TestDeleteArtifact: + def _seed(self, conn, **over): + conv = _make_conversation(conn) + defaults = dict( + conversation_id=str(conv["id"]), + filename="f.bin", + storage_path="inputs/owner/artifacts/x/v1/f.bin", + ) + defaults.update(over) + return _make_artifact(conn, **defaults) + + def test_owner_deletes_and_reaps_bytes( + self, _patch_db, flask_app, token_owner, monkeypatch + ): + from application.api.user.artifacts.routes import GetArtifact + + art = self._seed(_patch_db) + storage = _FakeStorage() + monkeypatch.setattr( + "application.api.user.artifacts.routes.StorageCreator.get_storage", + lambda: storage, + ) + resp = _call(flask_app, GetArtifact, art["id"], token=token_owner, method="delete") + assert resp.status_code == 200 + assert ArtifactsRepository(_patch_db).get_artifact(art["id"]) is None + assert "inputs/owner/artifacts/x/v1/f.bin" in storage.deleted + + def test_stranger_denied(self, _patch_db, flask_app): + from application.api.user.artifacts.routes import GetArtifact + + art = self._seed(_patch_db) + resp = _call( + flask_app, GetArtifact, art["id"], token={"sub": STRANGER}, method="delete" + ) + assert resp.status_code == 403 + assert ArtifactsRepository(_patch_db).get_artifact(art["id"]) is not None + + def test_share_token_holder_denied(self, _patch_db, flask_app): + from application.api.user.artifacts.routes import GetArtifact + + conv = _make_conversation(_patch_db) + art = _make_artifact(_patch_db, conversation_id=str(conv["id"])) + share = SharedConversationsRepository(_patch_db).create(str(conv["id"]), OWNER) + resp = _call( + flask_app, GetArtifact, art["id"], token=None, + query={"share_token": str(share["uuid"])}, method="delete", + ) + assert resp.status_code == 403 + assert ArtifactsRepository(_patch_db).get_artifact(art["id"]) is not None + + def test_unknown_artifact_404(self, _patch_db, flask_app, token_owner): + from application.api.user.artifacts.routes import GetArtifact + + resp = _call( + flask_app, GetArtifact, str(uuid.uuid4()), token=token_owner, method="delete" + ) + assert resp.status_code == 404 + + +@pytest.mark.unit +class TestConversationDeleteReapsArtifacts: + def test_delete_conversation_removes_artifacts_frees_quota_and_reaps_bytes( + self, _patch_db, flask_app, monkeypatch + ): + conv = _make_conversation(_patch_db) + art = _make_artifact( + _patch_db, conversation_id=str(conv["id"]), + storage_path="k/reap.bin", filename="r.bin", size=42, + ) + storage = _FakeStorage() + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + lambda: storage, + ) + + deleted = ConversationsRepository(_patch_db).delete(str(conv["id"]), OWNER) + assert deleted is True + + repo = ArtifactsRepository(_patch_db) + assert repo.get_artifact(art["id"]) is None # cascade-removed with the parent + assert repo.count_for_user(OWNER) == 0 # quota freed + assert "k/reap.bin" in storage.deleted # bytes reaped + + # --------------------------------------------------------------------------- # Malformed artifact_id path segment -> 404 (no CAST(:id AS uuid) txn poison) # --------------------------------------------------------------------------- diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index b28b101e..66f4fd02 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -327,6 +327,53 @@ class TestMcpOauthTask: assert result == {"url": "http://auth"} +class TestParseDocumentTask: + """parse_document runs on the parsing queue under a bounded time limit.""" + + @pytest.mark.unit + @patch("application.api.user.tasks.parse_document_worker") + def test_calls_parse_document_worker(self, mock_worker): + from application.api.user.tasks import parse_document + + mock_worker.return_value = {"status": "ok", "content": "hi"} + + result = parse_document( + "art-1", {"conversation_id": "c1"}, "user1", {"output": "markdown"}, + ) + + mock_worker.assert_called_once_with( + ANY, "art-1", {"conversation_id": "c1"}, "user1", + {"output": "markdown"}, + ) + assert result == {"status": "ok", "content": "hi"} + + @pytest.mark.unit + @patch("application.api.user.tasks.parse_document_worker") + def test_soft_time_limit_returns_clean_error(self, mock_worker): + from celery.exceptions import SoftTimeLimitExceeded + + from application.api.user.tasks import parse_document + + mock_worker.side_effect = SoftTimeLimitExceeded("parse") + + # The task must swallow the soft-limit signal and return the worker's + # clean error shape so the parsing-worker slot frees instead of crashing. + result = parse_document("art-1", {"conversation_id": "c1"}, "user1") + + assert result["status"] == "error" + assert "timed out" in result["error"] + + @pytest.mark.unit + def test_time_limits_derived_from_document_parse_timeout(self): + from application.api.user.tasks import parse_document + from application.core.settings import settings + + assert parse_document.soft_time_limit == settings.DOCUMENT_PARSE_TIMEOUT + assert parse_document.time_limit == settings.DOCUMENT_PARSE_TIMEOUT + 30 + # Hard limit must exceed the soft limit so the handler can unwind first. + assert parse_document.time_limit > parse_document.soft_time_limit + + class TestDurableTaskRetryPolicy: """The long-running tasks share a uniform retry policy.""" diff --git a/tests/llm/handlers/test_llm_handlers.py b/tests/llm/handlers/test_llm_handlers.py index 64369f91..90cbdf44 100644 --- a/tests/llm/handlers/test_llm_handlers.py +++ b/tests/llm/handlers/test_llm_handlers.py @@ -759,6 +759,66 @@ class TestHandleToolCalls: assistant_ids = {tc["id"] for tc in assistant_msg["tool_calls"]} assert tool_msg["tool_call_id"] in assistant_ids + def test_create_tool_message_failure_no_duplicate_assistant(self): + """A tool that succeeds but returns an unserializable result makes + create_tool_message raise *after* the success path appended the + assistant tool_calls message. The except branch must reuse that + message, not append a second one — exactly one assistant(tool_calls) + precedes the role:"tool" error, or the next completion 400s.""" + + class _SerializingHandler(ConcreteHandler): + def create_tool_message( + self, tool_call: ToolCall, result: Any + ) -> Dict: + # Mirror real handlers: JSON-encode non-str results, which + # raises on a value the encoder can't serialize. The error + # path passes a plain string, so only the success path trips. + import json as _json + + content = ( + result + if isinstance(result, str) + else _json.dumps(result) + ) + return { + "role": "tool", + "tool_call_id": tool_call.id, + "content": content, + } + + handler = _SerializingHandler() + agent = self._make_agent() + agent.tool_executor._name_to_tool = {} + + # Tool runs fine but returns a value json.dumps can't encode. + def fake_execute(tools_dict, call): + yield {"type": "tool_call", "data": {"status": "pending"}} + return (object(), call.id) + + agent._execute_tool_action = Mock(side_effect=fake_execute) + + call = ToolCall(id="c1", name="action_1", arguments="{}") + gen = handler.handle_tool_calls(agent, [call], {"1": {"name": "t"}}, []) + try: + while True: + next(gen) + except StopIteration as e: + messages, _pending = e.value + + # Exactly one assistant(tool_calls) for c1, then the tool error — the + # buggy path produced [assistant, assistant, tool]. + assert [m["role"] for m in messages] == ["assistant", "tool"] + assistant_msgs = [ + m for m in messages + if m.get("role") == "assistant" + and any(tc["id"] == "c1" for tc in m.get("tool_calls", [])) + ] + assert len(assistant_msgs) == 1 + tool_msg = messages[-1] + assert tool_msg["role"] == "tool" + assert tool_msg["tool_call_id"] == "c1" + assert "Error executing tool" in tool_msg["content"] + def test_partial_batch_failure_answers_every_tool_call(self): """One tool raises, others succeed: every assistant tool_call must get a matching role:"tool" reply (OpenAI rejects unanswered ids).""" diff --git a/tests/parser/test_document_reader.py b/tests/parser/test_document_reader.py index 06ced7d3..06019a7b 100644 --- a/tests/parser/test_document_reader.py +++ b/tests/parser/test_document_reader.py @@ -121,6 +121,94 @@ def test_fast_engine_uses_legacy_parser(monkeypatch): assert out["content"] == "legacy text" +# --------------------------------------------------------------------------- +# single Docling conversion: the default markdown+tables path must not re-convert +# --------------------------------------------------------------------------- +@pytest.mark.unit +def test_default_markdown_with_tables_converts_docling_once(monkeypatch): + """Default read (markdown, engine=auto, include_tables) converts the doc with Docling once. + + A Docling-backed parser already converts the whole document to produce its text; + collecting tables must reuse that single conversion instead of re-running + DocumentConverter. Counts conversions across both sites to prove no double pass. + """ + import sys + import types + + from application.parser.file.docling_parser import DoclingParser + + counter = {"instances": 0, "converts": 0} + + class _FakeTable: + def export_to_dataframe(self): + raise RuntimeError("no dataframe") # force the markdown fallback path + + def export_to_markdown(self): + return "| h |\n| - |\n| v |" + + class _FakeDoc: + tables = [_FakeTable()] + pages = {"1": {}} + + def export_to_markdown(self): + return "# single-pass content" + + def export_to_dict(self): + return {"texts": [{}], "tables": [{}], "pages": {"1": {}}} + + class _FakeResult: + document = _FakeDoc() + + class _CountingConverter: + def __init__(self, *args, **kwargs): + counter["instances"] += 1 + + def convert(self, *args, **kwargs): + counter["converts"] += 1 + return _FakeResult() + + fake_docling = types.ModuleType("docling") + fake_dc_module = types.ModuleType("docling.document_converter") + fake_dc_module.DocumentConverter = _CountingConverter + fake_docling.document_converter = fake_dc_module + monkeypatch.setitem(sys.modules, "docling", fake_docling) + monkeypatch.setitem(sys.modules, "docling.document_converter", fake_dc_module) + + class _FakeDoclingParser(DoclingParser): + """Docling-backed parser exposing its configured converter + export for reuse. + + The collapse path converts ONCE via ``self._converter`` and exports content via + ``_export_content`` (the configured pipeline/OCR), so both content and tables + come from a single conversion. + """ + + def __init__(self): + super().__init__() + self._parser_config = {"ready": True} # makes parser_config_set True + self._converter = _CountingConverter() # single configured conversion + + def _export_content(self, document): + return "# single-pass content" + + def parse_file(self, file, errors="ignore"): + # Only reached if the collapse path fails; the test asserts it doesn't. + return "# parse-file content" + + monkeypatch.setattr(dr, "get_default_file_extractor", lambda ocr_enabled=None: {".pdf": _FakeDoclingParser()}) + + # Defaults: output=markdown, engine=auto, include_tables=True -> the double-parse path. + out = parse_document_bytes(b"%PDF-1.4", "doc.pdf") + + assert out["output"] == "markdown" + # Content comes from the CONFIGURED parser's export (matches the legacy single + # parse), not a vanilla converter; the collapse path was taken, not the fallback. + assert "single-pass content" in out["content"] + assert "parse-file content" not in out["content"] + assert out["tables"] == [{"markdown": "| h |\n| - |\n| v |"}] + assert counter["converts"] == 1 # was 2 before the fix (content pass + tables pass) + assert counter["instances"] == 1 + + # --------------------------------------------------------------------------- # pages: page-range slice on a form-feed delimited blob # --------------------------------------------------------------------------- diff --git a/tests/storage/db/repositories/test_artifacts.py b/tests/storage/db/repositories/test_artifacts.py index 75ee3905..8c68c37f 100644 --- a/tests/storage/db/repositories/test_artifacts.py +++ b/tests/storage/db/repositories/test_artifacts.py @@ -422,3 +422,37 @@ class TestCascadeDelete: ) assert repo.get_artifact(created["id"]) is None assert repo.list_versions(created["id"]) == [] + + +class TestDeleteArtifact: + def test_delete_removes_versions_and_frees_quota(self, pg_conn): + repo = _repo(pg_conn) + art = repo.create_artifact( + "u-del", "document", conversation_id=_conversation_id(), + storage_path="k/v1.bin", size=100, filename="a.bin", + mime_type="application/octet-stream", + ) + repo.append_version(art["id"], storage_path="k/v2.bin", size=200, filename="a.bin") + assert repo.count_for_user("u-del") == 1 + assert repo.total_bytes_for_user("u-del") == 300 + + paths = repo.delete_artifact(art["id"]) + + assert set(paths) == {"k/v1.bin", "k/v2.bin"} # caller reaps these bytes + assert repo.get_artifact(art["id"]) is None + assert repo.list_versions(art["id"]) == [] # versions cascade + assert repo.count_for_user("u-del") == 0 # quota freed + assert repo.total_bytes_for_user("u-del") == 0 + + def test_delete_for_conversation_scopes_to_that_conversation(self, pg_conn): + repo = _repo(pg_conn) + conv = _conversation_id() + repo.create_artifact("u2", "document", conversation_id=conv, storage_path="k/a", size=10, filename="a") + repo.create_artifact("u2", "document", conversation_id=conv, storage_path="k/b", size=20, filename="b") + other = repo.create_artifact("u2", "document", conversation_id=_conversation_id(), storage_path="k/c", size=30, filename="c") + + n = repo.delete_for_conversation(conv) + + assert n == 2 + assert repo.get_artifact(other["id"]) is not None # other conversation untouched + assert repo.count_for_user("u2") == 1