fix: minor issue fixes

This commit is contained in:
Alex committed 2026-06-30 18:36:06 +02:00
1 parent 99484e02c8
commit 75e07af78f
16 files changed
+950 -43

No files matched your search

@@ -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
+74 -6
View File
@@ -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/<artifact_id>/versions/<int:version>")
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
+23 -1
View File
@@ -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)
+6
View File
@@ -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.
+41 -25
View File
@@ -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)
+54 -7
View File
@@ -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"]]
@@ -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(
@@ -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
# ------------------------------------------------------------------
+54
View File
@@ -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:
+57
View File
@@ -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:
+144 -1
View File
@@ -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
+112 -1
View File
@@ -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)
# ---------------------------------------------------------------------------
+47
View File
@@ -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."""
+60
View File
@@ -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)."""
+88
View File
@@ -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
# ---------------------------------------------------------------------------
@@ -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