diff --git a/application/agents/workflow_agent.py b/application/agents/workflow_agent.py index 56a1b9c1..5a644097 100644 --- a/application/agents/workflow_agent.py +++ b/application/agents/workflow_agent.py @@ -1,6 +1,6 @@ import logging from datetime import datetime, timezone -from typing import Any, Dict, Generator, List, Optional +from typing import Any, Dict, Generator, List, Optional, Tuple from application.agents.base import BaseAgent from application.agents.workflows.schemas import ( @@ -12,7 +12,9 @@ from application.agents.workflows.schemas import ( WorkflowRun, ) from application.agents.workflows.workflow_engine import WorkflowEngine +from application.core.settings import settings from application.logging import log_activity, LogContext +from application.sandbox.artifacts_capture import QuotaExceeded from application.storage.db.base_repository import looks_like_uuid from application.storage.db.repositories.workflow_edges import WorkflowEdgesRepository from application.storage.db.repositories.workflow_nodes import WorkflowNodesRepository @@ -44,6 +46,9 @@ class WorkflowAgent(BaseAgent): self._workflow_data = workflow self._engine: Optional[WorkflowEngine] = None self._run_persisted = False + # Set to a message when the input-document bridge fails fatally (quota), so + # the run is finalized FAILED instead of running with missing documents. + self._bridge_error: Optional[str] = None @log_activity() def gen( @@ -73,10 +78,33 @@ class WorkflowAgent(BaseAgent): ) self._run_persisted = pg_workflow_id is not None - input_documents = self._bridge_attachments( - run_user_id, persisted=self._run_persisted - ) + try: + input_documents, dropped = self._bridge_attachments( + run_user_id, persisted=self._run_persisted + ) + except QuotaExceeded as exc: + # The run's input documents exceed the uploader's artifact quota. Surface + # a clean error and finalize the pre-created RUNNING row as FAILED rather + # than executing nodes with silently-missing documents. + self._bridge_error = str(exc) + yield { + "type": "error", + "error": ( + "This run's input documents exceed your artifact storage quota. " + "Delete some artifacts and try again." + ), + } + self._finalize_workflow_run( + workflow_owner_id, run_user_id, pg_workflow_id, query + ) + return + # Non-fatal: some attachments were dropped (oversize / unreadable). Tell the + # user which, then still run with the documents that did bridge. + if dropped: + yield {"type": "error", "error": " ".join(dropped)} + + self._engine.run_persisted = self._run_persisted yield from self._engine.execute({"input_documents": input_documents}, query) self._finalize_workflow_run( workflow_owner_id, run_user_id, pg_workflow_id, query @@ -270,66 +298,77 @@ class WorkflowAgent(BaseAgent): def _bridge_attachments( self, run_user_id: Optional[str], *, persisted: bool - ) -> List[Dict[str, Any]]: + ) -> Tuple[List[Dict[str, Any]], List[str]]: """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. + Bytes are read server-side from each attachment's ``upload_path`` (bounded + by ``ARTIFACT_MAX_BYTES``, handle always closed) 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. + + Returns the bridged references and a list of user-facing notices for + attachments that were dropped (oversize / unreadable / unstorable). + ``QuotaExceeded`` is NOT swallowed: it propagates to the caller so the run + fails cleanly instead of running with silently-missing documents. """ if not self._engine or not self.attachments or not run_user_id: - return [] + return [], [] # Without a persisted run row the artifacts would be orphaned (no authz # parent), so skip the bridge for unowned/draft ids. if not persisted: - return [] + return [], [] from application.sandbox.artifacts_capture import persist_new_artifact from application.storage.storage_creator import StorageCreator storage = StorageCreator.get_storage() + max_bytes = int(getattr(settings, "ARTIFACT_MAX_BYTES", 0) or 0) + dropped: List[str] = [] if len(self.attachments) > _MAX_INPUT_DOCUMENTS: - dropped = len(self.attachments) - _MAX_INPUT_DOCUMENTS + over = len(self.attachments) - _MAX_INPUT_DOCUMENTS logger.warning( "Workflow run input documents exceed cap (%d); dropping %d attachment(s)", _MAX_INPUT_DOCUMENTS, - dropped, + over, + ) + dropped.append( + f"Only the first {_MAX_INPUT_DOCUMENTS} input document(s) were used; " + f"{over} additional attachment(s) were dropped." ) refs: List[Dict[str, Any]] = [] - for index, attachment in enumerate(self.attachments[:_MAX_INPUT_DOCUMENTS]): + for attachment in self.attachments[:_MAX_INPUT_DOCUMENTS]: upload_path = attachment.get("upload_path") or attachment.get("path") if not upload_path: continue filename = attachment.get("filename") or "attachment" mime_type = attachment.get("mime_type") or "application/octet-stream" - attachment_id = attachment.get("id", index) - try: - data = storage.get_file(upload_path).read() - except Exception as exc: - logger.error( - "Failed to read attachment %s for workflow run: %s", - attachment_id, - type(exc).__name__, - ) + # Reject oversize attachments via the authoritative ``size`` column + # BEFORE buffering the bytes into worker memory (a memory-DoS guard); + # the bounded read below backstops a missing/lying ``size``. + declared_size = attachment.get("size") + if max_bytes and isinstance(declared_size, (int, float)) and declared_size > max_bytes: + dropped.append(f'Document "{filename}" exceeds the artifact size limit and was skipped.') continue - try: - ref = persist_new_artifact( - user_id=run_user_id, - kind="file", - data=data, - filename=filename, - mime_type=mime_type, - title=filename, - workflow_run_id=self._engine.workflow_run_id, - ) - except Exception as exc: - logger.error( - "Failed to persist attachment %s artifact: %s", - attachment_id, - type(exc).__name__, - ) + data = self._read_attachment_bytes(storage, upload_path, max_bytes) + if data is None: + dropped.append(f'Document "{filename}" could not be read and was skipped.') continue + if max_bytes and len(data) > max_bytes: + dropped.append(f'Document "{filename}" exceeds the artifact size limit and was skipped.') + continue + # QuotaExceeded propagates (fatal); persist_new_artifact returns None on + # any other error, which we report as a per-attachment drop. + ref = persist_new_artifact( + user_id=run_user_id, + kind="file", + data=data, + filename=filename, + mime_type=mime_type, + title=filename, + workflow_run_id=self._engine.workflow_run_id, + ) if ref is None: + dropped.append(f'Document "{filename}" could not be stored and was skipped.') continue refs.append( { @@ -339,7 +378,25 @@ class WorkflowAgent(BaseAgent): "mime_type": ref["mime_type"], } ) - return refs + return refs, dropped + + @staticmethod + def _read_attachment_bytes(storage: Any, upload_path: str, max_bytes: int) -> Optional[bytes]: + """Read an attachment with a bounded read and a guaranteed handle close; None on failure.""" + try: + file_obj = storage.get_file(upload_path) + except Exception as exc: + logger.error("Failed to open attachment for workflow run: %s", type(exc).__name__) + return None + try: + return file_obj.read(max_bytes + 1) if max_bytes else file_obj.read() + except Exception as exc: + logger.error("Failed to read attachment for workflow run: %s", type(exc).__name__) + return None + finally: + close = getattr(file_obj, "close", None) + if callable(close): + close() def _finalize_workflow_run( self, @@ -410,6 +467,10 @@ class WorkflowAgent(BaseAgent): logger.error(f"Failed to save workflow run: {e}") def _determine_run_status(self) -> ExecutionStatus: + # A fatal input-document bridge failure (quota) means the engine never ran; + # the run is FAILED even though there is no per-node failure log entry. + if self._bridge_error is not None: + return ExecutionStatus.FAILED if not self._engine or not self._engine.execution_log: return ExecutionStatus.COMPLETED for log in self._engine.execution_log: diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index 03aaf1a4..3991fedb 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -69,6 +69,10 @@ class WorkflowEngine: # artifact's parent; mint one when the caller has not supplied a # pre-created ``workflow_runs`` id so the engine is self-contained. self.workflow_run_id: str = workflow_run_id or str(uuid.uuid4()) + # False when no ``workflow_runs`` row backs ``workflow_run_id`` (unsaved / + # embedded draft), so run-scoped artifacts must NOT be persisted as orphans. + # Defaults True; WorkflowAgent flips it off on the draft path. + self.run_persisted: bool = True self.state: WorkflowState = {} self.execution_log: List[Dict[str, Any]] = [] self._condition_result: Optional[str] = None @@ -299,8 +303,11 @@ class WorkflowEngine: node_agent = WorkflowNodeAgentFactory.create(**factory_kwargs) # Run-scope the node agent's tools so artifact_generator / code_executor # address artifacts by this workflow run: a short ref (A1) created by one - # node resolves for edit_artifact in a later node within the same run. - if getattr(node_agent, "tool_executor", None) is not None: + # node resolves for edit_artifact in a later node within the same run. Only + # when a workflow_runs row backs the run -- otherwise an artifact parented + # to this run id would be an orphan (403 on get/download); left unset, the + # run-scoped tools persist under a conversation parent or cleanly error. + if self.run_persisted and getattr(node_agent, "tool_executor", None) is not None: node_agent.tool_executor.workflow_run_id = self.workflow_run_id # Decide native-eligibility from the SAME supported-types list the provider @@ -396,14 +403,21 @@ class WorkflowEngine: manager.put_file(session_id, "state.json", state_json) pre_signatures = snapshot_signatures(manager, session_id) result = manager.exec(session_id, code, timeout=timeout) - artifacts = capture_artifacts( - manager, - session_id, - pre_signatures, - user_id=user_id, - workflow_run_id=self.workflow_run_id, - produced_by={"node_id": node.id, "node_type": NodeType.CODE.value}, - ) + if self.run_persisted: + artifacts = capture_artifacts( + manager, + session_id, + pre_signatures, + user_id=user_id, + workflow_run_id=self.workflow_run_id, + produced_by={"node_id": node.id, "node_type": NodeType.CODE.value}, + ) + else: + # No workflow_runs row backs this run (unsaved/embedded draft): an + # artifact parented to this run id would be an unreachable orphan + # (403 on get/download). Skip persistence; the sandbox still ran and + # ``_build_code_output`` handles the empty-artifacts case. + artifacts = [] finally: try: manager.close(session_id) diff --git a/application/api/user/artifacts/authz.py b/application/api/user/artifacts/authz.py index 09d34d80..a09f6bb5 100644 --- a/application/api/user/artifacts/authz.py +++ b/application/api/user/artifacts/authz.py @@ -24,10 +24,11 @@ class Principal: A decoded JWT is a full *owner* session (``agent_id`` is ``None``). An ``api_key`` is a low-trust, publicly-distributed *agent* credential -- it is embedded client-side in the widget and accepted from the query string -- so - it resolves to the owning ``user_id`` but stays **agent-scoped**: it may only - reach artifacts from that agent's own conversations, never the owner's whole - corpus, and may not perform mutating operations. This mirrors the scoping the - MCP resource layer already enforces (see ``artifact_resource_service``). + it resolves to the owning ``user_id`` but stays **agent-scoped**: it may reach + an artifact only when the request also carries that artifact's parent + ``conversation_id`` (the unguessable per-visitor secret), confining it to a + single conversation rather than the agent's whole corpus, and may not perform + mutating operations. """ user_id: Optional[str] = None @@ -81,17 +82,23 @@ def user_can_access_conversation( def authorize_artifact(conn, artifact: dict, principal: Principal) -> bool: """Authorize a READ of ``artifact`` for ``principal``; missing parent fails closed. - A low-trust agent api_key is confined to its own agent's conversations (owner - match + agent scope) and never inherits share-link access. A JWT owner, - ``shared_with`` collaborator, or share-token holder is authorized by resolving - the artifact's parent (conversation or workflow run). + A low-trust agent api_key is confined to a single conversation it proves by + carrying that artifact's parent ``conversation_id`` query param (owner match + + matching conversation + agent scope) and never inherits share-link access. A + JWT owner, ``shared_with`` collaborator, or share-token holder is authorized by + resolving the artifact's parent (conversation or workflow run). """ if principal.is_agent_scoped: - # An agent key is not the owner's session: require both that the artifact - # belongs to the owner AND that its parent conversation is this agent's, - # so a public widget key cannot enumerate the owner's whole corpus. + # An agent key is not the owner's session and is embedded in public widget + # JS, so require the artifact to belong to the owner AND the request to + # carry the parent conversation_id (the unguessable per-visitor secret): + # the key alone, without that id, cannot download a known artifact. if str(artifact.get("user_id")) != str(principal.user_id): return False + req_conv = request.args.get("conversation_id") + parent = artifact.get("conversation_id") + if not req_conv or parent is None or str(parent) != str(req_conv): + return False return ArtifactsRepository(conn).artifact_in_agent_scope( str(artifact.get("id")), str(principal.agent_id) ) diff --git a/application/api/user/artifacts/routes.py b/application/api/user/artifacts/routes.py index b1a7b217..9f9ea3fc 100644 --- a/application/api/user/artifacts/routes.py +++ b/application/api/user/artifacts/routes.py @@ -124,21 +124,23 @@ class ListArtifacts(Resource): with db_readonly() as conn: repo = ArtifactsRepository(conn) if principal.is_agent_scoped: - # A low-trust agent api_key only ever sees its OWN agent's - # artifacts (never the owner's whole corpus): scope to the - # agent, then optionally narrow to the requested conversation. - # Workflow runs are not agent-scoped, so reject that filter. + # A low-trust agent api_key is embedded in public widget JS and + # shared by every anonymous visitor, so it may NOT enumerate the + # agent's artifacts. The unguessable conversation_id (UUIDv4) is + # the only per-visitor secret: require it as a bearer capability + # and scope the query to it. Workflow runs are not agent-scoped, + # so reject that filter. if workflow_run_id: return make_response( jsonify({"success": False, "message": "Forbidden"}), 403 ) - rows = repo.list_artifacts_for_agent(principal.agent_id, user_id) - if conversation_id: - rows = [ - r - for r in rows - if str(r.get("conversation_id")) == str(conversation_id) - ] + if not conversation_id: + return make_response( + jsonify({"success": False, "message": "Forbidden"}), 403 + ) + rows = repo.list_artifacts_for_agent( + principal.agent_id, user_id, conversation_id=conversation_id + ) elif conversation_id: if not user_can_access_conversation( conn, conversation_id, user_id, share_token diff --git a/application/services/artifact_resource_service.py b/application/services/artifact_resource_service.py index b9d86968..d1356c07 100644 --- a/application/services/artifact_resource_service.py +++ b/application/services/artifact_resource_service.py @@ -133,6 +133,11 @@ def list_artifact_resources(api_key: Optional[str]) -> List[Resource]: return [] try: with db_readonly() as conn: + # Known residual enumeration gap (tracked separately): unlike the HTTP + # artifact list route, MCP-over-key has no per-visitor conversation + # context, so this lists every artifact across the agent's + # conversations. Whether MCP-over-key should enumerate at all is a + # separate product decision; do not "fix" it by neutering the feature. rows = ArtifactsRepository(conn).list_artifacts_for_agent( str(agent["id"]), str(agent["user_id"]) ) diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py index b75af252..c3a515be 100644 --- a/application/storage/db/repositories/artifacts.py +++ b/application/storage/db/repositories/artifacts.py @@ -184,20 +184,29 @@ class ArtifactsRepository: ) return [_artifact_to_dict(r) for r in result.fetchall()] - def list_artifacts_for_agent(self, agent_id: str, user_id: str) -> list[dict]: - """List artifacts whose parent conversation belongs to ``agent_id`` (owner-scoped). + def list_artifacts_for_agent( + self, agent_id: str, user_id: str, conversation_id: Optional[str] = None + ) -> list[dict]: + """List an agent's artifacts (owner-scoped), optionally narrowed to one conversation. Scopes a per-agent api-key's artifact visibility to the conversations that - agent produced, so the key cannot enumerate the owner's whole corpus. + agent produced, so the key cannot enumerate the owner's whole corpus. When + ``conversation_id`` is given, the SQL narrows to that single conversation + (the per-visitor bearer capability for a public widget key). """ + clauses = ["c.agent_id = CAST(:agent_id AS uuid)", "a.user_id = :user_id"] + params: dict[str, Any] = {"agent_id": str(agent_id), "user_id": user_id} + if conversation_id is not None: + clauses.append("a.conversation_id = CAST(:conversation_id AS uuid)") + params["conversation_id"] = conversation_id result = self._conn.execute( text( "SELECT a.* FROM artifacts a " "JOIN conversations c ON a.conversation_id = c.id " - "WHERE c.agent_id = CAST(:agent_id AS uuid) AND a.user_id = :user_id " + f"WHERE {' AND '.join(clauses)} " "ORDER BY a.created_at DESC, a.id DESC" ), - {"agent_id": str(agent_id), "user_id": user_id}, + params, ) return [_artifact_to_dict(r) for r in result.fetchall()] diff --git a/tests/agents/test_workflow_agent_graph.py b/tests/agents/test_workflow_agent_graph.py index 9bdc2909..f5f17c3a 100644 --- a/tests/agents/test_workflow_agent_graph.py +++ b/tests/agents/test_workflow_agent_graph.py @@ -21,6 +21,7 @@ def _make_agent(*, workflow_id=None, workflow=None, workflow_owner=None, agent._workflow_data = workflow agent._engine = None agent._run_persisted = False + agent._bridge_error = None agent.decoded_token = decoded_token or {} return agent diff --git a/tests/agents/test_workflow_code_node.py b/tests/agents/test_workflow_code_node.py index 77a079c2..00840a8f 100644 --- a/tests/agents/test_workflow_code_node.py +++ b/tests/agents/test_workflow_code_node.py @@ -90,7 +90,12 @@ class _FakeManager: @pytest.fixture() def patch_sandbox(monkeypatch): """Patch the sandbox manager + capture helper; return knobs to drive them.""" - state = {"result": _Result(ok=True, stdout="ok"), "captured": [], "snapshot_calls": 0} + state = { + "result": _Result(ok=True, stdout="ok"), + "captured": [], + "snapshot_calls": 0, + "capture_calls": 0, + } manager_holder = {} @@ -99,6 +104,10 @@ def patch_sandbox(monkeypatch): manager_holder["manager"] = manager return manager + def _capture(*a, **k): + state["capture_calls"] += 1 + return list(state["captured"]) + monkeypatch.setattr( "application.sandbox.sandbox_creator.SandboxCreator.get_manager", _get_manager ) @@ -107,8 +116,7 @@ def patch_sandbox(monkeypatch): lambda *a, **k: state.__setitem__("snapshot_calls", state["snapshot_calls"] + 1) or {}, ) monkeypatch.setattr( - "application.sandbox.artifacts_capture.capture_artifacts", - lambda *a, **k: list(state["captured"]), + "application.sandbox.artifacts_capture.capture_artifacts", _capture ) state["manager_holder"] = manager_holder return state @@ -148,6 +156,40 @@ def test_code_node_no_artifacts_still_writes_status(patch_sandbox): assert engine.state["out"] == {"artifacts": [], "status": "ok"} +def test_code_node_captures_when_run_persisted(patch_sandbox): + # The default persisted run captures produced files into artifacts. + engine = _engine() # run_persisted defaults True + patch_sandbox["captured"] = [] + node = _code_node(output_variable="out", code="print('x')") + + list(engine._execute_code_node(node)) + + assert patch_sandbox["capture_calls"] == 1 + + +def test_code_node_skips_capture_when_run_not_persisted(patch_sandbox): + # An unsaved/draft run has no workflow_runs row, so capturing would create + # orphan artifacts (403 on get/download): capture is skipped entirely and the + # node emits the empty-artifacts shape. The sandbox still ran. + engine = _engine() + engine.run_persisted = False + ref = { + "artifact_id": "art-1", "version": 1, "filename": "r.pdf", + "mime_type": "application/pdf", "size": 3, + } + patch_sandbox["captured"] = [ref] + node = _code_node(output_variable="out", code="open('r.pdf','wb').write(b'x')") + + list(engine._execute_code_node(node)) + + assert patch_sandbox["capture_calls"] == 0 + assert engine.state["out"] == {"artifacts": [], "status": "ok"} + # The sandbox session still opened and closed (only persistence was skipped). + manager = patch_sandbox["manager_holder"]["manager"] + assert manager.opened == ["11111111-1111-1111-1111-111111111111"] + assert manager.closed == ["11111111-1111-1111-1111-111111111111"] + + def test_code_node_reads_prior_state_from_state_json(patch_sandbox): # Prior state is staged as DATA in state.json (workspace root = kernel cwd) so # node code reads it with json.load(open("state.json")) -- e.g. state["decision"]. diff --git a/tests/agents/test_workflow_input_documents.py b/tests/agents/test_workflow_input_documents.py index ebb5ea72..4414708e 100644 --- a/tests/agents/test_workflow_input_documents.py +++ b/tests/agents/test_workflow_input_documents.py @@ -82,6 +82,7 @@ def _agent(workflow_id, attachments, owner: str = OWNER) -> WorkflowAgent: agent._workflow_data = None agent._engine = None agent._run_persisted = False + agent._bridge_error = None return agent @@ -356,6 +357,99 @@ def test_no_attachments_run_still_works(pg_engine, tmp_path, monkeypatch): assert n == 0 +def test_quota_exceeded_fails_run_and_does_not_execute(pg_engine, tmp_path, monkeypatch): + """Over quota: the bridge raises, an error is surfaced, the run is finalized FAILED, + and the engine never executes with silently-missing documents.""" + storage = _wire(pg_engine, tmp_path, monkeypatch) + wf_id = _make_workflow(pg_engine) + attachments = [_stage_attachment(storage, b"doc", "d.txt", "text/plain")] + agent = _agent(wf_id, attachments) + _patch_engine(monkeypatch) + + from application.sandbox.artifacts_capture import QuotaExceeded + + def _raise_quota(**kwargs): + raise QuotaExceeded("artifact storage quota reached") + + monkeypatch.setattr( + "application.sandbox.artifacts_capture.persist_new_artifact", _raise_quota + ) + + events = list(agent._gen_inner("summarize", log_context=None)) + engine = _RecordingEngine.instances[-1] + + # A fatal error was surfaced on the stream. + errors = [e for e in events if e.get("type") == "error"] + assert errors and "quota" in errors[0]["error"].lower() + # The engine never executed (execute was not reached -> no captured inputs). + assert engine.captured_inputs is None + # The pre-created RUNNING row was finalized FAILED (not left dangling), and no + # artifact rows were persisted for the run. + with pg_engine.connect() as conn: + run = WorkflowRunsRepository(conn).get(engine.workflow_run_id) + n = conn.execute( + text("SELECT count(*) FROM artifacts WHERE workflow_run_id = CAST(:r AS uuid)"), + {"r": engine.workflow_run_id}, + ).scalar() + assert run is not None + assert run["status"] == "failed" + assert n == 0 + + +def test_oversize_declared_attachment_skipped_with_notice(pg_engine, tmp_path, monkeypatch): + """A declared-oversize attachment is dropped with a surfaced notice; the run still proceeds.""" + from application.core.settings import settings + + storage = _wire(pg_engine, tmp_path, monkeypatch) + wf_id = _make_workflow(pg_engine) + att = _stage_attachment(storage, b"tiny", "big.txt", "text/plain") + att["size"] = int(settings.ARTIFACT_MAX_BYTES) + 1 # declared past the per-file cap + agent = _agent(wf_id, [att]) + _patch_engine(monkeypatch) + + events = list(agent._gen_inner("summarize", log_context=None)) + engine = _RecordingEngine.instances[-1] + + # The oversize doc was dropped -> no input documents bridged, but the run ran. + assert engine.captured_inputs is not None + assert engine.captured_inputs["input_documents"] == [] + # A non-fatal notice naming the dropped document was surfaced. + errors = [e for e in events if e.get("type") == "error"] + assert errors and "big.txt" in errors[0]["error"] + # Nothing was persisted for the oversize doc. + with pg_engine.connect() as conn: + n = conn.execute( + text("SELECT count(*) FROM artifacts WHERE workflow_run_id = CAST(:r AS uuid)"), + {"r": engine.workflow_run_id}, + ).scalar() + assert n == 0 + + +def test_read_attachment_bytes_closes_handle_on_bounded_read(): + """The bounded read pulls at most max_bytes+1 and always closes the storage handle.""" + + class _Handle: + def __init__(self, data: bytes) -> None: + self._data = data + self.closed = False + + def read(self, n: int = -1) -> bytes: + return self._data if n is None or n < 0 else self._data[:n] + + def close(self) -> None: + self.closed = True + + handle = _Handle(b"x" * 100) + + class _Storage: + def get_file(self, _path): + return handle + + data = WorkflowAgent._read_attachment_bytes(_Storage(), "p", max_bytes=10) + assert data == b"x" * 11 # bounded to max_bytes + 1 (backstops a lying size) + assert handle.closed is True # handle is never left open + + 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()), []) diff --git a/tests/agents/test_workflow_run_scoped_artifacts.py b/tests/agents/test_workflow_run_scoped_artifacts.py index 678758c8..2fd1d726 100644 --- a/tests/agents/test_workflow_run_scoped_artifacts.py +++ b/tests/agents/test_workflow_run_scoped_artifacts.py @@ -97,6 +97,27 @@ def test_agent_node_run_scopes_tool_executor(monkeypatch): assert stub.tool_executor.conversation_id is None +def test_agent_node_skips_run_scope_when_not_persisted(monkeypatch): + """An unpersisted run must NOT stamp workflow_run_id (it would orphan artifacts).""" + engine = _create_engine(workflow_run_id="22222222-2222-2222-2222-222222222222") + engine.run_persisted = False + node = _agent_node() + stub = _StubNodeAgent([{"answer": "done"}]) + + monkeypatch.setattr( + WorkflowNodeAgentFactory, "create", staticmethod(lambda **kwargs: stub) + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", lambda _provider: None + ) + + list(engine._execute_agent_node(node)) + + # Left unset: the run-scoped tools then persist under a conversation parent or + # cleanly error, never orphaning an artifact under a nonexistent run row. + assert stub.tool_executor.workflow_run_id is None + + # --------------------------------------------------------------------------- # ToolExecutor stamping # --------------------------------------------------------------------------- diff --git a/tests/api/user/test_artifacts_routes.py b/tests/api/user/test_artifacts_routes.py index 5ba5df26..e7d34c2f 100644 --- a/tests/api/user/test_artifacts_routes.py +++ b/tests/api/user/test_artifacts_routes.py @@ -677,9 +677,30 @@ class TestMalformedArtifactId: # --------------------------------------------------------------------------- @pytest.mark.unit class TestApiKeyPrincipal: - def test_api_key_reads_artifact_in_its_agent_scope( + def test_api_key_reads_artifact_with_matching_conversation_id( self, _patch_db, flask_app, monkeypatch ): + # The public widget key reads an artifact ONLY when the request also + # carries the parent conversation_id (the per-visitor bearer capability). + from application.api.user.artifacts.routes import GetArtifact + + agent = _make_agent(_patch_db) + conv = _make_agent_conversation(_patch_db, agent["id"]) + art = _make_artifact(_patch_db, conversation_id=str(conv["id"])) + _wire_api_key(monkeypatch, _patch_db) + + resp = _call( + flask_app, GetArtifact, art["id"], token=None, + query={"api_key": "secret-key", "conversation_id": str(conv["id"])}, + ) + assert resp.status_code == 200 + assert resp.json["artifact"]["id"] == str(art["id"]) + + def test_api_key_get_without_conversation_id_denied( + self, _patch_db, flask_app, monkeypatch + ): + # Regression (critical IDOR): the public widget key alone -- without the + # unguessable per-visitor conversation_id -- cannot fetch a known artifact. from application.api.user.artifacts.routes import GetArtifact agent = _make_agent(_patch_db) @@ -691,15 +712,15 @@ class TestApiKeyPrincipal: flask_app, GetArtifact, art["id"], token=None, query={"api_key": "secret-key"}, ) - assert resp.status_code == 200 - assert resp.json["artifact"]["id"] == str(art["id"]) + assert resp.status_code == 403 def test_api_key_cannot_read_owner_artifact_outside_agent_scope( self, _patch_db, flask_app, monkeypatch ): # Regression (critical IDOR): a public widget key resolves to the owner # but must NOT reach an artifact from the owner's OTHER (non-agent) - # conversation -- it is confined to its own agent's conversations. + # conversation -- even when it supplies that conversation_id, the agent + # scope gate still blocks it. from application.api.user.artifacts.routes import GetArtifact _make_agent(_patch_db) # the widget agent (key=secret-key) @@ -709,31 +730,106 @@ class TestApiKeyPrincipal: resp = _call( flask_app, GetArtifact, art["id"], token=None, - query={"api_key": "secret-key"}, + query={"api_key": "secret-key", "conversation_id": str(other_conv["id"])}, ) assert resp.status_code == 403 - def test_api_key_list_is_scoped_to_agent_not_whole_corpus( + def test_api_key_download_requires_conversation_id( self, _patch_db, flask_app, monkeypatch ): - # Regression (critical IDOR): the no-filter list must return only the - # agent's own conversations' artifacts, never the owner's whole corpus. + # Download follows the same bearer-capability rule as get: the key alone is + # denied; the key + matching conversation_id is served. + from application.api.user.artifacts.routes import DownloadArtifact + + agent = _make_agent(_patch_db) + conv = _make_agent_conversation(_patch_db, agent["id"]) + art = _make_artifact( + _patch_db, conversation_id=str(conv["id"]), + filename="f.bin", storage_path="inputs/owner/artifacts/x/v1/f.bin", + ) + _wire_api_key(monkeypatch, _patch_db) + storage = _FakeStorage(b"BYTES") + 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, + ) + + denied = _call( + flask_app, DownloadArtifact, art["id"], token=None, + query={"api_key": "secret-key"}, + ) + assert denied.status_code == 403 + + ok = _call( + flask_app, DownloadArtifact, art["id"], token=None, + query={"api_key": "secret-key", "conversation_id": str(conv["id"])}, + ) + assert ok.status_code == 200 + assert ok.data == b"BYTES" + + def test_api_key_list_without_conversation_id_forbidden( + self, _patch_db, flask_app, monkeypatch + ): + # Regression (critical IDOR): the public widget key cannot enumerate the + # agent's artifacts -- a list with no conversation_id is refused outright. from application.api.user.artifacts.routes import ListArtifacts agent = _make_agent(_patch_db) - agent_conv = _make_agent_conversation(_patch_db, agent["id"]) - in_scope = _make_artifact(_patch_db, conversation_id=str(agent_conv["id"])) - other_conv = _make_conversation(_patch_db, user_id=OWNER) - out_of_scope = _make_artifact(_patch_db, conversation_id=str(other_conv["id"])) + conv = _make_agent_conversation(_patch_db, agent["id"]) + _make_artifact(_patch_db, conversation_id=str(conv["id"])) + _wire_api_key(monkeypatch, _patch_db) + + resp = _call(flask_app, ListArtifacts, token=None, query={"api_key": "secret-key"}) + assert resp.status_code == 403 + + def test_api_key_list_with_conversation_id_scoped_to_it( + self, _patch_db, flask_app, monkeypatch + ): + # With the per-visitor conversation_id the list returns only that + # conversation's artifacts -- not the agent's other conversations. + from application.api.user.artifacts.routes import ListArtifacts + + agent = _make_agent(_patch_db) + conv_a = _make_agent_conversation(_patch_db, agent["id"]) + conv_b = _make_agent_conversation(_patch_db, agent["id"]) + in_scope = _make_artifact(_patch_db, conversation_id=str(conv_a["id"])) + other = _make_artifact(_patch_db, conversation_id=str(conv_b["id"])) _wire_api_key(monkeypatch, _patch_db) resp = _call( - flask_app, ListArtifacts, token=None, query={"api_key": "secret-key"}, + flask_app, ListArtifacts, token=None, + query={"api_key": "secret-key", "conversation_id": str(conv_a["id"])}, ) assert resp.status_code == 200 ids = {a["id"] for a in resp.json["artifacts"]} assert str(in_scope["id"]) in ids - assert str(out_of_scope["id"]) not in ids + assert str(other["id"]) not in ids + + def test_api_key_list_foreign_conversation_id_returns_empty( + self, _patch_db, flask_app, monkeypatch + ): + # A conversation_id belonging to a DIFFERENT agent leaks nothing: the JOIN + # filters it out, yielding an empty 200 rather than a cross-visitor leak. + from application.api.user.artifacts.routes import ListArtifacts + + agent = _make_agent(_patch_db) + agent_conv = _make_agent_conversation(_patch_db, agent["id"]) + _make_artifact(_patch_db, conversation_id=str(agent_conv["id"])) + other_agent = _make_agent(_patch_db, user_id=OWNER, key="other-key") + other_conv = _make_agent_conversation(_patch_db, other_agent["id"]) + _make_artifact(_patch_db, conversation_id=str(other_conv["id"])) + _wire_api_key(monkeypatch, _patch_db) + + resp = _call( + flask_app, ListArtifacts, token=None, + query={"api_key": "secret-key", "conversation_id": str(other_conv["id"])}, + ) + assert resp.status_code == 200 + assert resp.json["artifacts"] == [] def test_api_key_cannot_delete_even_in_scope( self, _patch_db, flask_app, monkeypatch diff --git a/tests/storage/db/repositories/test_artifacts.py b/tests/storage/db/repositories/test_artifacts.py index 8c68c37f..31302d15 100644 --- a/tests/storage/db/repositories/test_artifacts.py +++ b/tests/storage/db/repositories/test_artifacts.py @@ -8,7 +8,9 @@ import pytest from sqlalchemy import text from sqlalchemy.exc import IntegrityError +from application.storage.db.repositories.agents import AgentsRepository from application.storage.db.repositories.artifacts import ArtifactsRepository +from application.storage.db.repositories.conversations import ConversationsRepository def _repo(conn) -> ArtifactsRepository: @@ -224,6 +226,57 @@ class TestListArtifacts: repo.list_artifacts() +class TestListArtifactsForAgent: + """Agent-scoped listing + optional conversation_id narrowing (public widget key path).""" + + @staticmethod + def _agent_with_conv(pg_conn, user_id: str, key: str): + agent = AgentsRepository(pg_conn).create(user_id, "widget", "active", key=key) + agent_id = str(agent["id"]) + conv = ConversationsRepository(pg_conn).create(user_id, name="c", agent_id=agent_id) + return agent_id, str(conv["id"]) + + def test_two_arg_returns_all_agent_artifacts(self, pg_conn): + # Back-compat: the 2-arg call (used by the MCP resource layer) still sees + # every artifact across the agent's conversations. + repo = _repo(pg_conn) + agent_id, conv_a = self._agent_with_conv(pg_conn, "owner-a", "k-a") + conv_b = str( + ConversationsRepository(pg_conn).create("owner-a", name="c2", agent_id=agent_id)["id"] + ) + repo.create_artifact("owner-a", "document", conversation_id=conv_a) + repo.create_artifact("owner-a", "document", conversation_id=conv_b) + # A conversation not owned by this agent is excluded by the JOIN. + stray = str(ConversationsRepository(pg_conn).create("owner-a", name="c3")["id"]) + repo.create_artifact("owner-a", "document", conversation_id=stray) + + rows = repo.list_artifacts_for_agent(agent_id, "owner-a") + assert len(rows) == 2 + + def test_conversation_id_narrows_to_single_conversation(self, pg_conn): + repo = _repo(pg_conn) + agent_id, conv_a = self._agent_with_conv(pg_conn, "owner-b", "k-b") + conv_b = str( + ConversationsRepository(pg_conn).create("owner-b", name="c2", agent_id=agent_id)["id"] + ) + a1 = repo.create_artifact("owner-b", "document", conversation_id=conv_a) + repo.create_artifact("owner-b", "document", conversation_id=conv_b) + + rows = repo.list_artifacts_for_agent(agent_id, "owner-b", conversation_id=conv_a) + assert [r["id"] for r in rows] == [a1["id"]] + + def test_foreign_conversation_id_yields_empty(self, pg_conn): + # A conversation_id that is not this agent's leaks nothing (JOIN filters it). + repo = _repo(pg_conn) + agent_id, conv_a = self._agent_with_conv(pg_conn, "owner-c", "k-c") + repo.create_artifact("owner-c", "document", conversation_id=conv_a) + other_agent_id, other_conv = self._agent_with_conv(pg_conn, "owner-c", "k-c2") + repo.create_artifact("owner-c", "document", conversation_id=other_conv) + + rows = repo.list_artifacts_for_agent(agent_id, "owner-c", conversation_id=other_conv) + assert rows == [] + + class TestQuotaAccounting: def test_count_for_user_counts_only_that_user(self, pg_conn): repo = _repo(pg_conn)