From 2d79ff2227404216061857cbeb56b46097a295cd Mon Sep 17 00:00:00 2001 From: Alex Date: Sat, 4 Jul 2026 15:48:40 +0200 Subject: [PATCH] feat: more QoL fixes --- application/sandbox/manager.py | 8 ++ .../services/artifact_resource_service.py | 28 ++++--- .../storage/db/repositories/artifacts.py | 84 +++++++++++++++++-- .../storage/db/repositories/workflows.py | 32 ++++++- tests/sandbox/test_sandbox_manager.py | 16 ++++ .../test_artifact_resource_service.py | 41 ++++++++- .../storage/db/repositories/test_workflows.py | 54 ++++++++++++ 7 files changed, 240 insertions(+), 23 deletions(-) diff --git a/application/sandbox/manager.py b/application/sandbox/manager.py index 2ddafa54..c471eb6e 100644 --- a/application/sandbox/manager.py +++ b/application/sandbox/manager.py @@ -93,6 +93,14 @@ class SandboxManager: now = time.monotonic() if session is not None and session.ready: session.last_access = now + # Reuse must honor an explicit longer keep-alive: a run_code + # persist/ttl on a session first opened by another tool (e.g. + # artifact_generator at the 60s exec timeout) would otherwise be + # dropped, and the kernel + its background state reaped early. + # Extend only (never shrink) so a default reuse can't cut short a + # session another caller kept alive. + if ttl is not None: + session.ttl = max(session.ttl, self._clamp_ttl(ttl)) return session.handle reaped = self._reap_locked(now) if session is None: diff --git a/application/services/artifact_resource_service.py b/application/services/artifact_resource_service.py index d1356c07..7e401f2e 100644 --- a/application/services/artifact_resource_service.py +++ b/application/services/artifact_resource_service.py @@ -3,11 +3,13 @@ The MCP server (``application/mcp_server.py``) authenticates a request with ``Authorization: Bearer ``; that key resolves to the owning ``user_id`` via ``AgentsRepository.find_by_key`` (the same api_key->owner path -as the HTTP artifact routes). Resources are scoped strictly to that principal: -``resources/list`` returns only the principal's owned artifacts, and -``resources/read`` re-checks ownership before serving any bytes. An -unresolvable principal yields an empty list / a denied read -- never another -principal's artifact. +as the HTTP artifact routes). Resources are scoped strictly to that principal +AND to the owner's own conversations: because the api_key is embedded publicly +in the widget and MCP carries no per-visitor conversation context, shared-usage +(widget/visitor) conversations are excluded. ``resources/list`` returns only the +owner's own-chat artifacts, and ``resources/read`` re-checks ownership + owner +scope before serving any bytes. An unresolvable principal yields an empty list / +a denied read -- never another principal's or another visitor's artifact. ``resources/list`` returns FastMCP ``Resource`` objects so they pass straight through the server's list/dedupe/wire pipeline; ``resources/read`` is served by @@ -133,13 +135,12 @@ 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. + # MCP-over-key has no per-visitor conversation context, and the api_key + # is embedded publicly in the widget, so scope to the owner's OWN + # (non-shared-usage) conversations only. Without this, the key would + # enumerate artifacts from every visitor's conversation with the agent. rows = ArtifactsRepository(conn).list_artifacts_for_agent( - str(agent["id"]), str(agent["user_id"]) + str(agent["id"]), str(agent["user_id"]), owner_only=True ) except Exception: logger.exception("artifact resource: list failed") @@ -209,7 +210,10 @@ def read_artifact_resource(api_key: Optional[str], uri: str) -> ArtifactReadResu raise ResourceDenied("forbidden") # Agent scope is the second: a per-agent key only reads artifacts from # its own conversations, not the owner's other agents / workflow runs. - if not repo.artifact_in_agent_scope(artifact_id, agent_id): + # ``owner_only`` further excludes shared-usage (widget/visitor) + # conversations: the public key reads the owner's own chats only, never + # another visitor's artifact (parity with the HTTP route's confinement). + if not repo.artifact_in_agent_scope(artifact_id, agent_id, owner_only=True): raise ResourceDenied("forbidden") version_row = repo.get_version(artifact_id, version) except (DataError, DBAPIError) as exc: diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py index ebe26aa6..88fd7643 100644 --- a/application/storage/db/repositories/artifacts.py +++ b/application/storage/db/repositories/artifacts.py @@ -13,12 +13,15 @@ only metadata + the ``BaseStorage`` key are stored here, never binary bytes. from __future__ import annotations import json +import logging from typing import Any, Optional from sqlalchemy import Connection, text from application.storage.db.base_repository import row_to_dict +logger = logging.getLogger(__name__) + def _artifact_to_dict(row: Any) -> dict: """row_to_dict for an ``artifacts`` row (``metadata`` key preserved).""" @@ -199,20 +202,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, conversation_id: Optional[str] = None + self, + agent_id: str, + user_id: str, + conversation_id: Optional[str] = None, + owner_only: bool = False, ) -> 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. When ``conversation_id`` is given, the SQL narrows to that single conversation - (the per-visitor bearer capability for a public widget key). + (the per-visitor bearer capability for a public widget key). When + ``owner_only`` is set, shared-usage (widget/visitor) conversations are + excluded so a credential with no per-visitor conversation context (the MCP + Bearer key) sees only the owner's own chats, never other visitors' artifacts. """ 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 + if owner_only: + clauses.append("c.is_shared_usage = false") result = self._conn.execute( text( "SELECT a.* FROM artifacts a " @@ -224,13 +236,26 @@ class ArtifactsRepository: ) return [_artifact_to_dict(r) for r in result.fetchall()] - def artifact_in_agent_scope(self, artifact_id: str, agent_id: str) -> bool: - """True if ``artifact_id``'s parent conversation belongs to ``agent_id``.""" + def artifact_in_agent_scope( + self, artifact_id: str, agent_id: str, owner_only: bool = False + ) -> bool: + """True if ``artifact_id``'s parent conversation belongs to ``agent_id``. + + With ``owner_only`` set, shared-usage (widget/visitor) conversations are + excluded so a per-visitor-less credential (the MCP Bearer key) can only + confirm scope for the owner's own chats, never another visitor's artifact. + """ + clauses = [ + "a.id = CAST(:id AS uuid)", + "c.agent_id = CAST(:agent_id AS uuid)", + ] + if owner_only: + clauses.append("c.is_shared_usage = false") result = self._conn.execute( text( "SELECT 1 FROM artifacts a " "JOIN conversations c ON a.conversation_id = c.id " - "WHERE a.id = CAST(:id AS uuid) AND c.agent_id = CAST(:agent_id AS uuid) " + f"WHERE {' AND '.join(clauses)} " "LIMIT 1" ), {"id": artifact_id, "agent_id": str(agent_id)}, @@ -445,6 +470,55 @@ class ArtifactsRepository: ) return result.rowcount + def storage_paths_for_workflow(self, workflow_id: str) -> list[str]: + """Return every stored version path for artifacts under a workflow's runs (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 " + "JOIN workflow_runs r ON r.id = a.workflow_run_id " + "WHERE r.workflow_id = CAST(:wid AS uuid) " + "AND v.storage_path IS NOT NULL" + ), + {"wid": workflow_id}, + ) + return [r[0] for r in result.fetchall()] + + def delete_for_workflow(self, workflow_id: str) -> int: + """Delete all artifacts (+ versions, via cascade) parented to a workflow's runs. + + ``artifacts.workflow_run_id`` is a bare uuid column (no FK), so deleting the + workflow (which cascade-deletes its ``workflow_runs``) does NOT reap these + rows -- callers invoke this to reclaim the rows + their quota. Must run + BEFORE the workflow delete, while the run rows still resolve the subquery. + """ + result = self._conn.execute( + text( + "DELETE FROM artifacts WHERE workflow_run_id IN " + "(SELECT id FROM workflow_runs WHERE workflow_id = CAST(:wid AS uuid))" + ), + {"wid": workflow_id}, + ) + return result.rowcount + + @staticmethod + def reap_storage(paths: list) -> None: + """Best-effort delete a batch of artifact byte objects; never raises.""" + if not paths: + return + try: + from application.storage.storage_creator import StorageCreator + + storage = StorageCreator.get_storage() + except Exception: + logger.warning("artifact cleanup: storage unavailable", exc_info=True) + return + for path in paths: + try: + storage.delete_file(path) + except Exception: + logger.warning("artifact cleanup: failed to delete bytes %s", path, exc_info=True) + def get_version(self, artifact_id: str, version: int) -> Optional[dict]: """Fetch a single version row by artifact id and version number.""" result = self._conn.execute( diff --git a/application/storage/db/repositories/workflows.py b/application/storage/db/repositories/workflows.py index 58cc5fd9..bc9e4cc9 100644 --- a/application/storage/db/repositories/workflows.py +++ b/application/storage/db/repositories/workflows.py @@ -13,7 +13,7 @@ from typing import Optional from sqlalchemy import Connection, text from sqlalchemy.dialects.postgresql import insert as pg_insert -from application.storage.db.base_repository import row_to_dict +from application.storage.db.base_repository import looks_like_uuid, row_to_dict from application.storage.db.models import workflows_table @@ -116,6 +116,31 @@ class WorkflowsRepository: return row[0] if row else None def delete(self, workflow_id: str, user_id: str) -> bool: + # Run artifacts parent to ``workflow_runs`` (which cascade-delete with the + # workflow), but ``artifacts.workflow_run_id`` is a bare uuid with no FK, so + # the rows + bytes + the quota they consume would leak. Reclaim them BEFORE + # deleting the workflow, while the run rows still resolve the subquery. + # Ownership is confirmed first so this never reaps another user's artifacts. + from application.storage.db.repositories.artifacts import ArtifactsRepository + + if not looks_like_uuid(workflow_id): + return False + owned = ( + self._conn.execute( + text( + "SELECT 1 FROM workflows " + "WHERE id = CAST(:id AS uuid) AND user_id = :user_id" + ), + {"id": workflow_id, "user_id": user_id}, + ).fetchone() + is not None + ) + if not owned: + return False + + artifacts = ArtifactsRepository(self._conn) + paths = artifacts.storage_paths_for_workflow(workflow_id) + artifacts.delete_for_workflow(workflow_id) result = self._conn.execute( text( "DELETE FROM workflows " @@ -123,7 +148,10 @@ class WorkflowsRepository: ), {"id": workflow_id, "user_id": user_id}, ) - return result.rowcount > 0 + deleted = result.rowcount > 0 + if deleted: + artifacts.reap_storage(paths) + return deleted def delete_by_legacy_id(self, legacy_mongo_id: str, user_id: str) -> bool: """Delete a workflow addressed by the Mongo ObjectId string. diff --git a/tests/sandbox/test_sandbox_manager.py b/tests/sandbox/test_sandbox_manager.py index 82fcb8a7..02b9abcb 100644 --- a/tests/sandbox/test_sandbox_manager.py +++ b/tests/sandbox/test_sandbox_manager.py @@ -120,6 +120,22 @@ def test_ttl_honored_when_below_max(backend): assert mgr.ttl_for("conv-1") == 120 +def test_reuse_extends_ttl_but_never_shrinks(backend): + # A tool (e.g. artifact_generator) may open the shared session at the short + # exec timeout; a later run_code(persist, ttl=...) reuses it and must be able + # to extend the keep-alive, or the kernel + background state are reaped early. + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("conv-1", ttl=60) + assert mgr.ttl_for("conv-1") == 60 + mgr.open("conv-1", ttl=600) # reuse: explicit longer ttl extends + assert mgr.ttl_for("conv-1") == 600 + mgr.open("conv-1", ttl=120) # reuse: shorter ttl must not shrink + assert mgr.ttl_for("conv-1") == 600 + mgr.open("conv-1") # reuse: default (ttl=None) leaves it untouched + assert mgr.ttl_for("conv-1") == 600 + assert backend.open_calls == ["conv-1"] # still opened only once + + def test_ttl_default_used_for_nonpositive(backend): mgr = SandboxManager(backend, max_ttl=300, default_ttl=200) mgr.open("conv-1", ttl=0) diff --git a/tests/services/test_artifact_resource_service.py b/tests/services/test_artifact_resource_service.py index 1df1dfaf..2bd26903 100644 --- a/tests/services/test_artifact_resource_service.py +++ b/tests/services/test_artifact_resource_service.py @@ -32,6 +32,9 @@ ART_BIN = "22222222-2222-4222-8222-222222222222" ART_FOREIGN = "33333333-3333-4333-8333-333333333333" # Owned by OWNER but produced by a DIFFERENT agent — must be invisible to owner-key. ART_OTHER_AGENT = "44444444-4444-4444-8444-444444444444" +# In the owner's agent scope but from a shared-usage (widget/visitor) conversation +# — must be invisible to the public MCP key (owner-only scope). +ART_SHARED = "55555555-5555-4555-8555-555555555555" @contextmanager @@ -63,16 +66,23 @@ class _FakeArtifacts: def __init__(self, conn): pass - def list_artifacts_for_agent(self, agent_id, user_id): + def list_artifacts_for_agent(self, agent_id, user_id, conversation_id=None, owner_only=False): return [ a for a in self.artifacts.values() - if a.get("agent_id") == agent_id and a["user_id"] == user_id + if a.get("agent_id") == agent_id + and a["user_id"] == user_id + and (conversation_id is None or a.get("conversation_id") == conversation_id) + and not (owner_only and a.get("is_shared_usage")) ] - def artifact_in_agent_scope(self, artifact_id, agent_id): + def artifact_in_agent_scope(self, artifact_id, agent_id, owner_only=False): art = self.artifacts.get(artifact_id) - return art is not None and art.get("agent_id") == agent_id + return ( + art is not None + and art.get("agent_id") == agent_id + and not (owner_only and art.get("is_shared_usage")) + ) def get_artifact(self, artifact_id): return self.artifacts.get(artifact_id) @@ -112,12 +122,22 @@ def _wire(monkeypatch): "title": "other-agent", "current_version": 1, }, + ART_SHARED: { + "id": ART_SHARED, + "user_id": OWNER, + "agent_id": "agent-owner", + "kind": "data", + "title": "visitor-doc", + "current_version": 1, + "is_shared_usage": True, + }, } _FakeArtifacts.versions = { (ART_TEXT, 2): {"mime_type": "text/csv", "storage_path": "k/text.csv", "preview_text": None}, (ART_BIN, 1): {"mime_type": "image/png", "storage_path": "k/chart.png", "preview_text": None}, (ART_FOREIGN, 1): {"mime_type": "text/plain", "storage_path": "k/secret.txt", "preview_text": None}, (ART_OTHER_AGENT, 1): {"mime_type": "text/plain", "storage_path": "k/other.txt", "preview_text": None}, + (ART_SHARED, 1): {"mime_type": "text/plain", "storage_path": "k/visitor.txt", "preview_text": None}, } monkeypatch.setattr(svc, "db_readonly", _fake_conn) monkeypatch.setattr(svc, "AgentsRepository", _FakeAgents) @@ -133,6 +153,13 @@ class TestListArtifactResources: assert uris == {f"artifact://{ART_TEXT}/v2", f"artifact://{ART_BIN}/v1"} assert f"artifact://{ART_FOREIGN}/v1" not in uris + def test_excludes_shared_usage_visitor_artifacts(self): + # The MCP Bearer key is embedded publicly in the widget and has no + # per-visitor conversation context, so shared-usage (visitor) artifacts + # must never be enumerated — only the owner's own-chat artifacts. + uris = {str(r.uri) for r in svc.list_artifact_resources("owner-key")} + assert f"artifact://{ART_SHARED}/v1" not in uris + def test_unresolved_principal_is_empty(self): assert svc.list_artifact_resources("bogus-key") == [] @@ -185,6 +212,12 @@ class TestReadArtifactResource: with pytest.raises(svc.ResourceDenied): svc.read_artifact_resource("owner-key", f"artifact://{ART_OTHER_AGENT}/v1") + def test_read_denies_shared_usage_visitor_artifact(self): + # In the owner's agent scope but produced in a shared-usage (widget/visitor) + # conversation: the public MCP key must not read another visitor's artifact. + with pytest.raises(svc.ResourceDenied): + svc.read_artifact_resource("owner-key", f"artifact://{ART_SHARED}/v1") + def test_foreign_owner_is_denied(self): with pytest.raises(svc.ResourceDenied): svc.read_artifact_resource("owner-key", f"artifact://{ART_FOREIGN}/v1") diff --git a/tests/storage/db/repositories/test_workflows.py b/tests/storage/db/repositories/test_workflows.py index 27f507a0..a09cc39e 100644 --- a/tests/storage/db/repositories/test_workflows.py +++ b/tests/storage/db/repositories/test_workflows.py @@ -109,3 +109,57 @@ class TestDelete: repo = _repo(pg_conn) created = repo.create("user-1", "wf") assert repo.delete(created["id"], "other") is False + + def test_delete_reaps_run_artifacts(self, pg_conn, monkeypatch): + # Run artifacts have no FK cascade off the workflow, so deleting the + # workflow must explicitly reclaim their rows (quota) and bytes. + from application.storage.db.repositories.artifacts import ArtifactsRepository + from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository + from application.storage.storage_creator import StorageCreator + + deleted_paths: list[str] = [] + + class _RecordingStorage: + def delete_file(self, path): + deleted_paths.append(path) + + monkeypatch.setattr( + StorageCreator, "get_storage", staticmethod(lambda: _RecordingStorage()) + ) + + repo = _repo(pg_conn) + artifacts = ArtifactsRepository(pg_conn) + runs = WorkflowRunsRepository(pg_conn) + + wf = repo.create("user-1", "wf") + run = runs.create(wf["id"], "user-1", "running") + art = artifacts.create_artifact( + "user-1", "document", workflow_run_id=run["id"], storage_path="k/run.bin" + ) + assert artifacts.get_artifact(art["id"]) is not None + + assert repo.delete(wf["id"], "user-1") is True + + # Rows (and their quota) reclaimed, and the bytes reaped best-effort. + assert artifacts.get_artifact(art["id"]) is None + assert artifacts.count_for_user("user-1") == 0 + assert deleted_paths == ["k/run.bin"] + + def test_delete_wrong_user_keeps_run_artifacts(self, pg_conn): + # The ownership guard must run before any artifact cleanup, so a + # non-owner delete can never reap another user's run artifacts. + from application.storage.db.repositories.artifacts import ArtifactsRepository + from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository + + repo = _repo(pg_conn) + artifacts = ArtifactsRepository(pg_conn) + runs = WorkflowRunsRepository(pg_conn) + + wf = repo.create("user-1", "wf") + run = runs.create(wf["id"], "user-1", "running") + art = artifacts.create_artifact( + "user-1", "document", workflow_run_id=run["id"], storage_path="k/run.bin" + ) + + assert repo.delete(wf["id"], "intruder") is False + assert artifacts.get_artifact(art["id"]) is not None