From b144265094c1452ce2c68aa2ce24a68ad7593a0e Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 29 Jun 2026 19:21:35 +0200 Subject: [PATCH] fix: workflow security pass --- application/agents/tools/read_document.py | 9 +- application/agents/workflow_agent.py | 85 ++++++++++++++----- .../agents/workflows/workflow_engine.py | 16 +++- application/api/user/artifacts/authz.py | 26 ++++++ application/api/user/artifacts/routes.py | 6 +- .../services/artifact_resource_service.py | 28 ++++-- deployment/docker-compose.yaml | 22 +++++ tests/agents/test_workflow_agent_graph.py | 12 +-- tests/agents/test_workflow_input_documents.py | 71 ++++++++++++++++ tests/agents/tools/test_attachment_bridge.py | 2 +- tests/agents/tools/test_read_document_tool.py | 16 +++- tests/api/user/test_artifacts_routes.py | 38 +++++++++ .../test_artifact_resource_service.py | 17 +++- tests/services/test_mcp_server.py | 36 +++++++- 14 files changed, 340 insertions(+), 44 deletions(-) diff --git a/application/agents/tools/read_document.py b/application/agents/tools/read_document.py index 3bcf6b05..6f031fa0 100644 --- a/application/agents/tools/read_document.py +++ b/application/agents/tools/read_document.py @@ -206,7 +206,14 @@ class ReadDocumentTool(Tool): async_result = parse_document.apply_async( args=[artifact_id, parent, self.user_id, options], queue=queue ) - result = async_result.get(timeout=timeout) + # ``read_document`` runs both from the web request path AND from inside a + # Celery task (headless/scheduled agents). In a prefork worker + # ``task_join_will_block()`` is process-wide, so the default + # ``disable_sync_subtasks=True`` makes ``get()`` raise RuntimeError + # ("Never call result.get() within a task!"). The dedicated ``parsing`` + # queue + separate workers already avoid the real self-deadlock, so opt + # out of the blanket guard explicitly. + result = async_result.get(timeout=timeout, disable_sync_subtasks=False) except (CeleryTimeoutError, TimeoutError): return {"status": "error", "error": f"document parsing timed out after {int(timeout)}s."} except Exception as exc: diff --git a/application/agents/workflow_agent.py b/application/agents/workflow_agent.py index 0d4dafd9..56a1b9c1 100644 --- a/application/agents/workflow_agent.py +++ b/application/agents/workflow_agent.py @@ -60,14 +60,27 @@ class WorkflowAgent(BaseAgent): return self._engine = WorkflowEngine(graph, self) - owner_id = self._resolve_owner_id() - pg_workflow_id = self._precreate_workflow_run(owner_id, query) + # Two distinct identities: the workflow *owner* (A) owns the workflow + # definition and is used to resolve the workflow row; the *runner* (B, + # the caller) owns the run and its artifacts. They are the same user + # except for a shared agent, where B != A. Owning run artifacts by the + # runner means quota is charged to the uploader and the caller can read + # the outputs of the run they triggered (authz is run.user_id == caller). + workflow_owner_id = self._resolve_owner_id() + run_user_id = self._resolve_run_user_id(workflow_owner_id) + pg_workflow_id = self._precreate_workflow_run( + workflow_owner_id, run_user_id, query + ) self._run_persisted = pg_workflow_id is not None - input_documents = self._bridge_attachments(owner_id, persisted=self._run_persisted) + input_documents = self._bridge_attachments( + run_user_id, persisted=self._run_persisted + ) yield from self._engine.execute({"input_documents": input_documents}, query) - self._finalize_workflow_run(owner_id, pg_workflow_id, query) + self._finalize_workflow_run( + workflow_owner_id, run_user_id, pg_workflow_id, query + ) def _load_workflow_graph(self) -> Optional[WorkflowGraph]: if self._workflow_data: @@ -193,12 +206,22 @@ class WorkflowAgent(BaseAgent): return None def _resolve_owner_id(self) -> Optional[str]: - """Resolve the run owner from the explicit workflow owner or the token ``sub``.""" + """Resolve the workflow *owner* (used to resolve the owned workflow row).""" owner_id = self.workflow_owner if not owner_id and isinstance(self.decoded_token, dict): owner_id = self.decoded_token.get("sub") return owner_id + def _resolve_run_user_id(self, workflow_owner_id: Optional[str]) -> Optional[str]: + """Resolve the *runner* (caller) who owns the run and its artifacts. + + Equals the owner for a user running their own workflow (and for external + API-key calls, where the key owner is the caller); for a shared agent it + is the calling user, so their uploads/outputs are owned by and readable + to them rather than silently accruing under the agent owner's account. + """ + return getattr(self, "initial_user_id", None) or getattr(self, "user", None) or workflow_owner_id + def _resolve_owned_workflow_pg_id( self, conn: Any, owner_id: Optional[str] ) -> Optional[str]: @@ -212,18 +235,29 @@ class WorkflowAgent(BaseAgent): workflow_row = wf_repo.get_by_legacy_id(self.workflow_id, owner_id) return str(workflow_row["id"]) if workflow_row is not None else None - def _precreate_workflow_run(self, owner_id: Optional[str], query: str) -> Optional[str]: - """Insert the run row up front so run-scoped artifacts are authz-reachable mid-run.""" - if not self._engine or not self.workflow_id or not owner_id: + def _precreate_workflow_run( + self, + workflow_owner_id: Optional[str], + run_user_id: Optional[str], + query: str, + ) -> Optional[str]: + """Insert the run row up front so run-scoped artifacts are authz-reachable mid-run. + + The workflow row is resolved against its *owner*; the run is owned by the + *runner* so artifact access (``run.user_id == caller``) tracks the caller. + """ + if not self._engine or not self.workflow_id or not workflow_owner_id or not run_user_id: return None try: with db_session() as conn: - pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, owner_id) + pg_workflow_id = self._resolve_owned_workflow_pg_id( + conn, workflow_owner_id + ) if pg_workflow_id is None: return None WorkflowRunsRepository(conn).create( pg_workflow_id, - owner_id, + run_user_id, ExecutionStatus.RUNNING.value, run_id=self._engine.workflow_run_id, inputs={"query": query}, @@ -235,15 +269,16 @@ class WorkflowAgent(BaseAgent): return None def _bridge_attachments( - self, owner_id: Optional[str], *, persisted: bool + self, run_user_id: Optional[str], *, persisted: bool ) -> List[Dict[str, Any]]: """Stage uploaded attachments as run-scoped artifacts the nodes can read. Bytes are read server-side from each attachment's ``upload_path`` and re-persisted through ``persist_new_artifact`` (size/sha256/storage key all derived server-side); only the resulting references enter the run state. + Artifacts are owned by the *runner* (the uploader), not the workflow owner. """ - if not self._engine or not self.attachments or not owner_id: + if not self._engine or not self.attachments or not run_user_id: return [] # Without a persisted run row the artifacts would be orphaned (no authz # parent), so skip the bridge for unowned/draft ids. @@ -279,7 +314,7 @@ class WorkflowAgent(BaseAgent): continue try: ref = persist_new_artifact( - user_id=owner_id, + user_id=run_user_id, kind="file", data=data, filename=filename, @@ -307,15 +342,23 @@ class WorkflowAgent(BaseAgent): return refs def _finalize_workflow_run( - self, owner_id: Optional[str], pg_workflow_id: Optional[str], query: str + self, + workflow_owner_id: Optional[str], + run_user_id: Optional[str], + pg_workflow_id: Optional[str], + query: str, ) -> None: - """Write the run's terminal status/result; upsert the row if pre-creation was skipped.""" + """Write the run's terminal status/result; upsert the row if pre-creation was skipped. + + The run is owned by the *runner* (so it stays readable to the caller and + matches the pre-created row); the workflow row is resolved by its *owner*. + """ if not self._engine: return try: run = WorkflowRun( workflow_id=self.workflow_id or "unknown", - user=owner_id, + user=run_user_id, status=self._determine_run_status(), inputs={"query": query}, outputs=self._serialize_state(self._engine.state), @@ -325,11 +368,13 @@ class WorkflowAgent(BaseAgent): ) steps_json = [step.model_dump(mode="json") for step in run.steps] - if not self.workflow_id or not owner_id: + if not self.workflow_id or not workflow_owner_id or not run_user_id: return with db_session() as conn: if pg_workflow_id is None: - pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, owner_id) + pg_workflow_id = self._resolve_owned_workflow_pg_id( + conn, workflow_owner_id + ) if pg_workflow_id is None: return runs_repo = WorkflowRunsRepository(conn) @@ -337,7 +382,7 @@ class WorkflowAgent(BaseAgent): if self._run_persisted: updated = runs_repo.finalize( self._engine.workflow_run_id, - owner_id, + run_user_id, run.status.value, result=run.outputs, steps=steps_json, @@ -352,7 +397,7 @@ class WorkflowAgent(BaseAgent): if not self._run_persisted or not updated: runs_repo.create( pg_workflow_id, - owner_id, + run_user_id, run.status.value, run_id=self._engine.workflow_run_id, inputs=run.inputs, diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index 746a0b11..57a1d380 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -43,6 +43,11 @@ TEMPLATE_RESERVED_NAMESPACES = {"agent", "artifacts", "system", "source", "tools # only accepts [A-Za-z0-9_-]+, so any disallowed character is stripped before binding. _SESSION_ID_RE = re.compile(r"[^A-Za-z0-9_-]+") +# State keys never staged into a code node's ``state.json``. ``chat_history`` is +# the caller's conversation and must not be reachable by owner-authored, +# egress-open sandbox code (see ``_json_safe_state``). +_CODE_STATE_EXCLUDED_KEYS = frozenset({"chat_history"}) + class WorkflowEngine: MAX_EXECUTION_STEPS = 50 @@ -701,13 +706,20 @@ class WorkflowEngine: return _SESSION_ID_RE.sub("-", str(self.workflow_run_id)) or str(uuid.uuid4()) def _json_safe_state(self) -> Dict[str, Any]: - """Project ``self.state`` to a JSON-safe dict (the code node reads it from state.json).""" + """Project ``self.state`` to a JSON-safe dict (the code node reads it from state.json). + + ``chat_history`` is excluded: it is the *caller's* full conversation, and a + code node (authored by the workflow owner, who may differ from the runner + in a shared agent) has no legitimate need for it. Since sandbox egress is + open by design, staging it would let owner-authored code exfiltrate the + runner's conversation. Node outputs and ``query`` are still exposed. + """ projection: Dict[str, Any] = {} for key, value in self.state.items(): if not isinstance(key, str): continue normalized_key = key.strip() - if not normalized_key: + if not normalized_key or normalized_key in _CODE_STATE_EXCLUDED_KEYS: continue projection[normalized_key] = value return projection diff --git a/application/api/user/artifacts/authz.py b/application/api/user/artifacts/authz.py index f267131e..d9441f4a 100644 --- a/application/api/user/artifacts/authz.py +++ b/application/api/user/artifacts/authz.py @@ -66,3 +66,29 @@ def authorize_artifact(conn, artifact: dict, user_id: Optional[str]) -> bool: return run is not None and run.get("user_id") == user_id # No parent row reachable -> deny (e.g. deleted conversation/run). return False + + +def authorize_artifact_write(conn, artifact: dict, user_id: Optional[str]) -> bool: + """Authorize a *mutating* artifact operation (e.g. restore). + + Stricter than :func:`authorize_artifact`: a write requires an authenticated + owner of the parent. Share links and ``shared_with`` collaborators inherit + read/download access only, so an anonymous (``user_id`` is ``None``) or + share-token-only caller is denied -- they can read an artifact but never + mutate it. + """ + if not user_id: + return False + conversation_id = artifact.get("conversation_id") + workflow_run_id = artifact.get("workflow_run_id") + + if conversation_id is not None: + return ( + ConversationsRepository(conn).get_owned(str(conversation_id), user_id) + is not None + ) + if workflow_run_id is not None: + run = WorkflowRunsRepository(conn).get(str(workflow_run_id)) + return run is not None and run.get("user_id") == user_id + # No parent row reachable -> deny (e.g. deleted conversation/run). + return False diff --git a/application/api/user/artifacts/routes.py b/application/api/user/artifacts/routes.py index fee5c4d1..19824e10 100644 --- a/application/api/user/artifacts/routes.py +++ b/application/api/user/artifacts/routes.py @@ -11,6 +11,7 @@ from flask_restx import Namespace, Resource from application.api import api from application.api.user.artifacts.authz import ( authorize_artifact, + authorize_artifact_write, resolve_authenticated_user, user_can_access_conversation, ) @@ -341,7 +342,10 @@ class RestoreArtifact(Resource): return make_response( jsonify({"success": False, "message": "Artifact not found"}), 404 ) - if not authorize_artifact(conn, artifact, user_id): + # Restore is a WRITE (it appends a new current version); share + # links / shared_with collaborators inherit read access only, so + # gate on the stricter owner-required write check. + if not authorize_artifact_write(conn, artifact, user_id): return make_response( jsonify({"success": False, "message": "Forbidden"}), 403 ) diff --git a/application/services/artifact_resource_service.py b/application/services/artifact_resource_service.py index 92711aae..3ef88398 100644 --- a/application/services/artifact_resource_service.py +++ b/application/services/artifact_resource_service.py @@ -9,8 +9,9 @@ as the HTTP artifact routes). Resources are scoped strictly to that principal: unresolvable principal yields an empty list / a denied read -- never another principal's artifact. -Returns plain ``mcp.types`` objects so both the FastMCP middleware adapter and -the unit tests can consume them without depending on FastMCP internals. +``resources/list`` returns FastMCP ``Resource`` objects so they pass straight +through the server's list/dedupe/wire pipeline; ``resources/read`` is served by +the MCP middleware, which streams the real bytes for each ``artifact://`` uri. """ from __future__ import annotations @@ -21,7 +22,8 @@ import re from dataclasses import dataclass from typing import List, Optional -import mcp.types as mt +from fastmcp.resources.base import Resource +from fastmcp.resources.types import TextResource from sqlalchemy.exc import DataError, DBAPIError from application.core.settings import settings @@ -109,8 +111,17 @@ def _resource_uri(artifact_id: str, version: int) -> str: return f"artifact://{artifact_id}/v{version}" -def list_artifact_resources(api_key: Optional[str]) -> List[mt.Resource]: - """List the calling principal's artifacts as MCP resources (empty if unresolved).""" +def list_artifact_resources(api_key: Optional[str]) -> List[Resource]: + """List the calling principal's artifacts as FastMCP resources (empty if unresolved). + + Returns FastMCP ``Resource`` objects (not raw ``mcp.types.Resource``): the + server's ``resources/list`` pipeline reads FastMCP-only attributes + (``.version``, ``.auth``, ``is_enabled``) and calls ``to_mcp_resource()`` for + the wire encoding, so a raw ``mcp.types.Resource`` mixed into the list would + raise ``AttributeError``. The placeholder ``text`` is never served -- the + middleware's ``on_read_resource`` intercepts every ``artifact://`` read and + streams the real bytes via :func:`read_artifact_resource`. + """ user_id = _resolve_principal(api_key) if not user_id: return [] @@ -121,18 +132,19 @@ def list_artifact_resources(api_key: Optional[str]) -> List[mt.Resource]: logger.exception("artifact resource: list failed") return [] - resources: List[mt.Resource] = [] + resources: List[Resource] = [] for row in rows[:_RESOURCE_LIST_LIMIT]: artifact_id = str(row.get("id")) version = row.get("current_version") or 1 title = row.get("title") or f"artifact-{artifact_id}" resources.append( - mt.Resource( + TextResource( uri=_resource_uri(artifact_id, version), name=title, title=title, description=f"DocsGPT {row.get('kind') or 'file'} artifact", - mimeType=_kind_mime_hint(row.get("kind")), + mime_type=_kind_mime_hint(row.get("kind")), + text="", ) ) return resources diff --git a/deployment/docker-compose.yaml b/deployment/docker-compose.yaml index 7e0ce7fe..5c8cf648 100644 --- a/deployment/docker-compose.yaml +++ b/deployment/docker-compose.yaml @@ -32,6 +32,11 @@ services: - SANDBOX_KERNEL_NAME=docsgpt-python ports: - "7091:7091" + networks: + # `default` reaches the internet + 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 @@ -50,6 +55,9 @@ services: # For heavy/OCR parsing run a separate worker with `-Q parsing` (see # deployment/sandbox/README.md). command: celery -A application.app.celery worker -l INFO -B -Q docsgpt,parsing + networks: + - default + - data-net env_file: - ../.env environment: @@ -106,11 +114,18 @@ services: redis: image: redis:6-alpine + # Data plane only: NOT reachable from the code-exec sandbox (which has no + # auth on this broker, so a reachable Redis = Celery task injection -> 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 @@ -130,6 +145,13 @@ networks: # this network; the runner still reaches the internet via the default bridge). sandbox-net: internal: true + # Internal-only data plane. redis (broker/cache) and postgres live here with + # backend/worker, but the code-exec sandbox is NOT attached, so arbitrary + # sandboxed code cannot reach the unauthenticated broker (Celery task + # injection -> worker RCE) or the database. For full internet-egress + # containment additionally apply deployment/optional/docker-compose.optional.sandbox-egress.yaml. + data-net: + internal: true volumes: postgres_data: diff --git a/tests/agents/test_workflow_agent_graph.py b/tests/agents/test_workflow_agent_graph.py index 853bf29e..9bdc2909 100644 --- a/tests/agents/test_workflow_agent_graph.py +++ b/tests/agents/test_workflow_agent_graph.py @@ -185,7 +185,7 @@ class TestSaveWorkflowRun: def test_returns_when_no_engine(self): agent = _make_agent(workflow_id="x", workflow_owner="u") # _engine is None - agent._finalize_workflow_run(agent.workflow_owner, None, "query") + agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "query") # should not raise def test_returns_when_no_workflow_id(self): @@ -194,7 +194,7 @@ class TestSaveWorkflowRun: agent._engine.execution_log = [] agent._engine.state = {} agent._engine.get_execution_summary.return_value = [] - agent._finalize_workflow_run(agent.workflow_owner, None, "query") + agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "query") def test_returns_when_workflow_missing_in_db(self, pg_conn): agent = _make_agent( @@ -207,7 +207,7 @@ class TestSaveWorkflowRun: agent._engine.get_execution_summary.return_value = [] with _patch_db(pg_conn): # Should just return None since workflow not found in DB - agent._finalize_workflow_run(agent.workflow_owner, None, "q") + agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "q") def test_creates_run_row(self, pg_conn): from application.storage.db.repositories.workflows import ( @@ -234,7 +234,9 @@ class TestSaveWorkflowRun: agent._engine.workflow_run_id = run_id with _patch_db(pg_conn): - agent._finalize_workflow_run(agent.workflow_owner, None, "my query") + agent._finalize_workflow_run( + agent.workflow_owner, agent.workflow_owner, None, "my query" + ) runs = WorkflowRunsRepository(pg_conn).list_for_workflow(str(wf["id"])) assert len(runs) >= 1 @@ -256,7 +258,7 @@ class TestSaveWorkflowRun: "application.agents.workflow_agent.db_session", _broken ): # Should not raise - agent._finalize_workflow_run(agent.workflow_owner, None, "q") + agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "q") class TestDetermineRunStatus: diff --git a/tests/agents/test_workflow_input_documents.py b/tests/agents/test_workflow_input_documents.py index c6c2cac0..b91f5326 100644 --- a/tests/agents/test_workflow_input_documents.py +++ b/tests/agents/test_workflow_input_documents.py @@ -25,6 +25,8 @@ from application.storage.storage_creator import StorageCreator pytestmark = pytest.mark.integration OWNER = "user-bridge" +# A distinct caller for the shared-agent case (caller != workflow owner). +RUNNER = "user-runner" def _wire(pg_engine, tmp_path, monkeypatch) -> LocalStorage: @@ -181,6 +183,75 @@ def test_attachments_bridge_to_run_scoped_artifacts(pg_engine, tmp_path, monkeyp assert fh.read() == a1 +def test_shared_agent_run_and_artifacts_owned_by_caller(pg_engine, tmp_path, monkeypatch): + """Shared agent (caller != owner): the run + bridged artifacts are owned by the caller. + + The workflow row is still resolved against its owner, but the run.user_id and + the artifact owner are the caller, so quota is charged to the uploader and the + caller (not the agent owner) can read the run's artifacts. + """ + storage = _wire(pg_engine, tmp_path, monkeypatch) + wf_id = _make_workflow(pg_engine, owner=OWNER) + + attachments = [_stage_attachment(storage, b"caller-doc", "c.txt", "text/plain")] + agent = _agent(wf_id, attachments, owner=OWNER) + # Simulate a shared-agent invocation: caller identity differs from the owner. + agent.initial_user_id = RUNNER + agent.user = RUNNER + + _patch_engine(monkeypatch) + list(agent._gen_inner("summarize", log_context=None)) + engine = _RecordingEngine.instances[-1] + run_id = engine.workflow_run_id + + with pg_engine.connect() as conn: + run = WorkflowRunsRepository(conn).get(run_id) + assert run is not None + # The run is owned by the caller, not the workflow owner. + assert run["user_id"] == RUNNER + + refs = engine.captured_inputs["input_documents"] + assert len(refs) == 1 + owner_row = conn.execute( + text("SELECT user_id FROM artifacts WHERE workflow_run_id = CAST(:r AS uuid)"), + {"r": run_id}, + ).fetchone() + assert owner_row[0] == RUNNER + + # authz: the caller can reach the run's artifacts; the agent owner cannot. + # ``authorize_artifact`` uses the passed conn but reads ``request.args`` for a + # share token, so it needs a request context. + from flask import Flask + + from application.api.user.artifacts.authz import authorize_artifact + + app = Flask(__name__) + with app.test_request_context(): + with pg_engine.connect() as conn: + artifact = ArtifactsRepository(conn).get_artifact(refs[0]["artifact_id"]) + assert authorize_artifact(conn, artifact, RUNNER) is True + assert authorize_artifact(conn, artifact, OWNER) is False + + +def test_code_state_excludes_chat_history(pg_engine, tmp_path, monkeypatch): + """A code node's state.json projection omits the caller's chat_history.""" + _wire(pg_engine, tmp_path, monkeypatch) + wf_id = _make_workflow(pg_engine) + agent = _agent(wf_id, [], owner=OWNER) + agent.chat_history = [{"prompt": "secret question", "response": "secret answer"}] + + _patch_engine(monkeypatch) + list(agent._gen_inner("do it", log_context=None)) + engine = _RecordingEngine.instances[-1] + + projected = engine._json_safe_state() + # chat_history is set in state but must never be staged for sandboxed code. + assert "chat_history" in engine.state + assert "chat_history" not in projected + # Legitimate state (the query, node inputs) is still exposed. + assert projected.get("query") == "do it" + + def test_attachments_capped_per_run(pg_engine, tmp_path, monkeypatch): """More than the cap of attachments bridges only the cap; the rest are dropped.""" storage = _wire(pg_engine, tmp_path, monkeypatch) diff --git a/tests/agents/tools/test_attachment_bridge.py b/tests/agents/tools/test_attachment_bridge.py index 3e38b864..4d7f595b 100644 --- a/tests/agents/tools/test_attachment_bridge.py +++ b/tests/agents/tools/test_attachment_bridge.py @@ -417,7 +417,7 @@ def test_read_document_bridges_attachment_then_enqueues(monkeypatch): captured: Dict[str, Any] = {} class _AR: - def get(self, timeout=None): + def get(self, timeout=None, disable_sync_subtasks=True): return {"status": "ok", "content": "parsed", "truncated": False} import application.api.user.tasks as tasks diff --git a/tests/agents/tools/test_read_document_tool.py b/tests/agents/tools/test_read_document_tool.py index c6d46a45..f4160b58 100644 --- a/tests/agents/tools/test_read_document_tool.py +++ b/tests/agents/tools/test_read_document_tool.py @@ -60,8 +60,15 @@ class _FakeAsyncResult: def __init__(self, payload=None, exc=None): self._payload = payload self._exc = exc + self.get_kwargs = None - def get(self, timeout=None): + def get(self, timeout=None, disable_sync_subtasks=True): + # ``read_document`` must pass ``disable_sync_subtasks=False`` so the await + # works from inside a Celery worker (headless/scheduled agents). + self.get_kwargs = { + "timeout": timeout, + "disable_sync_subtasks": disable_sync_subtasks, + } if self._exc is not None: raise self._exc return self._payload @@ -76,7 +83,9 @@ def _patch_task(monkeypatch, *, payload=None, exc=None): def _apply_async(args=None, queue=None, **kw): captured["args"] = args captured["queue"] = queue - return _FakeAsyncResult(payload=payload, exc=exc) + result = _FakeAsyncResult(payload=payload, exc=exc) + captured["result"] = result + return result monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async) return captured @@ -170,6 +179,9 @@ def test_resolves_input_then_enqueues_and_returns_payload(monkeypatch): assert captured["queue"] == rd.settings.DOCUMENT_PARSE_QUEUE options = captured["args"][3] assert options["output"] == "markdown" and options["persist"] is False + # The await must opt out of Celery's worker-self-deadlock guard, else it + # raises RuntimeError from inside a prefork worker (headless/scheduled agents). + assert captured["result"].get_kwargs["disable_sync_subtasks"] is False @pytest.mark.unit diff --git a/tests/api/user/test_artifacts_routes.py b/tests/api/user/test_artifacts_routes.py index f65725e4..0322830a 100644 --- a/tests/api/user/test_artifacts_routes.py +++ b/tests/api/user/test_artifacts_routes.py @@ -467,6 +467,44 @@ class TestRestoreArtifact: ) assert resp.status_code == 403 + def test_restore_anonymous_share_token_holder_denied( + self, _patch_db, flask_app + ): + # A share link inherits read/download access only; restore is a WRITE and + # an anonymous link holder must NOT be able to mutate the artifact. + from application.api.user.artifacts.routes import RestoreArtifact + + conv = _make_conversation(_patch_db) + art = _make_artifact( + _patch_db, conversation_id=str(conv["id"]), spec={"body": "v1"} + ) + ArtifactsRepository(_patch_db).append_version(art["id"], spec={"body": "v2"}) + share = SharedConversationsRepository(_patch_db).create(str(conv["id"]), OWNER) + + resp = _call( + flask_app, RestoreArtifact, art["id"], token=None, + query={"share_token": str(share["uuid"])}, + json_body={"version": 1}, method="post", + ) + assert resp.status_code == 403 + # The artifact is unchanged (still at version 2, not reverted/appended). + assert ArtifactsRepository(_patch_db).get_artifact(art["id"])["current_version"] == 2 + + def test_restore_shared_with_collaborator_denied(self, _patch_db, flask_app): + # A read-only ``shared_with`` collaborator can GET but not restore. + from application.api.user.artifacts.routes import RestoreArtifact + + conv = _make_conversation(_patch_db) + ConversationsRepository(_patch_db).add_shared_user(str(conv["id"]), SHARED_USER) + art = _make_artifact(_patch_db, conversation_id=str(conv["id"])) + ArtifactsRepository(_patch_db).append_version(art["id"], spec={"body": "v2"}) + + resp = _call( + flask_app, RestoreArtifact, art["id"], token={"sub": SHARED_USER}, + json_body={"version": 1}, method="post", + ) + assert resp.status_code == 403 + # --------------------------------------------------------------------------- # Malformed artifact_id path segment -> 404 (no CAST(:id AS uuid) txn poison) diff --git a/tests/services/test_artifact_resource_service.py b/tests/services/test_artifact_resource_service.py index c9093f4a..a4e9ed06 100644 --- a/tests/services/test_artifact_resource_service.py +++ b/tests/services/test_artifact_resource_service.py @@ -120,8 +120,21 @@ class TestListArtifactResources: out = {str(r.uri): r for r in svc.list_artifact_resources("owner-key")} assert out[f"artifact://{ART_TEXT}/v2"].name == "notes" # The list row never advertises a wildcard/wrong type; the image kind - # falls back to the generic octet-stream hint. - assert out[f"artifact://{ART_BIN}/v1"].mimeType == "application/octet-stream" + # falls back to the generic octet-stream hint. FastMCP Resource uses + # ``mime_type`` (and must survive ``to_mcp_resource()`` for the wire). + assert out[f"artifact://{ART_BIN}/v1"].mime_type == "application/octet-stream" + + def test_resources_are_fastmcp_and_render_to_wire(self): + # Regression: ``resources/list`` must yield FastMCP Resource objects so + # the server's list pipeline (reads ``.version``/``.auth``, calls + # ``to_mcp_resource``) does not raise AttributeError on a raw mcp.types. + from fastmcp.resources.base import Resource as FastMCPResource + + out = svc.list_artifact_resources("owner-key") + assert out and all(isinstance(r, FastMCPResource) for r in out) + for r in out: + wire = r.to_mcp_resource() + assert str(wire.uri).startswith("artifact://") @pytest.mark.unit diff --git a/tests/services/test_mcp_server.py b/tests/services/test_mcp_server.py index 9a12bc5f..ede019aa 100644 --- a/tests/services/test_mcp_server.py +++ b/tests/services/test_mcp_server.py @@ -144,11 +144,14 @@ async def _call_next_empty(context): class TestArtifactResourcesMiddleware: @pytest.mark.asyncio async def test_list_appends_principal_artifacts(self): - import mcp.types as mt + from fastmcp.resources.types import TextResource + from application.mcp_server import ArtifactResourcesMiddleware mw = ArtifactResourcesMiddleware() - res = mt.Resource(uri="artifact://a/v1", name="x", mimeType="text/plain") + res = TextResource( + uri="artifact://a/v1", name="x", mime_type="text/plain", text="" + ) with ( patch("application.mcp_server._extract_bearer_token", return_value="k"), patch( @@ -158,6 +161,35 @@ class TestArtifactResourcesMiddleware: out = await mw.on_list_resources(SimpleNamespace(), _call_next_empty) assert [str(r.uri) for r in out] == ["artifact://a/v1"] + @pytest.mark.asyncio + async def test_list_resources_wire_pipeline_accepts_artifact(self): + # Regression for the resources/list crash: drive the real server pipeline + # (is_enabled / auth filter / dedupe / to_mcp_resource), which raises + # AttributeError on a raw mcp.types.Resource. A FastMCP resource passes. + from fastmcp.resources.types import TextResource + + from application.mcp_server import mcp + + res = TextResource( + uri="artifact://a/v1", + name="x", + title="x", + description="d", + mime_type="text/plain", + text="", + ) + with ( + patch("application.mcp_server._extract_bearer_token", return_value="k"), + patch( + "application.mcp_server.list_artifact_resources", return_value=[res] + ), + ): + resources = await mcp.list_resources() + assert "artifact://a/v1" in [str(r.uri) for r in resources] + # Every listed resource must encode to the MCP wire shape. + for r in resources: + assert str(r.to_mcp_resource().uri) + @pytest.mark.asyncio async def test_read_text_wraps_as_text_contents(self): from application.mcp_server import ArtifactResourcesMiddleware