From 31525621e083f8af8aefd31708b920b83060f4b9 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 24 Jun 2026 10:21:48 +0100 Subject: [PATCH 01/59] Add artifacts data model (tables, migration, repository) Introduce artifacts and artifact_versions tables: an append-only, pass-by-reference store with a stable identity row and immutable versions that atomically bump current_version. Access is parent-derived (conversation or workflow run); a CHECK enforces at least one parent and a nullable team_id is reserved for future team scoping. Adds the 0025_artifacts migration, ArtifactsRepository with a parent-scoped safe getter alongside the unscoped one, and repo tests. --- .../alembic/versions/0025_artifacts.py | 111 +++++++ application/storage/db/models.py | 68 +++++ .../storage/db/repositories/artifacts.py | 270 ++++++++++++++++++ .../storage/db/repositories/test_artifacts.py | 241 ++++++++++++++++ 4 files changed, 690 insertions(+) create mode 100644 application/alembic/versions/0025_artifacts.py create mode 100644 application/storage/db/repositories/artifacts.py create mode 100644 tests/storage/db/repositories/test_artifacts.py diff --git a/application/alembic/versions/0025_artifacts.py b/application/alembic/versions/0025_artifacts.py new file mode 100644 index 00000000..d36eeece --- /dev/null +++ b/application/alembic/versions/0025_artifacts.py @@ -0,0 +1,111 @@ +"""0025 artifacts — append-only artifact identity + version store. + +Adds the long-anticipated ``artifacts`` entity referenced by the schema header +in ``models.py`` but never built. Two tables: + +* ``artifacts`` — one identity row per logical artifact. The stable ``id`` is the + handle passed around (chat/workflow ``state``/message bodies carry only this + reference, never bytes). ``current_version`` mirrors + ``workflows.current_graph_version`` (atomic increment on each new version). + Authz is parent-derived: whoever can reach the ``conversation_id`` (chat) or + ``workflow_run_id`` (run) can reach its artifacts, so a CHECK requires at least + one parent. ``user_id`` is kept for ownership/quota only; ``team_id`` is a + nullable forward-compat hook for the in-flight Teams work (no FK, matching how + ``teams`` itself leaves owner/member columns FK-free). +* ``artifact_versions`` — append-only, never mutated. Each edit appends a row and + bumps ``artifacts.current_version``. ``UNIQUE(artifact_id, version)`` enforces + monotonic, gap-checkable versions. ``storage_path`` is the ``BaseStorage`` key + (NULL when the version is spec-only). Pass-by-reference: bytes live in storage, + only metadata + the key live here. + +Revision ID: 0025_artifacts +Revises: 0024_wiki_pages_updated_via +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0025_artifacts" +down_revision: Union[str, None] = "0024_wiki_pages_updated_via" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # --- artifacts ----------------------------------------------------------- + op.execute( + """ + CREATE TABLE IF NOT EXISTS artifacts ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id TEXT NOT NULL, + conversation_id UUID, + workflow_run_id UUID, + team_id UUID, + message_id UUID, + kind TEXT NOT NULL, + title TEXT, + metadata JSONB, + current_version INTEGER NOT NULL DEFAULT 1, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CONSTRAINT artifacts_parent_present_check + CHECK (conversation_id IS NOT NULL OR workflow_run_id IS NOT NULL) + ); + """ + ) + op.execute("CREATE INDEX IF NOT EXISTS artifacts_user_idx ON artifacts (user_id);") + op.execute( + "CREATE INDEX IF NOT EXISTS artifacts_conversation_idx " + "ON artifacts (conversation_id) WHERE conversation_id IS NOT NULL;" + ) + op.execute( + "CREATE INDEX IF NOT EXISTS artifacts_workflow_run_idx " + "ON artifacts (workflow_run_id) WHERE workflow_run_id IS NOT NULL;" + ) + # Reuse the shared set_updated_at() trigger fn defined in 0001. + op.execute("DROP TRIGGER IF EXISTS artifacts_set_updated_at ON artifacts;") + op.execute( + """ + CREATE TRIGGER artifacts_set_updated_at + BEFORE UPDATE ON artifacts + FOR EACH ROW EXECUTE FUNCTION set_updated_at(); + """ + ) + + # --- artifact_versions --------------------------------------------------- + op.execute( + """ + CREATE TABLE IF NOT EXISTS artifact_versions ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + artifact_id UUID NOT NULL REFERENCES artifacts(id) ON DELETE CASCADE, + version INTEGER NOT NULL, + mime_type TEXT, + filename TEXT, + storage_path TEXT, + size BIGINT, + sha256 TEXT, + spec JSONB, + preview_text TEXT, + produced_by JSONB, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CONSTRAINT artifact_versions_artifact_version_uidx + UNIQUE (artifact_id, version) + ); + """ + ) + op.execute( + "CREATE INDEX IF NOT EXISTS artifact_versions_artifact_idx " + "ON artifact_versions (artifact_id);" + ) + + +def downgrade() -> None: + op.execute("DROP INDEX IF EXISTS artifact_versions_artifact_idx;") + op.execute("DROP TABLE IF EXISTS artifact_versions;") + op.execute("DROP TRIGGER IF EXISTS artifacts_set_updated_at ON artifacts;") + op.execute("DROP INDEX IF EXISTS artifacts_workflow_run_idx;") + op.execute("DROP INDEX IF EXISTS artifacts_conversation_idx;") + op.execute("DROP INDEX IF EXISTS artifacts_user_idx;") + op.execute("DROP TABLE IF EXISTS artifacts;") diff --git a/application/storage/db/models.py b/application/storage/db/models.py index 50df6bcb..261e5c1b 100644 --- a/application/storage/db/models.py +++ b/application/storage/db/models.py @@ -401,6 +401,74 @@ attachments_table = Table( Column("legacy_mongo_id", Text), ) +# Identity row, one per logical artifact. The stable ``id`` is the handle passed +# around (chat/workflow state/message bodies carry only this reference, never +# bytes). Authz is parent-derived: whoever can reach ``conversation_id`` (chat) +# or ``workflow_run_id`` (run) can reach its artifacts, so a CHECK requires at +# least one parent. ``user_id`` is ownership/quota only; ``team_id`` is a +# nullable forward-compat hook for Teams (no FK, matching ``teams`` itself). +artifacts_table = Table( + "artifacts", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()), + Column("user_id", Text, nullable=False), + Column("conversation_id", UUID(as_uuid=True)), + Column("workflow_run_id", UUID(as_uuid=True)), + Column("team_id", UUID(as_uuid=True)), + Column("message_id", UUID(as_uuid=True)), + Column("kind", Text, nullable=False), + Column("title", Text), + Column("metadata", JSONB), + Column("current_version", Integer, nullable=False, server_default="1"), + Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + CheckConstraint( + "conversation_id IS NOT NULL OR workflow_run_id IS NOT NULL", + name="artifacts_parent_present_check", + ), +) + +Index( + "artifacts_conversation_idx", + artifacts_table.c.conversation_id, + postgresql_where=artifacts_table.c.conversation_id.isnot(None), +) +Index( + "artifacts_workflow_run_idx", + artifacts_table.c.workflow_run_id, + postgresql_where=artifacts_table.c.workflow_run_id.isnot(None), +) +Index("artifacts_user_idx", artifacts_table.c.user_id) + +# Append-only version history; never mutated. Each edit appends a row and bumps +# ``artifacts.current_version``. ``UNIQUE(artifact_id, version)`` keeps versions +# monotonic. ``storage_path`` is the ``BaseStorage`` key (NULL when spec-only); +# bytes live in storage, only metadata + the key live here (pass-by-reference). +artifact_versions_table = Table( + "artifact_versions", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()), + Column( + "artifact_id", + UUID(as_uuid=True), + ForeignKey("artifacts.id", ondelete="CASCADE"), + nullable=False, + ), + Column("version", Integer, nullable=False), + Column("mime_type", Text), + Column("filename", Text), + Column("storage_path", Text), + Column("size", BigInteger), + Column("sha256", Text), + Column("spec", JSONB), + Column("preview_text", Text), + Column("produced_by", JSONB), + Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + UniqueConstraint("artifact_id", "version", name="artifact_versions_artifact_version_uidx"), +) + +Index("artifact_versions_artifact_idx", artifact_versions_table.c.artifact_id) + memories_table = Table( "memories", metadata, diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py new file mode 100644 index 00000000..86c81863 --- /dev/null +++ b/application/storage/db/repositories/artifacts.py @@ -0,0 +1,270 @@ +"""Repository for the ``artifacts`` / ``artifact_versions`` tables. + +Append-only artifact store. An ``artifacts`` row is the stable identity; each +edit appends an immutable ``artifact_versions`` row and atomically bumps +``artifacts.current_version`` (mirroring ``workflows.increment_graph_version``). + +Authz is parent-derived: reads are not gated on ``user_id`` — callers resolve +access via the artifact's ``conversation_id`` / ``workflow_run_id`` parent. +``user_id`` is carried for ownership / quota accounting only. Pass-by-reference: +only metadata + the ``BaseStorage`` key are stored here, never binary bytes. +""" + +from __future__ import annotations + +import json +from typing import Any, Optional + +from sqlalchemy import Connection, text + +from application.storage.db.base_repository import row_to_dict + + +def _artifact_to_dict(row: Any) -> dict: + """row_to_dict for an ``artifacts`` row (``metadata`` key preserved).""" + return row_to_dict(row) + + +def _version_to_dict(row: Any) -> dict: + """row_to_dict for an ``artifact_versions`` row (``spec``/``produced_by`` keys preserved).""" + return row_to_dict(row) + + +class ArtifactsRepository: + """CRUD for artifact identities and their append-only versions.""" + + def __init__(self, conn: Connection) -> None: + """Bind the repository to an open SQLAlchemy connection.""" + self._conn = conn + + def create_artifact( + self, + user_id: str, + kind: str, + *, + conversation_id: Optional[str] = None, + workflow_run_id: Optional[str] = None, + team_id: Optional[str] = None, + message_id: Optional[str] = None, + title: Optional[str] = None, + metadata: Any = None, + mime_type: Optional[str] = None, + filename: Optional[str] = None, + storage_path: Optional[str] = None, + size: Optional[int] = None, + sha256: Optional[str] = None, + spec: Any = None, + preview_text: Optional[str] = None, + produced_by: Any = None, + ) -> dict: + """Create the identity row and its version 1 atomically; return the artifact dict.""" + artifact = self._conn.execute( + text( + """ + INSERT INTO artifacts ( + user_id, conversation_id, workflow_run_id, team_id, + message_id, kind, title, metadata, current_version + ) + VALUES ( + :user_id, + CAST(:conversation_id AS uuid), + CAST(:workflow_run_id AS uuid), + CAST(:team_id AS uuid), + CAST(:message_id AS uuid), + :kind, :title, CAST(:metadata AS jsonb), 1 + ) + RETURNING * + """ + ), + { + "user_id": user_id, + "conversation_id": conversation_id, + "workflow_run_id": workflow_run_id, + "team_id": team_id, + "message_id": message_id, + "kind": kind, + "title": title, + "metadata": json.dumps(metadata) if metadata is not None else None, + }, + ).fetchone() + artifact_dict = _artifact_to_dict(artifact) + self._insert_version( + artifact_id=artifact_dict["id"], + version=1, + mime_type=mime_type, + filename=filename, + storage_path=storage_path, + size=size, + sha256=sha256, + spec=spec, + preview_text=preview_text, + produced_by=produced_by, + ) + return artifact_dict + + def get_artifact(self, artifact_id: str) -> Optional[dict]: + """Fetch one artifact by id, unscoped (prefer ``get_artifact_in_parent`` at request boundaries).""" + result = self._conn.execute( + text("SELECT * FROM artifacts WHERE id = CAST(:id AS uuid)"), + {"id": artifact_id}, + ) + row = result.fetchone() + return _artifact_to_dict(row) if row is not None else None + + def get_artifact_in_parent( + self, + artifact_id: str, + *, + conversation_id: Optional[str] = None, + workflow_run_id: Optional[str] = None, + ) -> Optional[dict]: + """Fetch an artifact only if it belongs to the given parent; the safe gate for request handlers.""" + if conversation_id is None and workflow_run_id is None: + raise ValueError("get_artifact_in_parent requires conversation_id or workflow_run_id") + clauses = ["id = CAST(:id AS uuid)"] + params: dict[str, Any] = {"id": artifact_id} + if conversation_id is not None: + clauses.append("conversation_id = CAST(:conversation_id AS uuid)") + params["conversation_id"] = conversation_id + if workflow_run_id is not None: + clauses.append("workflow_run_id = CAST(:workflow_run_id AS uuid)") + params["workflow_run_id"] = workflow_run_id + result = self._conn.execute( + text(f"SELECT * FROM artifacts WHERE {' AND '.join(clauses)}"), + params, + ) + row = result.fetchone() + return _artifact_to_dict(row) if row is not None else None + + def list_artifacts( + self, + conversation_id: Optional[str] = None, + workflow_run_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> list[dict]: + """List artifacts filtered by parent / owner (newest first); at least one filter is required.""" + if conversation_id is None and workflow_run_id is None and user_id is None: + raise ValueError("list_artifacts requires at least one of conversation_id, workflow_run_id, user_id") + clauses: list[str] = [] + params: dict[str, Any] = {} + if conversation_id is not None: + clauses.append("conversation_id = CAST(:conversation_id AS uuid)") + params["conversation_id"] = conversation_id + if workflow_run_id is not None: + clauses.append("workflow_run_id = CAST(:workflow_run_id AS uuid)") + params["workflow_run_id"] = workflow_run_id + if user_id is not None: + clauses.append("user_id = :user_id") + params["user_id"] = user_id + where = f"WHERE {' AND '.join(clauses)}" if clauses else "" + result = self._conn.execute( + text(f"SELECT * FROM artifacts {where} ORDER BY created_at DESC"), + params, + ) + return [_artifact_to_dict(r) for r in result.fetchall()] + + 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( + text( + "SELECT * FROM artifact_versions " + "WHERE artifact_id = CAST(:artifact_id AS uuid) AND version = :version" + ), + {"artifact_id": artifact_id, "version": version}, + ) + row = result.fetchone() + return _version_to_dict(row) if row is not None else None + + def list_versions(self, artifact_id: str) -> list[dict]: + """List every version of an artifact, oldest first.""" + result = self._conn.execute( + text( + "SELECT * FROM artifact_versions " + "WHERE artifact_id = CAST(:artifact_id AS uuid) ORDER BY version ASC" + ), + {"artifact_id": artifact_id}, + ) + return [_version_to_dict(r) for r in result.fetchall()] + + def append_version( + self, + artifact_id: str, + *, + mime_type: Optional[str] = None, + filename: Optional[str] = None, + storage_path: Optional[str] = None, + size: Optional[int] = None, + sha256: Optional[str] = None, + spec: Any = None, + preview_text: Optional[str] = None, + produced_by: Any = None, + ) -> dict: + """Append a new version and atomically bump ``current_version``; return the version dict.""" + new_version = self._conn.execute( + text( + "UPDATE artifacts " + "SET current_version = current_version + 1, updated_at = now() " + "WHERE id = CAST(:id AS uuid) " + "RETURNING current_version" + ), + {"id": artifact_id}, + ).fetchone() + if new_version is None: + raise ValueError(f"artifact {artifact_id} not found") + return self._insert_version( + artifact_id=artifact_id, + version=new_version[0], + mime_type=mime_type, + filename=filename, + storage_path=storage_path, + size=size, + sha256=sha256, + spec=spec, + preview_text=preview_text, + produced_by=produced_by, + ) + + def _insert_version( + self, + *, + artifact_id: str, + version: int, + mime_type: Optional[str], + filename: Optional[str], + storage_path: Optional[str], + size: Optional[int], + sha256: Optional[str], + spec: Any, + preview_text: Optional[str], + produced_by: Any, + ) -> dict: + """Insert one append-only version row; UNIQUE(artifact_id, version) guards duplicates.""" + result = self._conn.execute( + text( + """ + INSERT INTO artifact_versions ( + artifact_id, version, mime_type, filename, storage_path, + size, sha256, spec, preview_text, produced_by + ) + VALUES ( + CAST(:artifact_id AS uuid), :version, :mime_type, :filename, + :storage_path, :size, :sha256, CAST(:spec AS jsonb), + :preview_text, CAST(:produced_by AS jsonb) + ) + RETURNING * + """ + ), + { + "artifact_id": artifact_id, + "version": version, + "mime_type": mime_type, + "filename": filename, + "storage_path": storage_path, + "size": size, + "sha256": sha256, + "spec": json.dumps(spec) if spec is not None else None, + "preview_text": preview_text, + "produced_by": json.dumps(produced_by) if produced_by is not None else None, + }, + ) + return _version_to_dict(result.fetchone()) diff --git a/tests/storage/db/repositories/test_artifacts.py b/tests/storage/db/repositories/test_artifacts.py new file mode 100644 index 00000000..15bb9aca --- /dev/null +++ b/tests/storage/db/repositories/test_artifacts.py @@ -0,0 +1,241 @@ +"""Tests for ArtifactsRepository against a real Postgres instance.""" + +from __future__ import annotations + +import uuid + +import pytest +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError + +from application.storage.db.repositories.artifacts import ArtifactsRepository + + +def _repo(conn) -> ArtifactsRepository: + return ArtifactsRepository(conn) + + +def _conversation_id() -> str: + return str(uuid.uuid4()) + + +class TestCreateArtifact: + def test_creates_identity_and_auto_v1(self, pg_conn): + repo = _repo(pg_conn) + conv = _conversation_id() + artifact = repo.create_artifact( + "user-1", + "presentation", + conversation_id=conv, + title="Q3 deck", + metadata={"source": "chat"}, + mime_type="application/vnd.ms-powerpoint", + filename="deck.pptx", + storage_path="inputs/user-1/artifacts/x/v1/deck.pptx", + size=1234, + sha256="abc", + spec={"slides": []}, + ) + assert artifact["user_id"] == "user-1" + assert artifact["kind"] == "presentation" + assert artifact["title"] == "Q3 deck" + assert artifact["metadata"] == {"source": "chat"} + assert artifact["current_version"] == 1 + assert artifact["id"] is not None + assert artifact["_id"] == artifact["id"] + + v1 = repo.get_version(artifact["id"], 1) + assert v1 is not None + assert v1["version"] == 1 + assert v1["filename"] == "deck.pptx" + assert v1["storage_path"] == "inputs/user-1/artifacts/x/v1/deck.pptx" + assert v1["size"] == 1234 + assert v1["sha256"] == "abc" + assert v1["spec"] == {"slides": []} + + def test_spec_only_version_allows_null_storage_path(self, pg_conn): + repo = _repo(pg_conn) + artifact = repo.create_artifact( + "user-1", + "document", + workflow_run_id=_conversation_id(), + spec={"body": "draft"}, + ) + v1 = repo.get_version(artifact["id"], 1) + assert v1["storage_path"] is None + assert v1["spec"] == {"body": "draft"} + + def test_produced_by_and_preview_round_trip(self, pg_conn): + repo = _repo(pg_conn) + artifact = repo.create_artifact( + "user-1", + "code", + conversation_id=_conversation_id(), + preview_text="head of output", + produced_by={"tool_id": "code_executor", "session_id": "s1"}, + ) + v1 = repo.get_version(artifact["id"], 1) + assert v1["preview_text"] == "head of output" + assert v1["produced_by"] == {"tool_id": "code_executor", "session_id": "s1"} + + def test_missing_both_parents_rejected(self, pg_conn): + repo = _repo(pg_conn) + try: + with pg_conn.begin_nested(): + repo.create_artifact("user-1", "document") + except IntegrityError: + pass + else: + pytest.fail("expected IntegrityError when no parent is present") + + +class TestGetArtifact: + def test_get_existing(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create_artifact( + "user-1", "document", conversation_id=_conversation_id() + ) + fetched = repo.get_artifact(created["id"]) + assert fetched["id"] == created["id"] + + def test_get_missing_returns_none(self, pg_conn): + repo = _repo(pg_conn) + assert repo.get_artifact(str(uuid.uuid4())) is None + + def test_get_not_gated_on_user(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create_artifact( + "owner", "document", conversation_id=_conversation_id() + ) + # Parent-derived authz: the repository never filters reads by user_id. + assert repo.get_artifact(created["id"])["user_id"] == "owner" + + +class TestGetArtifactInParent: + def test_returns_when_parent_matches(self, pg_conn): + repo = _repo(pg_conn) + conv = _conversation_id() + created = repo.create_artifact("user-1", "document", conversation_id=conv) + fetched = repo.get_artifact_in_parent(created["id"], conversation_id=conv) + assert fetched is not None + assert fetched["id"] == created["id"] + + def test_returns_none_on_parent_mismatch(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create_artifact( + "user-1", "document", conversation_id=_conversation_id() + ) + # Correct id but a different conversation -> denied (no cross-parent leak). + assert repo.get_artifact_in_parent(created["id"], conversation_id=_conversation_id()) is None + + def test_requires_a_parent(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create_artifact( + "user-1", "document", conversation_id=_conversation_id() + ) + with pytest.raises(ValueError): + repo.get_artifact_in_parent(created["id"]) + + +class TestAppendVersion: + def test_increments_current_version(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create_artifact( + "user-1", "document", conversation_id=_conversation_id() + ) + assert created["current_version"] == 1 + + v2 = repo.append_version(created["id"], spec={"body": "edit"}, filename="d.md") + assert v2["version"] == 2 + + refreshed = repo.get_artifact(created["id"]) + assert refreshed["current_version"] == 2 + assert refreshed["updated_at"] >= created["updated_at"] + + def test_two_appends_yield_v2_then_v3(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create_artifact( + "user-1", "document", conversation_id=_conversation_id() + ) + assert repo.append_version(created["id"])["version"] == 2 + assert repo.append_version(created["id"])["version"] == 3 + assert repo.get_artifact(created["id"])["current_version"] == 3 + + def test_append_to_missing_artifact_raises(self, pg_conn): + repo = _repo(pg_conn) + with pytest.raises(ValueError): + repo.append_version(str(uuid.uuid4())) + + def test_duplicate_version_raises(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create_artifact( + "user-1", "document", conversation_id=_conversation_id() + ) + # Force a collision on (artifact_id, version) by re-inserting v1. + try: + with pg_conn.begin_nested(): + repo._insert_version( + artifact_id=created["id"], + version=1, + mime_type=None, + filename=None, + storage_path=None, + size=None, + sha256=None, + spec=None, + preview_text=None, + produced_by=None, + ) + except IntegrityError: + pass + else: + pytest.fail("expected IntegrityError on duplicate (artifact_id, version)") + + +class TestListArtifacts: + def test_filters_by_conversation(self, pg_conn): + repo = _repo(pg_conn) + conv_a = _conversation_id() + conv_b = _conversation_id() + repo.create_artifact("user-1", "document", conversation_id=conv_a) + repo.create_artifact("user-1", "document", conversation_id=conv_a) + repo.create_artifact("user-1", "document", conversation_id=conv_b) + + assert len(repo.list_artifacts(conversation_id=conv_a)) == 2 + assert len(repo.list_artifacts(conversation_id=conv_b)) == 1 + + def test_filters_by_workflow_run(self, pg_conn): + repo = _repo(pg_conn) + run = _conversation_id() + repo.create_artifact("user-1", "document", workflow_run_id=run) + repo.create_artifact("user-1", "document", conversation_id=_conversation_id()) + assert len(repo.list_artifacts(workflow_run_id=run)) == 1 + + def test_filters_by_user(self, pg_conn): + repo = _repo(pg_conn) + conv = _conversation_id() + repo.create_artifact("alice", "document", conversation_id=conv) + repo.create_artifact("bob", "document", conversation_id=conv) + assert len(repo.list_artifacts(conversation_id=conv, user_id="alice")) == 1 + + def test_requires_at_least_one_filter(self, pg_conn): + repo = _repo(pg_conn) + with pytest.raises(ValueError): + repo.list_artifacts() + + +class TestCascadeDelete: + def test_deleting_artifact_removes_versions(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create_artifact( + "user-1", "document", conversation_id=_conversation_id() + ) + repo.append_version(created["id"]) + assert len(repo.list_versions(created["id"])) == 2 + + pg_conn.execute( + text("DELETE FROM artifacts WHERE id = CAST(:id AS uuid)"), + {"id": created["id"]}, + ) + assert repo.get_artifact(created["id"]) is None + assert repo.list_versions(created["id"]) == [] From 922ed53a70b4f8e73e628cc21f6e078b334e02e5 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 24 Jun 2026 10:53:41 +0100 Subject: [PATCH 02/59] Add artifact REST endpoints (list, get, download, restore) Serve artifact metadata and bytes over HTTP with parent-derived authorization: conversation-parented artifacts inherit conversation access (owner, shared_with, or a public share token whose conversation matches the parent), workflow-run artifacts check run ownership, and access fails closed when the parent is missing or deleted. Adds list/get/versions/download/restore routes, an authenticated storage-agnostic download (sanitized Content-Disposition, 302 to a short-lived private S3 presigned URL when that strategy is configured), a generate_presigned_url primitive on the storage base and S3 backend, and generalizes the tools artifact endpoint for documents/files. Shared authorization helpers live in a dedicated module used by both surfaces. --- application/api/user/artifacts/__init__.py | 5 + application/api/user/artifacts/authz.py | 68 +++ application/api/user/artifacts/routes.py | 372 ++++++++++++ application/api/user/routes.py | 4 + application/api/user/tools/routes.py | 41 ++ application/storage/base.py | 17 + application/storage/s3.py | 10 + tests/api/user/test_artifacts_routes.py | 564 ++++++++++++++++++ .../user/test_tools_artifact_generalized.py | 107 ++++ 9 files changed, 1188 insertions(+) create mode 100644 application/api/user/artifacts/__init__.py create mode 100644 application/api/user/artifacts/authz.py create mode 100644 application/api/user/artifacts/routes.py create mode 100644 tests/api/user/test_artifacts_routes.py create mode 100644 tests/api/user/test_tools_artifact_generalized.py diff --git a/application/api/user/artifacts/__init__.py b/application/api/user/artifacts/__init__.py new file mode 100644 index 00000000..2c55b194 --- /dev/null +++ b/application/api/user/artifacts/__init__.py @@ -0,0 +1,5 @@ +"""Artifacts module.""" + +from .routes import artifacts_ns + +__all__ = ["artifacts_ns"] diff --git a/application/api/user/artifacts/authz.py b/application/api/user/artifacts/authz.py new file mode 100644 index 00000000..f267131e --- /dev/null +++ b/application/api/user/artifacts/authz.py @@ -0,0 +1,68 @@ +"""Parent-derived authorization helpers for artifact access.""" + +from __future__ import annotations + +from typing import Optional + +from flask import request + +from application.storage.db.repositories.agents import AgentsRepository +from application.storage.db.repositories.conversations import ConversationsRepository +from application.storage.db.repositories.shared_conversations import ( + SharedConversationsRepository, +) +from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository +from application.storage.db.session import db_readonly + + +def resolve_authenticated_user() -> Optional[str]: + """Resolve the caller to a user id via decoded JWT ``sub`` or api_key→owner.""" + decoded_token = getattr(request, "decoded_token", None) + if decoded_token: + return decoded_token.get("sub") + + api_key = request.args.get("api_key") or request.form.get("api_key") + if api_key: + with db_readonly() as conn: + agent = AgentsRepository(conn).find_by_key(api_key) + if agent: + return agent.get("user_id") + return None + + +def user_can_access_conversation( + conn, conversation_id: str, user_id: Optional[str], share_token: Optional[str] +) -> bool: + """Allow if the caller owns/shares the conversation, or holds a valid share token. + + Reuses ``ConversationsRepository.get`` (owner OR ``shared_with``) so artifact + access tracks message access; a publicly shared link inherits download access + via its share token (see ``SharedConversationsRepository.find_by_uuid``). + """ + if user_id: + if ConversationsRepository(conn).get(conversation_id, user_id) is not None: + return True + if share_token: + shared = SharedConversationsRepository(conn).find_by_uuid(share_token) + if shared and str(shared.get("conversation_id")) == str(conversation_id): + return True + return False + + +def authorize_artifact(conn, artifact: dict, user_id: Optional[str]) -> bool: + """Authorize an artifact by resolving its parent; missing parent fails closed.""" + conversation_id = artifact.get("conversation_id") + workflow_run_id = artifact.get("workflow_run_id") + share_token = request.args.get("share_token") + + if conversation_id is not None: + return user_can_access_conversation( + conn, str(conversation_id), user_id, share_token + ) + if workflow_run_id is not None: + if not user_id: + return False + 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 new file mode 100644 index 00000000..fee5c4d1 --- /dev/null +++ b/application/api/user/artifacts/routes.py @@ -0,0 +1,372 @@ +"""Artifact metadata and download routes (parent-derived authz).""" + +from __future__ import annotations + +import re +from typing import Optional + +from flask import current_app, jsonify, make_response, redirect, request +from flask_restx import Namespace, Resource + +from application.api import api +from application.api.user.artifacts.authz import ( + authorize_artifact, + resolve_authenticated_user, + user_can_access_conversation, +) +from application.core.settings import settings +from application.storage.db.base_repository import looks_like_uuid +from application.storage.db.repositories.artifacts import ArtifactsRepository +from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository +from application.storage.db.session import db_readonly, db_session +from application.storage.storage_creator import StorageCreator + +artifacts_ns = Namespace("artifacts", description="Artifact operations", path="/api") + +# Presigned-URL TTL for private S3 artifact downloads (seconds). +_PRESIGNED_URL_TTL = 300 + + +def _sanitize_header_filename(filename: Optional[str], fallback: str) -> str: + """Strip CRLF / quotes from a display filename for a Content-Disposition header.""" + if not filename: + return fallback + cleaned = re.sub(r'[\r\n"]', "", str(filename)).strip() + return cleaned or fallback + + +def _artifact_summary(artifact: dict) -> dict: + """Project an artifact identity row to its API metadata shape (owner id withheld).""" + return { + "id": str(artifact.get("id")), + "conversation_id": ( + str(artifact["conversation_id"]) + if artifact.get("conversation_id") is not None + else None + ), + "workflow_run_id": ( + str(artifact["workflow_run_id"]) + if artifact.get("workflow_run_id") is not None + else None + ), + "message_id": ( + str(artifact["message_id"]) + if artifact.get("message_id") is not None + else None + ), + "kind": artifact.get("kind"), + "title": artifact.get("title"), + "metadata": artifact.get("metadata"), + "current_version": artifact.get("current_version"), + "created_at": _iso(artifact.get("created_at")), + "updated_at": _iso(artifact.get("updated_at")), + } + + +def _version_summary(version: dict, *, include_spec: bool = False) -> dict: + """Project a version row to its API metadata shape (storage_path withheld).""" + out = { + "version": version.get("version"), + "mime_type": version.get("mime_type"), + "filename": version.get("filename"), + "size": version.get("size"), + "sha256": version.get("sha256"), + "preview_text": version.get("preview_text"), + "produced_by": version.get("produced_by"), + "created_at": _iso(version.get("created_at")), + } + if include_spec: + out["spec"] = version.get("spec") + return out + + +def _iso(value): + """ISO-format a datetime, passing through other values unchanged.""" + return value.isoformat() if hasattr(value, "isoformat") else value + + +@artifacts_ns.route("/artifacts") +class ListArtifacts(Resource): + @api.doc(description="List artifacts for a conversation, workflow run, or the caller") + def get(self): + user_id = resolve_authenticated_user() + conversation_id = request.args.get("conversation_id") + workflow_run_id = request.args.get("workflow_run_id") + share_token = request.args.get("share_token") + + if not user_id and not share_token: + return make_response( + jsonify({"success": False, "message": "Authentication required"}), 401 + ) + + # Gate UUID-shape before any CAST(:id AS uuid) reaches the repo, so a + # malformed id is rejected cleanly instead of poisoning the transaction. + if conversation_id and not looks_like_uuid(conversation_id): + return make_response( + jsonify({"success": False, "message": "Invalid conversation_id"}), 400 + ) + if workflow_run_id and not looks_like_uuid(workflow_run_id): + return make_response( + jsonify({"success": False, "message": "Invalid workflow_run_id"}), 400 + ) + + try: + with db_readonly() as conn: + if conversation_id: + if not user_can_access_conversation( + conn, conversation_id, user_id, share_token + ): + return make_response( + jsonify({"success": False, "message": "Forbidden"}), 403 + ) + rows = ArtifactsRepository(conn).list_artifacts( + conversation_id=conversation_id + ) + elif workflow_run_id: + run = WorkflowRunsRepository(conn).get(workflow_run_id) + if run is None or run.get("user_id") != user_id: + return make_response( + jsonify({"success": False, "message": "Forbidden"}), 403 + ) + rows = ArtifactsRepository(conn).list_artifacts( + workflow_run_id=workflow_run_id + ) + else: + if not user_id: + return make_response( + jsonify({"success": False, "message": "Authentication required"}), + 401, + ) + rows = ArtifactsRepository(conn).list_artifacts(user_id=user_id) + + return make_response( + jsonify( + {"success": True, "artifacts": [_artifact_summary(r) for r in rows]} + ), + 200, + ) + except Exception as err: + current_app.logger.error(f"Error listing artifacts: {err}", exc_info=True) + return make_response(jsonify({"success": False}), 400) + + +@artifacts_ns.route("/artifacts/") +class GetArtifact(Resource): + @api.doc(description="Get an artifact's metadata, version list, and current spec") + def get(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_readonly() 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 + ) + if not authorize_artifact(conn, artifact, user_id): + return make_response( + jsonify({"success": False, "message": "Forbidden"}), 403 + ) + versions = repo.list_versions(artifact_id) + current = repo.get_version(artifact_id, artifact.get("current_version")) + + payload = _artifact_summary(artifact) + payload["versions"] = [_version_summary(v) for v in versions] + payload["spec"] = current.get("spec") if current else None + return make_response(jsonify({"success": True, "artifact": payload}), 200) + except Exception as err: + current_app.logger.error(f"Error retrieving artifact: {err}", exc_info=True) + return make_response(jsonify({"success": False}), 400) + + +@artifacts_ns.route("/artifacts//versions/") +class GetArtifactVersion(Resource): + @api.doc(description="Get a single artifact version's metadata and spec") + def get(self, artifact_id: str, version: int): + 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_readonly() 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 + ) + if not authorize_artifact(conn, artifact, user_id): + return make_response( + jsonify({"success": False, "message": "Forbidden"}), 403 + ) + version_row = repo.get_version(artifact_id, version) + if version_row is None: + return make_response( + jsonify({"success": False, "message": "Version not found"}), 404 + ) + return make_response( + jsonify( + {"success": True, "version": _version_summary(version_row, include_spec=True)} + ), + 200, + ) + except Exception as err: + current_app.logger.error( + f"Error retrieving artifact version: {err}", exc_info=True + ) + return make_response(jsonify({"success": False}), 400) + + +@artifacts_ns.route("/artifacts//download") +class DownloadArtifact(Resource): + @api.doc(description="Download an artifact's bytes (302 to a presigned URL on S3)") + def get(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() + version_arg = request.args.get("version") + try: + with db_readonly() 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 + ) + if not authorize_artifact(conn, artifact, user_id): + return make_response( + jsonify({"success": False, "message": "Forbidden"}), 403 + ) + version = artifact.get("current_version") + if version_arg is not None: + try: + version = int(version_arg) + except ValueError: + return make_response( + jsonify({"success": False, "message": "Invalid version"}), 400 + ) + version_row = repo.get_version(artifact_id, version) + + if version_row is None: + return make_response( + jsonify({"success": False, "message": "Version not found"}), 404 + ) + # The object key is derived only from the stored path, never client input. + storage_path = version_row.get("storage_path") + if not storage_path: + return make_response( + jsonify({"success": False, "message": "No file for this version"}), 404 + ) + + filename = _sanitize_header_filename( + version_row.get("filename"), f"artifact-{artifact_id}" + ) + mime_type = version_row.get("mime_type") or "application/octet-stream" + storage = StorageCreator.get_storage() + + # With URL_STRATEGY=="s3" the contract is to hand back a presigned + # URL. If the active backend can't mint one, that's a config error: + # surface a 500 rather than silently proxying bytes from a backend + # the operator expected to be off the hot path. + if getattr(settings, "URL_STRATEGY", "backend") == "s3": + try: + url = storage.generate_presigned_url( + storage_path, expires_in=_PRESIGNED_URL_TTL + ) + except NotImplementedError: + current_app.logger.error( + "URL_STRATEGY=s3 but %s cannot mint presigned URLs", + type(storage).__name__, + ) + return make_response( + jsonify( + {"success": False, "message": "Storage misconfigured"} + ), + 500, + ) + return redirect(url, code=302) + + 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}"' + ) + return response + except FileNotFoundError: + return make_response( + jsonify({"success": False, "message": "File not found"}), 404 + ) + except Exception as err: + current_app.logger.error( + f"Error downloading artifact: {err}", exc_info=True + ) + return make_response(jsonify({"success": False}), 400) + + +@artifacts_ns.route("/artifacts//restore") +class RestoreArtifact(Resource): + @api.doc(description="Restore a prior version by appending it as the new current version") + def post(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() + data = request.get_json(silent=True) or {} + target_version = data.get("version") + if target_version is None: + return make_response( + jsonify({"success": False, "message": "Missing version"}), 400 + ) + try: + target_version = int(target_version) + except (ValueError, TypeError): + return make_response( + jsonify({"success": False, "message": "Invalid version"}), 400 + ) + + 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 + ) + if not authorize_artifact(conn, artifact, user_id): + return make_response( + jsonify({"success": False, "message": "Forbidden"}), 403 + ) + source = repo.get_version(artifact_id, target_version) + if source is None: + return make_response( + jsonify({"success": False, "message": "Version not found"}), 404 + ) + new_version = repo.append_version( + artifact_id, + mime_type=source.get("mime_type"), + filename=source.get("filename"), + storage_path=source.get("storage_path"), + size=source.get("size"), + sha256=source.get("sha256"), + spec=source.get("spec"), + preview_text=source.get("preview_text"), + produced_by=source.get("produced_by"), + ) + return make_response( + jsonify( + {"success": True, "version": _version_summary(new_version, include_spec=True)} + ), + 200, + ) + except Exception as err: + current_app.logger.error(f"Error restoring artifact: {err}", exc_info=True) + return make_response(jsonify({"success": False}), 400) diff --git a/application/api/user/routes.py b/application/api/user/routes.py index 2d8ed6c7..9464bef9 100644 --- a/application/api/user/routes.py +++ b/application/api/user/routes.py @@ -13,6 +13,7 @@ from .agents import ( agents_webhooks_ns, ) from .analytics import analytics_ns +from .artifacts import artifacts_ns from .attachments import attachments_ns from .conversations import conversations_ns from .me import me_ns @@ -31,6 +32,9 @@ user = Blueprint("user", __name__) # Analytics api.add_namespace(analytics_ns) +# Artifacts +api.add_namespace(artifacts_ns) + # Attachments api.add_namespace(attachments_ns) diff --git a/application/api/user/tools/routes.py b/application/api/user/tools/routes.py index 7897e2a6..5caa006f 100644 --- a/application/api/user/tools/routes.py +++ b/application/api/user/tools/routes.py @@ -15,9 +15,12 @@ from application.agents.default_tools import ( from application.agents.tools.spec_parser import parse_spec from application.agents.tools.tool_manager import ToolManager from application.api import api +from application.api.user.artifacts.authz import authorize_artifact from application.api.user.team_sharing import effective_write_owner, visible_with_access from application.core.url_validation import SSRFError, validate_url from application.security.encryption import decrypt_credentials, encrypt_credentials +from application.storage.db.base_repository import looks_like_uuid +from application.storage.db.repositories.artifacts import ArtifactsRepository from application.storage.db.repositories.notes import NotesRepository from application.storage.db.repositories.todos import TodosRepository from application.storage.db.repositories.user_tools import UserToolsRepository @@ -934,6 +937,44 @@ class GetArtifact(Resource): return make_response( jsonify({"success": True, "artifact": artifact}), 200 ) + + # Generalized document/file artifacts live in the + # ``artifacts`` store, authorized by their parent (not user_id). + # Shape-gate the UUID lookup so a non-UUID id (legacy note/todo + # ObjectId already handled above) never reaches the CAST that + # would poison this read-only transaction. + artifact_doc = None + if looks_like_uuid(artifact_id): + artifacts_repo = ArtifactsRepository(conn) + artifact_doc = artifacts_repo.get_artifact(artifact_id) + if artifact_doc and authorize_artifact(conn, artifact_doc, user_id): + current = artifacts_repo.get_version( + artifact_id, artifact_doc.get("current_version") + ) + artifact = { + "artifact_type": ( + "document" + if artifact_doc.get("kind") != "file" + else "file" + ), + "data": { + "id": str(artifact_doc.get("id")), + "kind": artifact_doc.get("kind"), + "title": artifact_doc.get("title"), + "current_version": artifact_doc.get("current_version"), + "mime_type": current.get("mime_type") if current else None, + "filename": current.get("filename") if current else None, + "size": current.get("size") if current else None, + "download_url": ( + f"/api/artifacts/{artifact_id}/download" + if current and current.get("storage_path") + else None + ), + }, + } + return make_response( + jsonify({"success": True, "artifact": artifact}), 200 + ) except Exception as err: current_app.logger.error( f"Error retrieving artifact: {err}", exc_info=True diff --git a/application/storage/base.py b/application/storage/base.py index 79bd7348..1dae6506 100644 --- a/application/storage/base.py +++ b/application/storage/base.py @@ -37,6 +37,23 @@ class BaseStorage(ABC): """ pass + def generate_presigned_url(self, path: str, expires_in: int = 300) -> str: + """Return a short-lived presigned download URL; not all backends support it. + + Args: + path: Path to the file + expires_in: TTL of the signed URL in seconds + + Returns: + str: A presigned URL granting time-limited read access + + Raises: + NotImplementedError: If the backend cannot mint presigned URLs. + """ + raise NotImplementedError( + f"{type(self).__name__} does not support presigned URLs" + ) + @abstractmethod def process_file(self, path: str, processor_func: Callable, **kwargs): """ diff --git a/application/storage/s3.py b/application/storage/s3.py index c2cb64b4..9c16fab2 100644 --- a/application/storage/s3.py +++ b/application/storage/s3.py @@ -121,6 +121,16 @@ class S3Storage(BaseStorage): file_obj.seek(0) return file_obj + def generate_presigned_url(self, path: str, expires_in: int = 300) -> str: + """Return a short-lived presigned GET URL for a private object (TTL <= 1h).""" + path = self._validate_path(path) + expires_in = min(expires_in, 3600) + return self.s3.generate_presigned_url( + "get_object", + Params={"Bucket": self.bucket_name, "Key": path}, + ExpiresIn=expires_in, + ) + def delete_file(self, path: str) -> bool: """Delete a file from S3 storage.""" path = self._validate_path(path) diff --git a/tests/api/user/test_artifacts_routes.py b/tests/api/user/test_artifacts_routes.py new file mode 100644 index 00000000..c2641857 --- /dev/null +++ b/tests/api/user/test_artifacts_routes.py @@ -0,0 +1,564 @@ +"""Route-layer tests for the artifacts API (parent-derived authz).""" + +from __future__ import annotations + +import io +import uuid +from contextlib import contextmanager + +import pytest +from flask import request +from sqlalchemy import text + +from application.storage.db.repositories.artifacts import ArtifactsRepository +from application.storage.db.repositories.conversations import ConversationsRepository +from application.storage.db.repositories.shared_conversations import ( + SharedConversationsRepository, +) +from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository + + +OWNER = "artifact_owner" +STRANGER = "artifact_stranger" +SHARED_USER = "artifact_shared_user" + + +@pytest.fixture +def _patch_db(pg_conn, monkeypatch): + """Redirect the routes' ``db_readonly`` / ``db_session`` to the test conn.""" + + @contextmanager + def _use_conn(): + yield pg_conn + + monkeypatch.setattr("application.api.user.artifacts.routes.db_readonly", _use_conn) + monkeypatch.setattr("application.api.user.artifacts.routes.db_session", _use_conn) + return pg_conn + + +@pytest.fixture +def token_owner(): + return {"sub": OWNER, "email": "owner@example.com"} + + +def _make_conversation(conn, user_id=OWNER): + return ConversationsRepository(conn).create(user_id, name="conv") + + +def _make_workflow(conn, user_id=OWNER): + res = conn.execute( + text("INSERT INTO workflows (user_id, name) VALUES (:u, 'wf') RETURNING id"), + {"u": user_id}, + ) + return str(res.fetchone()[0]) + + +def _make_artifact(conn, **kwargs): + return ArtifactsRepository(conn).create_artifact( + kwargs.pop("user_id", OWNER), + kwargs.pop("kind", "document"), + **kwargs, + ) + + +def _call(flask_app, resource_cls, *args, token=None, query=None, json_body=None, method="get"): + with flask_app.app_context(): + with flask_app.test_request_context(query_string=query or {}, json=json_body): + request.decoded_token = token + resource = resource_cls() + return getattr(resource, method)(*args) + + +# --------------------------------------------------------------------------- +# List +# --------------------------------------------------------------------------- +@pytest.mark.unit +class TestListArtifacts: + def test_owner_lists_conversation_artifacts( + self, _patch_db, flask_app, token_owner + ): + from application.api.user.artifacts.routes import ListArtifacts + + conv = _make_conversation(_patch_db) + conv_id = str(conv["id"]) + _make_artifact(_patch_db, conversation_id=conv_id, title="a1") + _make_artifact(_patch_db, conversation_id=conv_id, title="a2") + + resp = _call( + flask_app, ListArtifacts, token=token_owner, + query={"conversation_id": conv_id}, + ) + assert resp.status_code == 200 + assert len(resp.json["artifacts"]) == 2 + + def test_stranger_denied_conversation_list( + self, _patch_db, flask_app + ): + from application.api.user.artifacts.routes import ListArtifacts + + conv = _make_conversation(_patch_db) + _make_artifact(_patch_db, conversation_id=str(conv["id"])) + + resp = _call( + flask_app, ListArtifacts, + token={"sub": STRANGER}, + query={"conversation_id": str(conv["id"])}, + ) + assert resp.status_code == 403 + + def test_no_filter_scopes_to_user(self, _patch_db, flask_app, token_owner): + from application.api.user.artifacts.routes import ListArtifacts + + owner_conv = _make_conversation(_patch_db) + owned = _make_artifact( + _patch_db, conversation_id=str(owner_conv["id"]), user_id=OWNER + ) + stranger = _make_artifact( + _patch_db, conversation_id=str(_make_conversation(_patch_db)["id"]), + user_id=STRANGER, + ) + + resp = _call(flask_app, ListArtifacts, token=token_owner) + assert resp.status_code == 200 + # The owner id is accounting-only and must not leak in the summary. + assert all("user_id" not in a for a in resp.json["artifacts"]) + returned_ids = {a["id"] for a in resp.json["artifacts"]} + assert str(owned["id"]) in returned_ids + assert str(stranger["id"]) not in returned_ids + + def test_unauthenticated_401(self, _patch_db, flask_app): + from application.api.user.artifacts.routes import ListArtifacts + + resp = _call(flask_app, ListArtifacts, token=None) + assert resp.status_code == 401 + + def test_non_uuid_conversation_id_400(self, _patch_db, flask_app, token_owner): + from application.api.user.artifacts.routes import ListArtifacts + + # A malformed id must be rejected before reaching CAST(:id AS uuid), + # which would otherwise raise a DataError and poison the transaction. + resp = _call( + flask_app, ListArtifacts, token=token_owner, + query={"conversation_id": "not-a-uuid"}, + ) + assert resp.status_code == 400 + + def test_non_uuid_workflow_run_id_400(self, _patch_db, flask_app, token_owner): + from application.api.user.artifacts.routes import ListArtifacts + + resp = _call( + flask_app, ListArtifacts, token=token_owner, + query={"workflow_run_id": "1234"}, + ) + assert resp.status_code == 400 + + +# --------------------------------------------------------------------------- +# Get + versions +# --------------------------------------------------------------------------- +@pytest.mark.unit +class TestGetArtifact: + def test_owner_gets_artifact_with_versions( + self, _patch_db, flask_app, token_owner + ): + from application.api.user.artifacts.routes import GetArtifact + + 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"}) + + resp = _call(flask_app, GetArtifact, art["id"], token=token_owner) + assert resp.status_code == 200 + assert resp.json["artifact"]["current_version"] == 2 + assert len(resp.json["artifact"]["versions"]) == 2 + assert resp.json["artifact"]["spec"] == {"body": "v2"} + + def test_stranger_denied_403(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"])) + + resp = _call(flask_app, GetArtifact, art["id"], token={"sub": STRANGER}) + assert resp.status_code == 403 + + def test_missing_parent_fails_closed(self, _patch_db, flask_app, token_owner): + from application.api.user.artifacts.routes import GetArtifact + from sqlalchemy import text + + conv = _make_conversation(_patch_db) + art = _make_artifact(_patch_db, conversation_id=str(conv["id"])) + # Delete the parent conversation -> the artifact must no longer be reachable. + _patch_db.execute( + text("DELETE FROM conversations WHERE id = CAST(:id AS uuid)"), + {"id": str(conv["id"])}, + ) + + resp = _call(flask_app, GetArtifact, art["id"], token=token_owner) + assert resp.status_code == 403 + + 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) + assert resp.status_code == 404 + + def test_workflow_run_owner_access(self, _patch_db, flask_app, token_owner): + from application.api.user.artifacts.routes import GetArtifact + + wf_id = _make_workflow(_patch_db, OWNER) + run = WorkflowRunsRepository(_patch_db).create(wf_id, OWNER, "completed") + art = _make_artifact(_patch_db, workflow_run_id=str(run["id"])) + + resp = _call(flask_app, GetArtifact, art["id"], token=token_owner) + assert resp.status_code == 200 + + resp_other = _call( + flask_app, GetArtifact, art["id"], token={"sub": STRANGER} + ) + assert resp_other.status_code == 403 + + +@pytest.mark.unit +class TestGetArtifactVersion: + def test_version_returns_spec(self, _patch_db, flask_app, token_owner): + from application.api.user.artifacts.routes import GetArtifactVersion + + conv = _make_conversation(_patch_db) + art = _make_artifact( + _patch_db, conversation_id=str(conv["id"]), spec={"body": "one"} + ) + resp = _call(flask_app, GetArtifactVersion, art["id"], 1, token=token_owner) + assert resp.status_code == 200 + assert resp.json["version"]["spec"] == {"body": "one"} + + def test_missing_version_404(self, _patch_db, flask_app, token_owner): + from application.api.user.artifacts.routes import GetArtifactVersion + + conv = _make_conversation(_patch_db) + art = _make_artifact(_patch_db, conversation_id=str(conv["id"])) + resp = _call(flask_app, GetArtifactVersion, art["id"], 99, token=token_owner) + assert resp.status_code == 404 + + +# --------------------------------------------------------------------------- +# Shared-conversation inheritance +# --------------------------------------------------------------------------- +@pytest.mark.unit +class TestSharedAccess: + def test_shared_with_user_can_get(self, _patch_db, flask_app): + from application.api.user.artifacts.routes import GetArtifact + + 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"])) + + resp = _call(flask_app, GetArtifact, art["id"], token={"sub": SHARED_USER}) + assert resp.status_code == 200 + + def test_share_token_holder_can_download(self, _patch_db, flask_app, monkeypatch): + from application.api.user.artifacts.routes import DownloadArtifact + + conv = _make_conversation(_patch_db) + art = _make_artifact( + _patch_db, + conversation_id=str(conv["id"]), + filename="report.pdf", + mime_type="application/pdf", + storage_path="inputs/owner/artifacts/x/v1/report.pdf", + ) + share = SharedConversationsRepository(_patch_db).create( + str(conv["id"]), OWNER + ) + + storage = _FakeStorage(b"PDFDATA") + 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, + ) + + # Anonymous link holder (no JWT) supplies the share token. + resp = _call( + flask_app, DownloadArtifact, art["id"], token=None, + query={"share_token": str(share["uuid"])}, + ) + assert resp.status_code == 200 + assert resp.data == b"PDFDATA" + + +# --------------------------------------------------------------------------- +# Download +# --------------------------------------------------------------------------- +class _FakeStorage: + def __init__(self, data: bytes): + self._data = data + + def get_file(self, path): + return io.BytesIO(self._data) + + def generate_presigned_url(self, path, expires_in=300): + return f"https://signed.example/{path}?exp={expires_in}" + + +@pytest.mark.unit +class TestDownloadArtifact: + def _seed(self, conn, **over): + conv = _make_conversation(conn) + defaults = dict( + conversation_id=str(conv["id"]), + filename="deck.pptx", + mime_type="application/vnd.ms-powerpoint", + storage_path="inputs/owner/artifacts/x/v1/deck.pptx", + ) + defaults.update(over) + return _make_artifact(conn, **defaults) + + def test_local_streams_bytes_with_content_disposition( + self, _patch_db, flask_app, token_owner, monkeypatch + ): + from application.api.user.artifacts.routes import DownloadArtifact + + art = self._seed(_patch_db) + storage = _FakeStorage(b"BINARY") + 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.status_code == 200 + assert resp.data == b"BINARY" + assert resp.headers["Content-Disposition"] == 'attachment; filename="deck.pptx"' + assert resp.headers["Content-Type"] == "application/vnd.ms-powerpoint" + + def test_s3_strategy_redirects_to_presigned( + self, _patch_db, flask_app, token_owner, monkeypatch + ): + from application.api.user.artifacts.routes import DownloadArtifact + + art = self._seed(_patch_db) + storage = _FakeStorage(b"unused") + monkeypatch.setattr( + "application.api.user.artifacts.routes.StorageCreator.get_storage", + lambda: storage, + ) + monkeypatch.setattr( + "application.api.user.artifacts.routes.settings.URL_STRATEGY", + "s3", + raising=False, + ) + + resp = _call(flask_app, DownloadArtifact, art["id"], token=token_owner) + assert resp.status_code == 302 + assert resp.headers["Location"].startswith("https://signed.example/") + + def test_stranger_denied(self, _patch_db, flask_app, monkeypatch): + from application.api.user.artifacts.routes import DownloadArtifact + + art = self._seed(_patch_db) + resp = _call(flask_app, DownloadArtifact, art["id"], token={"sub": STRANGER}) + assert resp.status_code == 403 + + def test_s3_strategy_misconfigured_backend_500( + self, _patch_db, flask_app, token_owner, monkeypatch + ): + from application.api.user.artifacts.routes import DownloadArtifact + + art = self._seed(_patch_db) + + class _NoPresignStorage: + def get_file(self, path): + return io.BytesIO(b"unused") + + def generate_presigned_url(self, path, expires_in=300): + raise NotImplementedError("backend cannot mint presigned URLs") + + monkeypatch.setattr( + "application.api.user.artifacts.routes.StorageCreator.get_storage", + lambda: _NoPresignStorage(), + ) + monkeypatch.setattr( + "application.api.user.artifacts.routes.settings.URL_STRATEGY", + "s3", + raising=False, + ) + + resp = _call(flask_app, DownloadArtifact, art["id"], token=token_owner) + assert resp.status_code == 500 + + def test_crlf_filename_sanitized_in_header( + self, _patch_db, flask_app, token_owner, monkeypatch + ): + from application.api.user.artifacts.routes import DownloadArtifact + + art = self._seed(_patch_db, filename='a"\r\nInjected: x.txt') + storage = _FakeStorage(b"X") + 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) + disposition = resp.headers["Content-Disposition"] + assert "\r" not in disposition and "\n" not in disposition + assert disposition == 'attachment; filename="aInjected: x.txt"' + + +# --------------------------------------------------------------------------- +# Restore +# --------------------------------------------------------------------------- +@pytest.mark.unit +class TestRestoreArtifact: + def test_restore_appends_new_version(self, _patch_db, flask_app, token_owner): + 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"}) + + resp = _call( + flask_app, RestoreArtifact, art["id"], token=token_owner, + json_body={"version": 1}, method="post", + ) + assert resp.status_code == 200 + assert resp.json["version"]["version"] == 3 + assert resp.json["version"]["spec"] == {"body": "v1"} + assert ArtifactsRepository(_patch_db).get_artifact(art["id"])["current_version"] == 3 + + def test_restore_missing_version_field_400( + self, _patch_db, flask_app, token_owner + ): + from application.api.user.artifacts.routes import RestoreArtifact + + conv = _make_conversation(_patch_db) + art = _make_artifact(_patch_db, conversation_id=str(conv["id"])) + resp = _call( + flask_app, RestoreArtifact, art["id"], token=token_owner, + json_body={}, method="post", + ) + assert resp.status_code == 400 + + def test_restore_stranger_denied(self, _patch_db, flask_app): + from application.api.user.artifacts.routes import RestoreArtifact + + conv = _make_conversation(_patch_db) + art = _make_artifact(_patch_db, conversation_id=str(conv["id"])) + resp = _call( + flask_app, RestoreArtifact, art["id"], token={"sub": STRANGER}, + json_body={"version": 1}, method="post", + ) + assert resp.status_code == 403 + + +# --------------------------------------------------------------------------- +# Malformed artifact_id path segment -> 404 (no CAST(:id AS uuid) txn poison) +# --------------------------------------------------------------------------- +@pytest.mark.unit +class TestMalformedArtifactId: + @pytest.mark.parametrize( + "resource_name, extra_args, method, json_body", + [ + ("GetArtifact", (), "get", None), + ("GetArtifactVersion", (1,), "get", None), + ("DownloadArtifact", (), "get", None), + ("RestoreArtifact", (), "post", {"version": 1}), + ], + ) + def test_non_uuid_id_returns_404( + self, _patch_db, flask_app, token_owner, resource_name, + extra_args, method, json_body, + ): + import application.api.user.artifacts.routes as routes_mod + + resource_cls = getattr(routes_mod, resource_name) + resp = _call( + flask_app, resource_cls, "legacy-mongo-objectid-aabbcc", *extra_args, + token=token_owner, method=method, json_body=json_body, + ) + assert resp.status_code == 404 + + +# --------------------------------------------------------------------------- +# api_key -> agent-owner principal resolution +# --------------------------------------------------------------------------- +@pytest.mark.unit +class TestApiKeyPrincipal: + def test_api_key_owner_can_get_owned_artifact( + self, _patch_db, flask_app, monkeypatch + ): + from application.api.user.artifacts import authz + from application.api.user.artifacts.routes import GetArtifact + from application.storage.db.repositories.agents import AgentsRepository + + conv = _make_conversation(_patch_db, user_id=OWNER) + art = _make_artifact(_patch_db, conversation_id=str(conv["id"])) + + # The api_key path opens its own readonly conn inside authz; point it at + # the test conn and resolve the key to the artifact's owning agent. + @contextmanager + def _use_conn(): + yield _patch_db + + monkeypatch.setattr(authz, "db_readonly", _use_conn) + monkeypatch.setattr( + AgentsRepository, "find_by_key", + lambda self, key: {"user_id": OWNER} if key == "secret-key" else None, + ) + + # No JWT -> principal must be resolved from the api_key query param. + resp = _call( + 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"]) + + def test_api_key_resolving_to_stranger_denied( + self, _patch_db, flask_app, monkeypatch + ): + from application.api.user.artifacts import authz + from application.api.user.artifacts.routes import GetArtifact + from application.storage.db.repositories.agents import AgentsRepository + + conv = _make_conversation(_patch_db, user_id=OWNER) + art = _make_artifact(_patch_db, conversation_id=str(conv["id"])) + + @contextmanager + def _use_conn(): + yield _patch_db + + monkeypatch.setattr(authz, "db_readonly", _use_conn) + monkeypatch.setattr( + AgentsRepository, "find_by_key", + lambda self, key: {"user_id": STRANGER}, + ) + + resp = _call( + flask_app, GetArtifact, art["id"], token=None, + query={"api_key": "stranger-key"}, + ) + assert resp.status_code == 403 + + +# A full-app integration test (namespace registration + auth before_request) +# is covered by the e2e suite under tests/e2e; the unit tests above exercise +# the malformed-id gate and the api_key principal path directly against the +# Resource handlers without building the whole app. diff --git a/tests/api/user/test_tools_artifact_generalized.py b/tests/api/user/test_tools_artifact_generalized.py new file mode 100644 index 00000000..9f13bc4d --- /dev/null +++ b/tests/api/user/test_tools_artifact_generalized.py @@ -0,0 +1,107 @@ +"""Tests for the generalized GET /api/tools/artifact/ document branch.""" + +from __future__ import annotations + +import uuid +from contextlib import contextmanager + +import pytest +from flask import request + +from application.storage.db.repositories.artifacts import ArtifactsRepository +from application.storage.db.repositories.conversations import ConversationsRepository +from application.storage.db.repositories.notes import NotesRepository +from application.storage.db.repositories.user_tools import UserToolsRepository + + +@pytest.fixture +def _patch_db(pg_conn, monkeypatch): + """Point both the tools and artifacts route DB helpers at the test conn.""" + + @contextmanager + def _use_conn(): + yield pg_conn + + monkeypatch.setattr("application.api.user.tools.routes.db_readonly", _use_conn) + monkeypatch.setattr("application.api.user.artifacts.routes.db_readonly", _use_conn) + return pg_conn + + +def _get(flask_app, artifact_id, token): + from application.api.user.tools.routes import GetArtifact + + with flask_app.app_context(): + with flask_app.test_request_context(): + request.decoded_token = token + return GetArtifact().get(artifact_id) + + +@pytest.mark.unit +class TestGeneralizedToolsArtifact: + def test_document_branch_returns_metadata( + self, _patch_db, flask_app, decoded_token + ): + conv = ConversationsRepository(_patch_db).create( + decoded_token["sub"], name="conv" + ) + art = ArtifactsRepository(_patch_db).create_artifact( + decoded_token["sub"], + "document", + conversation_id=str(conv["id"]), + title="Report", + filename="report.pdf", + mime_type="application/pdf", + storage_path="inputs/u/artifacts/x/v1/report.pdf", + ) + + resp = _get(flask_app, art["id"], decoded_token) + assert resp.status_code == 200 + body = resp.json["artifact"] + assert body["artifact_type"] == "document" + assert body["data"]["title"] == "Report" + assert body["data"]["filename"] == "report.pdf" + assert body["data"]["download_url"] == f"/api/artifacts/{art['id']}/download" + + def test_file_kind_maps_to_file_type(self, _patch_db, flask_app, decoded_token): + conv = ConversationsRepository(_patch_db).create( + decoded_token["sub"], name="conv" + ) + art = ArtifactsRepository(_patch_db).create_artifact( + decoded_token["sub"], + "file", + conversation_id=str(conv["id"]), + storage_path="inputs/u/artifacts/y/v1/data.bin", + ) + resp = _get(flask_app, art["id"], decoded_token) + assert resp.status_code == 200 + assert resp.json["artifact"]["artifact_type"] == "file" + + def test_stranger_document_not_found(self, _patch_db, flask_app, decoded_token): + conv = ConversationsRepository(_patch_db).create("other_owner", name="conv") + art = ArtifactsRepository(_patch_db).create_artifact( + "other_owner", + "document", + conversation_id=str(conv["id"]), + storage_path="inputs/o/artifacts/z/v1/x.pdf", + ) + # Parent-derived authz denies a stranger -> falls through to 404. + resp = _get(flask_app, art["id"], decoded_token) + assert resp.status_code == 404 + + def test_note_branch_still_works(self, _patch_db, flask_app, decoded_token): + tool = UserToolsRepository(_patch_db).create( + user_id=decoded_token["sub"], name="notes_tool" + ) + note = NotesRepository(_patch_db).upsert( + user_id=decoded_token["sub"], + tool_id=str(tool["id"]), + title="t", + content="a\nb", + ) + resp = _get(flask_app, str(note["id"]), decoded_token) + assert resp.status_code == 200 + assert resp.json["artifact"]["artifact_type"] == "note" + + def test_unknown_id_404(self, _patch_db, flask_app, decoded_token): + resp = _get(flask_app, str(uuid.uuid4()), decoded_token) + assert resp.status_code == 404 From 3b795dbb83ca3b96a106828b8fabb63560dbde6e Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 24 Jun 2026 11:26:23 +0100 Subject: [PATCH 03/59] Add semi-persistent code-execution sandbox (Jupyter Kernel Gateway runner) Introduce a pluggable CodeSandbox abstraction with a SandboxManager and a Jupyter Kernel Gateway backend: the app is a client of a single always-on runner that executes code in stateful in-process kernels (no child-container spawning, no docker socket). Sessions bind to a conversation or workflow run with an agent-selectable TTL clamped by a global cap; execution enforces a wall-clock deadline with interrupt-on-timeout and capped output, and file transfer is workspace-contained with size and integrity checks. Adds a docsgpt-sandbox docker-compose service (resource-capped, read-only, internal network) plus settings, with a real local-gateway integration test. --- application/core/settings.py | 19 + application/requirements.txt | 3 +- application/sandbox/__init__.py | 1 + application/sandbox/base.py | 75 +++ application/sandbox/jupyter_gateway.py | 454 ++++++++++++++++++ application/sandbox/manager.py | 135 ++++++ application/sandbox/sandbox_creator.py | 56 +++ deployment/docker-compose.yaml | 33 ++ deployment/sandbox/Dockerfile | 48 ++ deployment/sandbox/README.md | 50 ++ tests/requirements.txt | 6 + tests/sandbox/__init__.py | 0 .../test_jupyter_gateway_integration.py | 292 +++++++++++ tests/sandbox/test_sandbox_manager.py | 161 +++++++ 14 files changed, 1332 insertions(+), 1 deletion(-) create mode 100644 application/sandbox/__init__.py create mode 100644 application/sandbox/base.py create mode 100644 application/sandbox/jupyter_gateway.py create mode 100644 application/sandbox/manager.py create mode 100644 application/sandbox/sandbox_creator.py create mode 100644 deployment/sandbox/Dockerfile create mode 100644 deployment/sandbox/README.md create mode 100644 tests/sandbox/__init__.py create mode 100644 tests/sandbox/test_jupyter_gateway_integration.py create mode 100644 tests/sandbox/test_sandbox_manager.py diff --git a/application/core/settings.py b/application/core/settings.py index a5d5b8ec..1350262e 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -324,6 +324,25 @@ class Settings(BaseSettings): SCHEDULE_ONCE_MAX_HORIZON: int = 31_536_000 SCHEDULE_RUN_OUTPUT_RETENTION_DAYS: int = 90 + # Code-execution sandbox (see artifacts-code-execution-spec.md §4 C2). + # The app is a CLIENT of an always-on runner; defaults are safe so app + # import never fails when the sandbox is unconfigured. + SANDBOX_BACKEND: str = "jupyter" # "jupyter" (self-host) — "daytona" reserved + # URL of the Jupyter Kernel Gateway runner (the docsgpt-sandbox service). + SANDBOX_GATEWAY_URL: str = "http://localhost:8888" + SANDBOX_GATEWAY_AUTH_TOKEN: Optional[str] = None # gateway auth token, if set + SANDBOX_KERNEL_NAME: str = "python3" # kernelspec name to launch per session + SANDBOX_MAX_TTL: int = 1200 # hard cap (s) on agent-selectable keep-alive TTL + SANDBOX_EXEC_TIMEOUT: int = 60 # default wall-clock cap (s) per exec call + SANDBOX_HTTP_TIMEOUT: int = 10 # fixed cap (s) for REST control calls (create/delete/alive/interrupt) + SANDBOX_MAX_OUTPUT_BYTES: int = 8 * 1024 * 1024 # cap on buffered stdout+stderr per exec + SANDBOX_MAX_FILE_BYTES: int = 10 * 1024 * 1024 # cap on get_file size routed through stdout + # 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. + SANDBOX_MEMORY: str = "1g" # docker mem_limit for the runner container + SANDBOX_CPUS: str = "1.0" # docker cpu quota for the runner container + @field_validator("POSTGRES_URI", mode="before") @classmethod def _normalize_postgres_uri_validator(cls, v): diff --git a/application/requirements.txt b/application/requirements.txt index 822259c6..9281e6bd 100644 --- a/application/requirements.txt +++ b/application/requirements.txt @@ -110,4 +110,5 @@ werkzeug>=3.1.0 yarl==1.23.0 markdownify==1.2.2 tldextract==5.3.1 -websockets==16.0 \ No newline at end of file +websockets==16.0 +websocket-client==1.9.0 \ No newline at end of file diff --git a/application/sandbox/__init__.py b/application/sandbox/__init__.py new file mode 100644 index 00000000..7f061acd --- /dev/null +++ b/application/sandbox/__init__.py @@ -0,0 +1 @@ +"""Pluggable code-execution sandbox abstraction for running untrusted LLM code.""" diff --git a/application/sandbox/base.py b/application/sandbox/base.py new file mode 100644 index 00000000..ff52911d --- /dev/null +++ b/application/sandbox/base.py @@ -0,0 +1,75 @@ +"""Backend-agnostic code-execution sandbox interface and result types.""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + + +@dataclass +class DisplayData: + """A single rich (non-stream) kernel output keyed by MIME bundle.""" + + data: Dict[str, Any] = field(default_factory=dict) + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class Plot: + """An image/plot output captured from execution, base64-encoded by format.""" + + format: str + content_base64: str + + +@dataclass +class ExecResult: + """Outcome of one ``exec`` call against a sandbox session.""" + + status: str = "ok" # "ok" | "error" + stdout: str = "" + stderr: str = "" + exit_code: int = 0 + execution_count: Optional[int] = None + error_name: Optional[str] = None + error_value: Optional[str] = None + traceback: List[str] = field(default_factory=list) + results: List[DisplayData] = field(default_factory=list) + display_data: List[DisplayData] = field(default_factory=list) + plots: List[Plot] = field(default_factory=list) + + @property + def ok(self) -> bool: + """True when the execution completed without raising.""" + return self.status == "ok" + + +class CodeSandbox(ABC): + """Common interface every sandbox backend (Jupyter, Daytona, ...) implements.""" + + @abstractmethod + def open(self, session_id: str) -> str: + """Create the underlying runtime for ``session_id`` and return its handle id.""" + + @abstractmethod + def attach(self, session_id: str) -> str: + """Reattach to ``session_id``'s runtime; MAY return a cold kernel (state not guaranteed).""" + + @abstractmethod + def close(self, session_id: str) -> None: + """Tear down the runtime bound to ``session_id``, freeing all resources.""" + + @abstractmethod + def exec(self, session_id: str, code: str, timeout: Optional[float] = None) -> ExecResult: + """Execute ``code`` in ``session_id``'s stateful runtime and return its result.""" + + @abstractmethod + def put_file(self, session_id: str, dest_path: str, data: bytes) -> None: + """Write ``data`` to ``dest_path``; precondition: a relative path contained in the workspace.""" + + @abstractmethod + def get_file(self, session_id: str, path: str) -> bytes: + """Read ``path`` from the workspace; precondition: a relative path contained in the workspace.""" + + @abstractmethod + def list_files(self, session_id: str) -> List[str]: + """List workspace-relative file paths for ``session_id`` (never escapes the workspace).""" diff --git a/application/sandbox/jupyter_gateway.py b/application/sandbox/jupyter_gateway.py new file mode 100644 index 00000000..b9960986 --- /dev/null +++ b/application/sandbox/jupyter_gateway.py @@ -0,0 +1,454 @@ +"""Jupyter Kernel Gateway sandbox: stateful in-process kernels over REST + WebSocket.""" + +import base64 +import hashlib +import json +import logging +import re +import threading +import time +import uuid +from typing import Dict, List, Optional +from urllib.parse import urlencode, urlparse, urlunparse + +import requests +import websocket + +from application.sandbox.base import ( + CodeSandbox, + DisplayData, + ExecResult, + Plot, +) + +logger = logging.getLogger(__name__) + +# Per-session workspace root inside the runner container. The kernel sets its +# cwd here so relative paths from LLM code and file in/out share one directory. +_WORKSPACE_ROOT = "/tmp/docsgpt-sandbox" + +# Marker the get_file helper prints around the base64 payload so we can extract +# it from interleaved stdout without a contents API (the gateway has none). +_FILE_BEGIN = "<<>>" +_FILE_END = "<<>>" + +# Session ids become filesystem path components; allow only safe characters. +_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9_-]+$") + +# Kernel-side helper that resolves a workspace-relative path and rejects any +# absolute path or one that escapes the per-session workspace (path traversal). +_CONTAINMENT_SNIPPET = ( + "def _resolve(_base, _rel):\n" + " import os as _os\n" + " if _os.path.isabs(_rel):\n" + " raise ValueError('absolute paths are not allowed')\n" + " _root = _os.path.realpath(_base)\n" + " _rp = _os.path.realpath(_os.path.join(_root, _rel))\n" + " if _rp != _root and not _rp.startswith(_root + _os.sep):\n" + " raise ValueError('path escapes the session workspace')\n" + " return _rp\n" +) + + +class _Kernel: + """Tracks one gateway kernel plus the per-session workspace it executes in.""" + + def __init__(self, kernel_id: str, workspace: str) -> None: + self.kernel_id = kernel_id + self.workspace = workspace + self.initialized = False + + +class JupyterKernelGatewaySandbox(CodeSandbox): + """Drives one always-on Jupyter Kernel Gateway, one stateful kernel per session.""" + + def __init__( + self, + gateway_url: str, + auth_token: Optional[str] = None, + kernel_name: str = "python3", + default_timeout: float = 60.0, + http_timeout: float = 10.0, + max_output_bytes: int = 8 * 1024 * 1024, + max_file_bytes: int = 10 * 1024 * 1024, + ) -> None: + """Configure the client; no kernel is created until ``open`` is called.""" + self._base_url = gateway_url.rstrip("/") + self._auth_token = auth_token + self._kernel_name = kernel_name + self._default_timeout = default_timeout + self._http_timeout = http_timeout + self._max_output_bytes = max_output_bytes + self._max_file_bytes = max_file_bytes + self._kernels: Dict[str, _Kernel] = {} + self._lock = threading.Lock() + + # -- HTTP helpers ---------------------------------------------------- + + def _headers(self) -> Dict[str, str]: + headers = {"Content-Type": "application/json"} + if self._auth_token: + headers["Authorization"] = f"token {self._auth_token}" + return headers + + def _ws_url(self, kernel_id: str) -> str: + parsed = urlparse(self._base_url) + scheme = "wss" if parsed.scheme == "https" else "ws" + path = f"/api/kernels/{kernel_id}/channels" + query = urlencode({"token": self._auth_token}) if self._auth_token else "" + return urlunparse((scheme, parsed.netloc, path, "", query, "")) + + def _get_kernel(self, session_id: str) -> _Kernel: + with self._lock: + kernel = self._kernels.get(session_id) + if kernel is None: + raise KeyError(f"No sandbox session open for {session_id!r}") + return kernel + + @staticmethod + def _validate_session_id(session_id: str) -> None: + """Reject session ids that would not be a safe filesystem path component.""" + if not _SESSION_ID_RE.match(session_id): + raise ValueError(f"Invalid session id {session_id!r}: expected [A-Za-z0-9_-]+") + + # -- Lifecycle ------------------------------------------------------- + + def open(self, session_id: str) -> str: + """Start a fresh kernel for ``session_id`` and prime its workspace cwd.""" + self._validate_session_id(session_id) + with self._lock: + existing = self._kernels.get(session_id) + if existing is not None: + return existing.kernel_id + + resp = requests.post( + f"{self._base_url}/api/kernels", + headers=self._headers(), + data=json.dumps({"name": self._kernel_name}), + timeout=self._http_timeout, + ) + resp.raise_for_status() + kernel_id = resp.json()["id"] + workspace = f"{_WORKSPACE_ROOT}/{session_id}" + kernel = _Kernel(kernel_id, workspace) + with self._lock: + self._kernels[session_id] = kernel + self._prime(kernel) + return kernel_id + + def attach(self, session_id: str) -> str: + """Reattach to a still-running kernel for ``session_id``; open a cold one if gone.""" + self._validate_session_id(session_id) + with self._lock: + existing = self._kernels.get(session_id) + if existing is not None and self._kernel_alive(existing.kernel_id): + return existing.kernel_id + if existing is not None: + with self._lock: + self._kernels.pop(session_id, None) + logger.warning("Re-attaching session %s to a cold kernel; previous state is lost", session_id) + return self.open(session_id) + + def close(self, session_id: str) -> None: + """Delete the gateway kernel for ``session_id`` and drop it from the registry.""" + with self._lock: + kernel = self._kernels.pop(session_id, None) + if kernel is None: + return + try: + requests.delete( + f"{self._base_url}/api/kernels/{kernel.kernel_id}", + headers=self._headers(), + timeout=self._http_timeout, + ) + except requests.RequestException as exc: # teardown is best-effort + logger.warning("Failed to delete kernel %s: %s", kernel.kernel_id, exc) + + def _kernel_alive(self, kernel_id: str) -> bool: + try: + resp = requests.get( + f"{self._base_url}/api/kernels/{kernel_id}", + headers=self._headers(), + timeout=self._http_timeout, + ) + return resp.status_code == 200 + except requests.RequestException: + return False + + def _interrupt(self, kernel_id: str) -> None: + """Best-effort interrupt so a timed-out/runaway kernel becomes reusable.""" + try: + requests.post( + f"{self._base_url}/api/kernels/{kernel_id}/interrupt", + headers=self._headers(), + timeout=self._http_timeout, + ) + except requests.RequestException as exc: + logger.warning("Failed to interrupt kernel %s: %s", kernel_id, exc) + + def _interrupt_and_drain(self, ws: websocket.WebSocket, msg_id: str, kernel_id: str) -> None: + """Interrupt the kernel then drain frames until it idles, leaving the session reusable.""" + self._interrupt(kernel_id) + drain_deadline = time.monotonic() + self._http_timeout + while time.monotonic() < drain_deadline: + try: + ws.settimeout(max(0.05, drain_deadline - time.monotonic())) + raw = ws.recv() + except (websocket.WebSocketTimeoutException, websocket.WebSocketConnectionClosedException): + return + if not raw: + continue + try: + msg = json.loads(raw) + except json.JSONDecodeError: + continue + if msg.get("parent_header", {}).get("msg_id") != msg_id: + continue + msg_type = msg.get("msg_type") or msg.get("header", {}).get("msg_type") + if msg_type == "status" and msg.get("content", {}).get("execution_state") == "idle": + return + + def _prime(self, kernel: _Kernel) -> None: + """Create the per-session workspace dir and chdir the kernel into it.""" + setup = ( + "import os as _os\n" + f"_os.makedirs({kernel.workspace!r}, exist_ok=True)\n" + f"_os.chdir({kernel.workspace!r})\n" + ) + result = self._run(kernel, setup, self._default_timeout) + if not result.ok: + raise RuntimeError(f"Sandbox workspace setup failed: {result.error_value}") + kernel.initialized = True + + # -- Execution ------------------------------------------------------- + + def exec(self, session_id: str, code: str, timeout: Optional[float] = None) -> ExecResult: + """Run ``code`` in the session's persistent kernel; state carries across calls.""" + kernel = self._get_kernel(session_id) + return self._run(kernel, code, timeout or self._default_timeout) + + def _run(self, kernel: _Kernel, code: str, timeout: float) -> ExecResult: + """Execute one ``execute_request`` over the WS channel and assemble the reply.""" + try: + ws = websocket.create_connection( + self._ws_url(kernel.kernel_id), + timeout=timeout, + header=self._ws_headers(), + ) + except Exception as exc: # noqa: BLE001 - connect failure -> error result, never raise + return _error_result(type(exc).__name__, str(exc) or "failed to open kernel channel") + try: + msg_id = uuid.uuid4().hex + ws.send(json.dumps(self._execute_request(msg_id, code))) + return self._collect(ws, msg_id, timeout, kernel.kernel_id) + finally: + try: + ws.close() + except Exception: # noqa: BLE001 - closing a socket must not mask results + pass + + def _ws_headers(self) -> List[str]: + if self._auth_token: + return [f"Authorization: token {self._auth_token}"] + return [] + + @staticmethod + def _execute_request(msg_id: str, code: str) -> dict: + """Build a Jupyter ``execute_request`` wire message for the shell channel.""" + return { + "header": { + "msg_id": msg_id, + "username": "docsgpt", + "session": uuid.uuid4().hex, + "msg_type": "execute_request", + "version": "5.3", + }, + "parent_header": {}, + "metadata": {}, + "content": { + "code": code, + "silent": False, + "store_history": True, + "user_expressions": {}, + "allow_stdin": False, + "stop_on_error": True, + }, + "channel": "shell", + } + + def _collect(self, ws: websocket.WebSocket, msg_id: str, timeout: float, kernel_id: str) -> ExecResult: + """Read iopub/shell frames until ``execute_reply``/idle, a wall-clock deadline, or a closed socket.""" + result = ExecResult() + stdout_parts: List[str] = [] + stderr_parts: List[str] = [] + buffered = 0 + truncated = False + reply_seen = False + idle_seen = False + deadline = time.monotonic() + timeout + + while not (reply_seen and idle_seen): + now = time.monotonic() + remaining = deadline - now + if remaining <= 0: + self._fail(result, "TimeoutError", f"execution exceeded {timeout}s") + self._interrupt_and_drain(ws, msg_id, kernel_id) + break + try: + ws.settimeout(remaining) + raw = ws.recv() + except websocket.WebSocketTimeoutException: + self._fail(result, "TimeoutError", f"execution exceeded {timeout}s") + self._interrupt_and_drain(ws, msg_id, kernel_id) + break + except websocket.WebSocketConnectionClosedException: + self._fail(result, "KernelDiedError", "kernel channel closed before completion") + break + if not raw: + continue + try: + msg = json.loads(raw) + except json.JSONDecodeError as exc: + self._fail(result, "ProtocolError", f"malformed kernel frame: {exc}") + break + if msg.get("parent_header", {}).get("msg_id") != msg_id: + continue + + msg_type = msg.get("msg_type") or msg.get("header", {}).get("msg_type") + content = msg.get("content", {}) + + if msg_type == "stream": + if not truncated: + text = content.get("text", "") + buffered += len(text) + if buffered > self._max_output_bytes: + truncated = True + self._interrupt_and_drain(ws, msg_id, kernel_id) # runaway output: stop and drain + break + elif content.get("name") == "stderr": + stderr_parts.append(text) + else: + stdout_parts.append(text) + elif msg_type in ("execute_result", "display_data"): + self._capture_rich(result, msg_type, content) + elif msg_type == "error": + result.status = "error" + result.exit_code = 1 + result.error_name = content.get("ename") + result.error_value = content.get("evalue") + result.traceback = content.get("traceback", []) + elif msg_type == "execute_reply": + reply_seen = True + result.execution_count = content.get("execution_count") + if content.get("status") == "error": + result.status = "error" + result.exit_code = result.exit_code or 1 + result.error_name = result.error_name or content.get("ename") + result.error_value = result.error_value or content.get("evalue") + if not result.traceback: + result.traceback = content.get("traceback", []) + elif msg_type == "status": + if content.get("execution_state") == "idle": + idle_seen = True + + if truncated: + stderr_parts.append(f"\n[output truncated at {self._max_output_bytes} bytes]") + result.stdout = "".join(stdout_parts) + result.stderr = "".join(stderr_parts) + return result + + @staticmethod + def _fail(result: ExecResult, name: str, value: str) -> None: + """Mark ``result`` as a failed exec with the given error name/value.""" + result.status = "error" + result.error_name = name + result.error_value = value + result.exit_code = -1 + + @staticmethod + def _capture_rich(result: ExecResult, msg_type: str, content: dict) -> None: + """Sort a rich output into results/display_data and pull out any image plots.""" + data = content.get("data", {}) or {} + metadata = content.get("metadata", {}) or {} + bundle = DisplayData(data=data, metadata=metadata) + if msg_type == "execute_result": + result.results.append(bundle) + else: + result.display_data.append(bundle) + for mime, payload in data.items(): + if mime.startswith("image/"): + result.plots.append(Plot(format=mime.split("/", 1)[1], content_base64=payload)) + + # -- File transfer --------------------------------------------------- + + def put_file(self, session_id: str, dest_path: str, data: bytes) -> None: + """Decode ``data`` inside the kernel and write it under the session workspace.""" + kernel = self._get_kernel(session_id) + encoded = base64.b64encode(data).decode("ascii") + code = ( + "import base64 as _b64, os as _os\n" + + _CONTAINMENT_SNIPPET + + f"_p = _resolve({kernel.workspace!r}, {dest_path!r})\n" + "_os.makedirs(_os.path.dirname(_p) or '.', exist_ok=True)\n" + f"_f = open(_p, 'wb'); _f.write(_b64.b64decode({encoded!r})); _f.close()\n" + ) + result = self._run(kernel, code, self._default_timeout) + if not result.ok: + raise IOError(f"put_file failed: {result.error_value}") + + def get_file(self, session_id: str, path: str) -> bytes: + """Read ``path`` inside the kernel and stream its base64 (with a length tag) over stdout.""" + kernel = self._get_kernel(session_id) + code = ( + "import base64 as _b64, hashlib as _hl, os as _os\n" + + _CONTAINMENT_SNIPPET + + f"_p = _resolve({kernel.workspace!r}, {path!r})\n" + f"_sz = _os.path.getsize(_p)\n" + f"if _sz > {self._max_file_bytes}:\n" + f" raise ValueError('file too large: %d > {self._max_file_bytes} bytes' % _sz)\n" + "_d = open(_p, 'rb').read()\n" + "_h = _hl.sha256(_d).hexdigest()\n" + f"print({_FILE_BEGIN!r} + str(len(_d)) + ':' + _h + ':'" + f" + _b64.b64encode(_d).decode('ascii') + {_FILE_END!r})\n" + ) + result = self._run(kernel, code, self._default_timeout) + if not result.ok: + raise IOError(f"get_file failed: {result.error_value}") + out = result.stdout + start = out.find(_FILE_BEGIN) + end = out.find(_FILE_END) + if start == -1 or end == -1: + raise IOError(f"get_file produced no payload for {path!r}") + payload = out[start + len(_FILE_BEGIN):end] + expected_len_s, expected_sha, encoded = payload.split(":", 2) + decoded = base64.b64decode(encoded) + if len(decoded) != int(expected_len_s) or hashlib.sha256(decoded).hexdigest() != expected_sha: + raise IOError(f"get_file integrity check failed for {path!r} (payload truncated)") + return decoded + + def list_files(self, session_id: str) -> List[str]: + """Walk the session workspace inside the kernel and return relative paths.""" + kernel = self._get_kernel(session_id) + code = ( + "import os as _os, json as _json\n" + f"_root = {kernel.workspace!r}\n" + "_out = []\n" + "for _dp, _dn, _fn in _os.walk(_root):\n" + " for _name in _fn:\n" + " _out.append(_os.path.relpath(_os.path.join(_dp, _name), _root))\n" + f"print({_FILE_BEGIN!r} + _json.dumps(_out) + {_FILE_END!r})\n" + ) + result = self._run(kernel, code, self._default_timeout) + if not result.ok: + raise IOError(f"list_files failed: {result.error_value}") + out = result.stdout + start = out.find(_FILE_BEGIN) + end = out.find(_FILE_END) + if start == -1 or end == -1: + return [] + return json.loads(out[start + len(_FILE_BEGIN):end]) + + +def _error_result(name: str, value: str) -> ExecResult: + """Build a failed ExecResult carrying the given error name/value.""" + return ExecResult(status="error", error_name=name, error_value=value, exit_code=-1) diff --git a/application/sandbox/manager.py b/application/sandbox/manager.py new file mode 100644 index 00000000..183a08a6 --- /dev/null +++ b/application/sandbox/manager.py @@ -0,0 +1,135 @@ +"""Process-wide registry binding session ids to sandbox handles with idle expiry.""" + +import threading +import time +from dataclasses import dataclass, field +from typing import Dict, List, Optional + +from application.sandbox.base import CodeSandbox, ExecResult + + +@dataclass +class _Session: + """Bookkeeping for one bound sandbox session: its TTL and access timestamps.""" + + session_id: str + ttl: float + created_at: float + last_access: float = field(default=0.0) + + def is_expired(self, now: float) -> bool: + """True when the session has been idle longer than its (clamped) TTL.""" + return (now - self.last_access) > self.ttl + + +class SandboxManager: + """Binds session ids to a shared backend, clamps TTLs, and reaps idle sessions.""" + + def __init__(self, backend: CodeSandbox, max_ttl: float, default_ttl: Optional[float] = None) -> None: + """Wrap ``backend`` with a registry clamped to ``max_ttl`` seconds per session.""" + self._backend = backend + self._max_ttl = max_ttl + self._default_ttl = default_ttl if default_ttl is not None else max_ttl + self._sessions: Dict[str, _Session] = {} + self._lock = threading.Lock() + + def _clamp_ttl(self, ttl: Optional[float]) -> float: + """Return an agent-requested TTL bounded to (0, ``max_ttl``].""" + if ttl is None: + ttl = self._default_ttl + if ttl <= 0: + return self._default_ttl + return min(ttl, self._max_ttl) + + def open(self, session_id: str, ttl: Optional[float] = None) -> str: + """Open (or reuse) the sandbox for ``session_id`` with a clamped TTL.""" + self.reap_expired() + with self._lock: + session = self._sessions.get(session_id) + now = time.monotonic() + if session is not None: + session.last_access = now + handle = self._backend.attach(session_id) + return handle + handle = self._backend.open(session_id) + self._sessions[session_id] = _Session( + session_id=session_id, + ttl=self._clamp_ttl(ttl), + created_at=now, + last_access=now, + ) + return handle + + def attach(self, session_id: str) -> str: + """Reattach to an existing session, refreshing its idle clock.""" + with self._lock: + session = self._sessions.get(session_id) + if session is None: + raise KeyError(f"No sandbox session bound for {session_id!r}") + session.last_access = time.monotonic() + return self._backend.attach(session_id) + + def exec(self, session_id: str, code: str, timeout: Optional[float] = None) -> ExecResult: + """Execute ``code`` in the bound session, touching its idle clock first.""" + with self._lock: + session = self._sessions.get(session_id) + if session is None: + raise KeyError(f"No sandbox session bound for {session_id!r}") + session.last_access = time.monotonic() + return self._backend.exec(session_id, code, timeout) + + def put_file(self, session_id: str, dest_path: str, data: bytes) -> None: + """Write ``data`` into the bound session's workspace.""" + self._touch(session_id) + self._backend.put_file(session_id, dest_path, data) + + def get_file(self, session_id: str, path: str) -> bytes: + """Read ``path`` from the bound session's workspace.""" + self._touch(session_id) + return self._backend.get_file(session_id, path) + + def list_files(self, session_id: str) -> List[str]: + """List files in the bound session's workspace.""" + self._touch(session_id) + return self._backend.list_files(session_id) + + def close(self, session_id: str) -> None: + """Tear down the backend runtime and drop the session from the registry.""" + with self._lock: + self._sessions.pop(session_id, None) + self._backend.close(session_id) + + def has_session(self, session_id: str) -> bool: + """True when ``session_id`` is currently bound in the registry.""" + with self._lock: + return session_id in self._sessions + + def ttl_for(self, session_id: str) -> Optional[float]: + """Return the clamped TTL bound to ``session_id``, or None if unbound.""" + with self._lock: + session = self._sessions.get(session_id) + return session.ttl if session else None + + def _touch(self, session_id: str) -> None: + with self._lock: + session = self._sessions.get(session_id) + if session is None: + raise KeyError(f"No sandbox session bound for {session_id!r}") + session.last_access = time.monotonic() + + def reap_expired(self) -> List[str]: + """Close sessions idle past their TTL and return the reaped session ids. + + Extension point: a later hardening slice runs this on a background + reaper and, before close, flushes workspace files to the artifact + store (persist-on-reap) so a fresh kernel can re-mount them by + reference on next access. For now it is a plain idle close. + """ + now = time.monotonic() + with self._lock: + expired = [sid for sid, s in self._sessions.items() if s.is_expired(now)] + for sid in expired: + self._sessions.pop(sid, None) + for sid in expired: + self._backend.close(sid) + return expired diff --git a/application/sandbox/sandbox_creator.py b/application/sandbox/sandbox_creator.py new file mode 100644 index 00000000..16ef4028 --- /dev/null +++ b/application/sandbox/sandbox_creator.py @@ -0,0 +1,56 @@ +"""Factory + process-wide singleton selecting a sandbox backend from settings.""" + +from typing import Callable, Dict + +from application.core.settings import settings +from application.sandbox.base import CodeSandbox +from application.sandbox.jupyter_gateway import JupyterKernelGatewaySandbox +from application.sandbox.manager import SandboxManager + + +def _make_jupyter() -> CodeSandbox: + """Build the Jupyter Kernel Gateway backend from ``SANDBOX_*`` settings.""" + return JupyterKernelGatewaySandbox( + gateway_url=settings.SANDBOX_GATEWAY_URL, + auth_token=settings.SANDBOX_GATEWAY_AUTH_TOKEN, + kernel_name=settings.SANDBOX_KERNEL_NAME, + default_timeout=float(settings.SANDBOX_EXEC_TIMEOUT), + http_timeout=float(settings.SANDBOX_HTTP_TIMEOUT), + max_output_bytes=int(settings.SANDBOX_MAX_OUTPUT_BYTES), + max_file_bytes=int(settings.SANDBOX_MAX_FILE_BYTES), + ) + + +class SandboxCreator: + """Resolves ``SANDBOX_BACKEND`` to a backend and caches a single manager.""" + + # Seam for a future "daytona" backend (Apache-2.0 SDK); not implemented here. + backends: Dict[str, Callable[[], CodeSandbox]] = { + "jupyter": _make_jupyter, + } + + _instance = None + + @classmethod + def get_manager(cls) -> SandboxManager: + """Return the process-wide ``SandboxManager``, building it on first use.""" + if cls._instance is None: + backend = cls.create_backend(getattr(settings, "SANDBOX_BACKEND", "jupyter")) + cls._instance = SandboxManager( + backend=backend, + max_ttl=float(settings.SANDBOX_MAX_TTL), + ) + return cls._instance + + @classmethod + def create_backend(cls, type_name: str) -> CodeSandbox: + """Instantiate the backend registered under ``type_name`` (case-insensitive).""" + factory = cls.backends.get(type_name.lower()) + if not factory: + raise ValueError(f"No sandbox backend found for type {type_name}") + return factory() + + @classmethod + def reset(cls) -> None: + """Drop the cached singleton (test/teardown hook).""" + cls._instance = None diff --git a/deployment/docker-compose.yaml b/deployment/docker-compose.yaml index 93c38240..e165a7f4 100644 --- a/deployment/docker-compose.yaml +++ b/deployment/docker-compose.yaml @@ -24,6 +24,8 @@ 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 ports: - "7091:7091" volumes: @@ -49,6 +51,7 @@ 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 volumes: - ../application/indexes:/app/indexes - ../application/inputs:/app/inputs @@ -59,6 +62,30 @@ services: postgres: condition: service_healthy + # Always-on code-execution runner (Jupyter Kernel Gateway). Sessions are + # in-process kernels, never child containers; the Docker socket is NOT + # mounted. On an internal-only network — no host port is published, so the + # runner is reachable only from backend/worker, not from the host/internet. + # Egress/SSRF blocks, the gVisor `runsc` runtime, and seccomp profile come in + # the hardening slice. + 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: + # Per-session workspaces (/tmp/docsgpt-sandbox/) and Jupyter + # runtime files live on tmpfs; the root FS is read-only everywhere else. + - /tmp + networks: + - sandbox-net + - default + redis: image: redis:6-alpine ports: @@ -80,6 +107,12 @@ services: timeout: 5s retries: 10 +networks: + # Internal-only network for the sandbox runner (no external egress route via + # this network; the runner still reaches the internet via the default bridge). + sandbox-net: + internal: true + volumes: postgres_data: diff --git a/deployment/sandbox/Dockerfile b/deployment/sandbox/Dockerfile new file mode 100644 index 00000000..8e36addf --- /dev/null +++ b/deployment/sandbox/Dockerfile @@ -0,0 +1,48 @@ +# docsgpt-sandbox: the always-on code-execution runner. +# +# One container runs a Jupyter Kernel Gateway; each sandbox session is an +# in-process kernel (a child process), NOT a child container. This image does +# NOT mount the Docker socket and never spawns containers — that avoids the +# host-root risk of Docker-in-Docker and is the core security win. +# +# The app (backend/worker) is the CLIENT and reaches this service over +# HTTP + WebSocket via SANDBOX_GATEWAY_URL. +# +# Pre-baked doc libs are all permissive (MIT/BSD/Apache) for speed/reliability; +# runtime `pip install` still works because egress is open. PyMuPDF and any +# other AGPL lib are intentionally excluded; Docling is deferred to its own slice. +# +# HARDENING (separate slice — NOT done here): run under the gVisor `runsc` +# runtime, add network-layer SSRF blocks (drop RFC1918 / link-local / +# 169.254.169.254), seccomp profile, read-only root FS + quota'd scratch dir, +# and cgroup CPU/mem/PID caps wired from SANDBOX_MEMORY / SANDBOX_CPUS. +FROM python:3.12-slim + +# Non-root user for the runner (untrusted code runs as this UID). +RUN useradd --create-home --uid 10001 sandbox + +# Exact pins for reproducible builds and no license drift (all MIT/BSD/Apache). +RUN pip install --no-cache-dir \ + jupyter-kernel-gateway==3.0.1 \ + ipykernel==6.29.5 \ + python-pptx==1.0.2 \ + python-docx==1.1.2 \ + openpyxl==3.1.5 \ + reportlab==4.2.5 \ + pandas==2.2.3 \ + matplotlib==3.9.2 + +USER sandbox +WORKDIR /home/sandbox + +EXPOSE 8888 + +# ip=0.0.0.0 so the backend can reach it over the internal compose network. +# allow_origin is intentionally NOT set to "*" — publishing port 8888 requires +# setting SANDBOX_GATEWAY_AUTH_TOKEN (see README). limit_rate=False raises the +# iopub data-rate cap so large get_file base64 payloads are not truncated; the +# get_file integrity check still guards against truncation if this is ever off. +CMD ["jupyter", "kernelgateway", \ + "--KernelGatewayApp.ip=0.0.0.0", \ + "--KernelGatewayApp.port=8888", \ + "--ZMQChannelsWebsocketConnection.limit_rate=False"] diff --git a/deployment/sandbox/README.md b/deployment/sandbox/README.md new file mode 100644 index 00000000..0a46116c --- /dev/null +++ b/deployment/sandbox/README.md @@ -0,0 +1,50 @@ +# docsgpt-sandbox runner + +Always-on Jupyter Kernel Gateway that executes sandboxed LLM code. The DocsGPT +backend/worker is the **client** and connects over HTTP + WebSocket via +`SANDBOX_GATEWAY_URL`. Each session is an **in-process kernel** (child process), +never a child container; the Docker socket is **not** mounted. + +## Run standalone for dev + +Build and run the runner on its own, then point the app at it: + +```bash +docker build -t docsgpt-sandbox deployment/sandbox +docker run --rm -p 8888:8888 docsgpt-sandbox +# in the app's .env: SANDBOX_GATEWAY_URL=http://localhost:8888 +``` + +Without Docker (matches the test harness) you can run the gateway directly from +a venv that has `jupyter-kernel-gateway` installed: + +```bash +jupyter kernelgateway --KernelGatewayApp.ip=0.0.0.0 --KernelGatewayApp.port=8888 \ + --ZMQChannelsWebsocketConnection.limit_rate=False +``` + +`--ZMQChannelsWebsocketConnection.limit_rate=False` raises the iopub data-rate +limit so large `get_file` base64 payloads aren't silently truncated. (On older +gateways the trait may live elsewhere; the client's `get_file` integrity check +catches any truncation regardless.) + +## Exposing the port requires auth + +The image does **not** set `--KernelGatewayApp.allow_origin=*`. If you publish +port 8888 (e.g. `docker run -p 8888:8888`), set `SANDBOX_GATEWAY_AUTH_TOKEN` +and launch the gateway with a matching `--KernelGatewayApp.auth_token` so the +runner is not an open arbitrary-code-execution endpoint. In compose the runner +stays on the internal-only network with no published port, so no token is +required there. + +## In docker-compose + +The `docsgpt-sandbox` service is defined in `deployment/docker-compose.yaml` on +an internal-only network. The backend and worker reach it at +`http://docsgpt-sandbox:8888`. + +## Hardening (separate slice — not in this image) + +Egress/SSRF blocks (drop RFC1918, link-local, `169.254.169.254`), the gVisor +`runsc` runtime, seccomp, read-only root FS, and cgroup CPU/mem/PID caps wired +from `SANDBOX_MEMORY` / `SANDBOX_CPUS` land in the hardening slice. diff --git a/tests/requirements.txt b/tests/requirements.txt index 2af47c93..f34900ea 100644 --- a/tests/requirements.txt +++ b/tests/requirements.txt @@ -5,3 +5,9 @@ coverage>=7.4.0 pytest-postgresql>=6.0.0 psycopg[binary]>=3.1 cryptography>=46.0.0 +# Runner engine for the sandbox integration test: tests launch a local +# `jupyter kernelgateway` process (no Docker). Not an app runtime dependency. +# ipykernel provides the python3 kernel the gateway launches per session. +jupyter-kernel-gateway>=2.5.0 +ipykernel>=6.0.0 +websocket-client>=1.6.0 diff --git a/tests/sandbox/__init__.py b/tests/sandbox/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/sandbox/test_jupyter_gateway_integration.py b/tests/sandbox/test_jupyter_gateway_integration.py new file mode 100644 index 00000000..20a01a77 --- /dev/null +++ b/tests/sandbox/test_jupyter_gateway_integration.py @@ -0,0 +1,292 @@ +"""End-to-end test against a locally launched Jupyter Kernel Gateway subprocess. + +Launches a real `jupyter kernelgateway` on an ephemeral port (no Docker — the +Docker credential helper hangs on this machine), points the backend at it, and +exercises the full session lifecycle. Skips gracefully if the gateway binary or +the websocket-client library is unavailable. +""" + +import shutil +import socket +import subprocess +import time + +import pytest + +requests = pytest.importorskip("requests") +pytest.importorskip("websocket") # websocket-client + +from application.sandbox.jupyter_gateway import JupyterKernelGatewaySandbox # noqa: E402 + +_GATEWAY_BIN = shutil.which("jupyter-kernelgateway") or shutil.which("jupyter") + +pytestmark = pytest.mark.skipif( + _GATEWAY_BIN is None, + reason="jupyter kernel gateway not installed (pip install jupyter-kernel-gateway)", +) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _gateway_cmd(port: int) -> list: + if _GATEWAY_BIN.endswith("jupyter-kernelgateway"): + base = [_GATEWAY_BIN] + else: + base = [_GATEWAY_BIN, "kernelgateway"] + return base + [ + "--KernelGatewayApp.ip=127.0.0.1", + f"--KernelGatewayApp.port={port}", + # Raise the iopub data-rate limit so large get_file payloads aren't truncated. + "--ZMQChannelsWebsocketConnection.limit_rate=False", + ] + + +@pytest.fixture(scope="module") +def gateway_url(): + port = _free_port() + proc = subprocess.Popen( + _gateway_cmd(port), + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + url = f"http://127.0.0.1:{port}" + deadline = time.time() + 30 + ready = False + try: + while time.time() < deadline: + if proc.poll() is not None: + pytest.skip("jupyter kernelgateway process exited during startup") + try: + resp = requests.get(f"{url}/api", timeout=1) + if resp.status_code == 200: + ready = True + break + except requests.RequestException: + time.sleep(0.3) + if not ready: + pytest.skip("jupyter kernelgateway did not become ready in time") + yield url + finally: + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +@pytest.fixture() +def sandbox(gateway_url): + sb = JupyterKernelGatewaySandbox(gateway_url=gateway_url, default_timeout=30.0) + yield sb + + +def test_open_and_close_lifecycle(sandbox): + session = "conv-lifecycle" + kernel_id = sandbox.open(session) + assert kernel_id + resp = requests.get(f"{sandbox._base_url}/api/kernels/{kernel_id}", timeout=5) + assert resp.status_code == 200 + sandbox.close(session) + resp = requests.get(f"{sandbox._base_url}/api/kernels/{kernel_id}", timeout=5) + assert resp.status_code == 404 # kernel torn down + + +def test_stateful_exec_carries_variables(sandbox): + session = "conv-state" + sandbox.open(session) + try: + first = sandbox.exec(session, "x = 1") + assert first.ok, first.error_value + second = sandbox.exec(session, "print(x + 41)") + assert second.ok, second.error_value + assert second.stdout.strip() == "42" + assert second.execution_count is not None + finally: + sandbox.close(session) + + +def test_stdout_capture(sandbox): + session = "conv-stdout" + sandbox.open(session) + try: + res = sandbox.exec(session, "print('hello sandbox')") + assert res.ok + assert "hello sandbox" in res.stdout + finally: + sandbox.close(session) + + +def test_error_capture_sets_failure_status(sandbox): + session = "conv-error" + sandbox.open(session) + try: + res = sandbox.exec(session, "raise ValueError('boom')") + assert not res.ok + assert res.status == "error" + assert res.error_name == "ValueError" + assert "boom" in (res.error_value or "") + assert res.traceback + finally: + sandbox.close(session) + + +def test_put_and_get_file_roundtrip_exact_bytes(sandbox): + session = "conv-files" + sandbox.open(session) + try: + payload = bytes(range(256)) * 4 # 1 KiB of every byte value + sandbox.put_file(session, "data/blob.bin", payload) + assert "data/blob.bin" in sandbox.list_files(session) + fetched = sandbox.get_file(session, "data/blob.bin") + assert fetched == payload + finally: + sandbox.close(session) + + +def test_file_written_by_code_is_readable(sandbox): + session = "conv-codefile" + sandbox.open(session) + try: + res = sandbox.exec(session, "open('out.txt', 'w').write('generated')") + assert res.ok, res.error_value + assert sandbox.get_file(session, "out.txt") == b"generated" + finally: + sandbox.close(session) + + +def test_attach_reuses_running_kernel(sandbox): + session = "conv-attach" + first_id = sandbox.open(session) + try: + sandbox.exec(session, "y = 7") + second_id = sandbox.attach(session) + assert second_id == first_id + res = sandbox.exec(session, "print(y)") + assert res.stdout.strip() == "7" + finally: + sandbox.close(session) + + +def test_exec_timeout_reports_error(sandbox): + session = "conv-timeout" + sandbox.open(session) + try: + res = sandbox.exec(session, "import time; time.sleep(5)", timeout=1.0) + assert not res.ok + assert res.error_name == "TimeoutError" + finally: + sandbox.close(session) + + +def test_kernel_killed_mid_exec_returns_error(sandbox): + session = "conv-killed" + kernel_id = sandbox.open(session) + try: + import threading + + def _kill(): + time.sleep(0.5) + requests.delete(f"{sandbox._base_url}/api/kernels/{kernel_id}", timeout=5) + + killer = threading.Thread(target=_kill) + killer.start() + res = sandbox.exec(session, "import time\nfor _ in range(50): time.sleep(0.2)", timeout=20.0) + killer.join() + assert not res.ok # error result, no raise / no hang + assert res.exit_code == -1 + finally: + sandbox.close(session) + + +def test_runaway_loop_times_out_and_session_stays_reusable(sandbox): + session = "conv-runaway" + sandbox.open(session) + try: + start = time.monotonic() + res = sandbox.exec(session, "while True:\n print('x')", timeout=2.0) + elapsed = time.monotonic() - start + assert not res.ok + assert res.error_name == "TimeoutError" + assert elapsed < 2.0 + 8.0 # returns within timeout + interrupt slack + # Interrupt should have freed the kernel: a fresh exec on the same session works. + again = sandbox.exec(session, "print(2 + 2)", timeout=10.0) + assert again.ok, again.error_value + assert again.stdout.strip() == "4" + finally: + sandbox.close(session) + + +def test_output_cap_truncates_runaway_output(gateway_url): + sb = JupyterKernelGatewaySandbox(gateway_url=gateway_url, default_timeout=30.0, max_output_bytes=50_000) + session = "conv-bigout" + sb.open(session) + try: + res = sb.exec(session, "for _ in range(100000): print('A' * 1000)", timeout=20.0) + assert "truncated" in res.stderr + assert len(res.stdout) <= 60_000 # capped well under the unbounded ~100MB + again = sb.exec(session, "print('alive')", timeout=10.0) + assert again.ok and "alive" in again.stdout + finally: + sb.close(session) + + +def test_put_file_rejects_path_traversal(sandbox): + session = "conv-trav-put" + sandbox.open(session) + try: + with pytest.raises(IOError): + sandbox.put_file(session, "../escape.txt", b"nope") + with pytest.raises(IOError): + sandbox.put_file(session, "/etc/escape.txt", b"nope") + # The escape target must not exist (no file written outside the workspace). + probe = sandbox.exec(session, "import os; print(os.path.exists('/etc/escape.txt'))") + assert probe.stdout.strip() == "False" + finally: + sandbox.close(session) + + +def test_get_file_rejects_path_traversal(sandbox): + session = "conv-trav-get" + sandbox.open(session) + try: + with pytest.raises(IOError): + sandbox.get_file(session, "../../etc/passwd") + with pytest.raises(IOError): + sandbox.get_file(session, "/etc/passwd") + finally: + sandbox.close(session) + + +def test_get_file_rejects_oversized_file(gateway_url): + sb = JupyterKernelGatewaySandbox(gateway_url=gateway_url, default_timeout=30.0, max_file_bytes=1024) + session = "conv-toobig" + sb.open(session) + try: + sb.put_file(session, "big.bin", b"x" * 4096) # 4 KiB > 1 KiB cap + with pytest.raises(IOError): + sb.get_file(session, "big.bin") + finally: + sb.close(session) + + +def test_get_file_integrity_roundtrip_large(gateway_url): + sb = JupyterKernelGatewaySandbox(gateway_url=gateway_url, default_timeout=60.0) + session = "conv-large" + sb.open(session) + try: + payload = bytes(range(256)) * 8192 # 2 MiB, exercises the iopub path + integrity check + sb.put_file(session, "big.bin", payload) + fetched = sb.get_file(session, "big.bin") + assert fetched == payload + finally: + sb.close(session) + + +def test_open_rejects_illegal_session_id(sandbox): + for bad in ["../evil", "a/b", "a b", "a;b", ""]: + with pytest.raises(ValueError): + sandbox.open(bad) diff --git a/tests/sandbox/test_sandbox_manager.py b/tests/sandbox/test_sandbox_manager.py new file mode 100644 index 00000000..c6736cbb --- /dev/null +++ b/tests/sandbox/test_sandbox_manager.py @@ -0,0 +1,161 @@ +"""Unit tests for SandboxManager and SandboxCreator using an in-memory backend.""" + +from typing import Dict, List + +import pytest + +from application.sandbox.base import CodeSandbox, ExecResult +from application.sandbox.manager import SandboxManager + + +class FakeBackend(CodeSandbox): + """In-memory backend recording calls; no network, fully deterministic.""" + + def __init__(self) -> None: + self.open_calls: List[str] = [] + self.attach_calls: List[str] = [] + self.closed: List[str] = [] + self.files: Dict[str, Dict[str, bytes]] = {} + self._handles: Dict[str, str] = {} + + def open(self, session_id: str) -> str: + self.open_calls.append(session_id) + self.files.setdefault(session_id, {}) + handle = f"handle-{session_id}-{len(self.open_calls)}" + self._handles[session_id] = handle + return handle + + def attach(self, session_id: str) -> str: + self.attach_calls.append(session_id) + if session_id in self._handles: + return self._handles[session_id] + return self.open(session_id) + + def close(self, session_id: str) -> None: + self.closed.append(session_id) + self.files.pop(session_id, None) + self._handles.pop(session_id, None) + + def exec(self, session_id, code, timeout=None) -> ExecResult: + return ExecResult(status="ok", stdout=f"ran:{code}") + + def put_file(self, session_id, dest_path, data) -> None: + self.files.setdefault(session_id, {})[dest_path] = data + + def get_file(self, session_id, path) -> bytes: + return self.files[session_id][path] + + def list_files(self, session_id) -> List[str]: + return list(self.files.get(session_id, {}).keys()) + + +@pytest.fixture() +def backend() -> FakeBackend: + return FakeBackend() + + +def test_open_registers_session_and_opens_backend(backend): + mgr = SandboxManager(backend, max_ttl=600) + handle = mgr.open("conv-1") + assert handle.startswith("handle-conv-1") + assert backend.open_calls == ["conv-1"] + assert mgr.has_session("conv-1") + + +def test_open_twice_reuses_via_attach(backend): + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("conv-1") + mgr.open("conv-1") + assert backend.open_calls == ["conv-1"] # opened once + assert backend.attach_calls == ["conv-1"] # second open reused via attach + + +def test_attach_reuse_requires_existing_session(backend): + mgr = SandboxManager(backend, max_ttl=600) + with pytest.raises(KeyError): + mgr.attach("missing") + mgr.open("conv-1") + assert mgr.attach("conv-1").startswith("handle-conv-1") + + +def test_ttl_clamped_to_max(backend): + mgr = SandboxManager(backend, max_ttl=300) + mgr.open("conv-1", ttl=99999) + assert mgr.ttl_for("conv-1") == 300 + + +def test_ttl_honored_when_below_max(backend): + mgr = SandboxManager(backend, max_ttl=300) + mgr.open("conv-1", ttl=120) + assert mgr.ttl_for("conv-1") == 120 + + +def test_ttl_default_used_for_nonpositive(backend): + mgr = SandboxManager(backend, max_ttl=300, default_ttl=200) + mgr.open("conv-1", ttl=0) + assert mgr.ttl_for("conv-1") == 200 + mgr.open("conv-2", ttl=-5) + assert mgr.ttl_for("conv-2") == 200 + + +def test_close_drops_registry_and_backend(backend): + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("conv-1") + mgr.close("conv-1") + assert not mgr.has_session("conv-1") + assert backend.closed == ["conv-1"] + + +def test_reap_expired_closes_idle_sessions(backend, monkeypatch): + clock = {"t": 1000.0} + monkeypatch.setattr("application.sandbox.manager.time.monotonic", lambda: clock["t"]) + mgr = SandboxManager(backend, max_ttl=100) + mgr.open("conv-1", ttl=50) + clock["t"] = 1051.0 # 51s idle > 50s ttl + reaped = mgr.reap_expired() + assert reaped == ["conv-1"] + assert backend.closed == ["conv-1"] + assert not mgr.has_session("conv-1") + + +def test_exec_and_file_roundtrip_through_manager(backend): + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("conv-1") + res = mgr.exec("conv-1", "1+1") + assert res.ok and res.stdout == "ran:1+1" + mgr.put_file("conv-1", "a.txt", b"hello") + assert mgr.get_file("conv-1", "a.txt") == b"hello" + assert mgr.list_files("conv-1") == ["a.txt"] + + +def test_file_ops_require_open_session(backend): + mgr = SandboxManager(backend, max_ttl=600) + with pytest.raises(KeyError): + mgr.get_file("missing", "a.txt") + + +def test_sandbox_creator_selects_jupyter_backend(monkeypatch): + from application.sandbox import sandbox_creator as sc + + sc.SandboxCreator.reset() + backend = sc.SandboxCreator.create_backend("jupyter") + from application.sandbox.jupyter_gateway import JupyterKernelGatewaySandbox + + assert isinstance(backend, JupyterKernelGatewaySandbox) + + +def test_sandbox_creator_unknown_backend_raises(): + from application.sandbox.sandbox_creator import SandboxCreator + + with pytest.raises(ValueError): + SandboxCreator.create_backend("does-not-exist") + + +def test_sandbox_creator_manager_is_singleton(): + from application.sandbox.sandbox_creator import SandboxCreator + + SandboxCreator.reset() + m1 = SandboxCreator.get_manager() + m2 = SandboxCreator.get_manager() + assert m1 is m2 + SandboxCreator.reset() From 182a14d257fe2e082d43fe29dbc5a68a196d8c65 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 24 Jun 2026 11:29:33 +0100 Subject: [PATCH 04/59] Add artifact preview, versions, and download to the sidebar Extend the artifact sidebar with a document/file view: a version selector, authenticated download honoring Content-Disposition (following the backend redirect to a presigned URL when configured), and version restore. Preview is chosen by kind - html, svg, and mermaid render inside a sandboxed iframe with a restrictive CSP (mermaid via its own sandboxed render so untrusted diagram text never reaches the app DOM), office documents and PDFs show a download card, and code or data render as text. Wires the new artifact endpoints into the API service and leaves the existing notes/todo path unchanged. --- frontend/src/api/endpoints.ts | 8 + frontend/src/api/services/userService.ts | 21 ++ frontend/src/components/ArtifactSidebar.tsx | 174 ++++++--- .../src/components/DocumentArtifactView.tsx | 344 ++++++++++++++++++ .../src/components/artifactViewUtils.test.ts | 192 ++++++++++ frontend/src/components/artifactViewUtils.ts | 227 ++++++++++++ 6 files changed, 907 insertions(+), 59 deletions(-) create mode 100644 frontend/src/components/DocumentArtifactView.tsx create mode 100644 frontend/src/components/artifactViewUtils.test.ts create mode 100644 frontend/src/components/artifactViewUtils.ts diff --git a/frontend/src/api/endpoints.ts b/frontend/src/api/endpoints.ts index 65a9dc15..771596af 100644 --- a/frontend/src/api/endpoints.ts +++ b/frontend/src/api/endpoints.ts @@ -110,6 +110,14 @@ const endpoints = { AGENT_FOLDER: (id: string) => `/api/agents/folders/${id}`, MOVE_AGENT_TO_FOLDER: '/api/agents/folders/move_agent', GET_ARTIFACT: (artifactId: string) => `/api/artifact/${artifactId}`, + GET_DOCUMENT_ARTIFACT: (artifactId: string) => + `/api/artifacts/${artifactId}`, + DOWNLOAD_ARTIFACT: (artifactId: string, version?: number) => + version != null + ? `/api/artifacts/${artifactId}/download?version=${version}` + : `/api/artifacts/${artifactId}/download`, + RESTORE_ARTIFACT: (artifactId: string) => + `/api/artifacts/${artifactId}/restore`, WORKFLOWS: '/api/workflows', WORKFLOW: (id: string) => `/api/workflows/${id}`, CUSTOM_MODELS: '/api/user/models', diff --git a/frontend/src/api/services/userService.ts b/frontend/src/api/services/userService.ts index e3f49659..af3118fe 100644 --- a/frontend/src/api/services/userService.ts +++ b/frontend/src/api/services/userService.ts @@ -322,6 +322,27 @@ const userService = { apiClient.post(endpoints.USER.MOVE_AGENT_TO_FOLDER, data, token), getArtifact: (artifactId: string, token: string | null): Promise => apiClient.get(endpoints.USER.GET_ARTIFACT(artifactId), token), + getDocumentArtifact: ( + artifactId: string, + token: string | null, + ): Promise => + apiClient.get(endpoints.USER.GET_DOCUMENT_ARTIFACT(artifactId), token), + downloadArtifact: ( + artifactId: string, + token: string | null, + version?: number, + ): Promise => + apiClient.get(endpoints.USER.DOWNLOAD_ARTIFACT(artifactId, version), token), + restoreArtifactVersion: ( + artifactId: string, + version: number, + token: string | null, + ): Promise => + apiClient.post( + endpoints.USER.RESTORE_ARTIFACT(artifactId), + { version }, + token, + ), getWorkflow: (id: string, token: string | null): Promise => apiClient.get(endpoints.USER.WORKFLOW(id), token), createWorkflow: (data: any, token: string | null): Promise => diff --git a/frontend/src/components/ArtifactSidebar.tsx b/frontend/src/components/ArtifactSidebar.tsx index 718831ed..4c003703 100644 --- a/frontend/src/components/ArtifactSidebar.tsx +++ b/frontend/src/components/ArtifactSidebar.tsx @@ -12,7 +12,9 @@ import remarkGfm from 'remark-gfm'; import userService from '../api/services/userService'; import { useDarkTheme } from '../hooks'; import { selectToken } from '../preferences/preferenceSlice'; +import { isDocumentArtifact, type DocumentArtifact } from './artifactViewUtils'; import CopyButton from './CopyButton'; +import DocumentArtifactView from './DocumentArtifactView'; import Spinner from './Spinner'; import { Button } from './ui/button'; import { Sheet, SheetContent } from './ui/sheet'; @@ -282,13 +284,16 @@ export default function ArtifactSidebar({ const currentFetchIdRef = React.useRef(null); const token = useSelector(selectToken); const [artifact, setArtifact] = useState(null); + const [documentArtifact, setDocumentArtifact] = + useState(null); const [loading, setLoading] = useState(false); const [error, setError] = useState(null); + const [refreshNonce, setRefreshNonce] = useState(0); const [effectiveArtifactId, setEffectiveArtifactId] = useState( artifactId, ); - const title = getArtifactTitle(artifact, toolName); + const title = documentArtifact?.title || getArtifactTitle(artifact, toolName); // Reset last successful todo artifact ID when conversation changes useEffect(() => { @@ -308,6 +313,7 @@ export default function ArtifactSidebar({ useEffect(() => { if (!isOpen || !effectiveArtifactId) { setArtifact(null); + setDocumentArtifact(null); setError(null); setLoading(false); currentFetchIdRef.current = null; @@ -315,76 +321,118 @@ export default function ArtifactSidebar({ } // Generate a unique ID for this fetch - const fetchId = `${effectiveArtifactId}-${Date.now()}`; + const artifactIdToFetch = effectiveArtifactId; + const fetchId = `${artifactIdToFetch}-${Date.now()}`; currentFetchIdRef.current = fetchId; setLoading(true); setError(null); - // Note: For todo artifacts, the endpoint always returns all todos for the tool; will be coversation scoped later + // Document/file artifacts live behind the generalized /api/artifacts/ + // endpoint; notes/todos behind the legacy /api/artifact/. Try the + // document path first and fall back to the legacy flow on a non-document + // shape so a single id resolves cleanly either way. userService - .getArtifact(effectiveArtifactId, token) - .then(async (res: any) => { - // Ignore if this is not the current fetch - if (currentFetchIdRef.current !== fetchId) return; - - const isResponseLike = res && typeof res.json === 'function'; - const status = isResponseLike ? res.status : undefined; - const ok = isResponseLike ? Boolean(res.ok) : true; - - let data: any = res; - if (isResponseLike) { - try { - data = await res.json(); - } catch { - data = null; - } + .getDocumentArtifact(artifactIdToFetch, token) + .then(async (res: Response) => { + if (currentFetchIdRef.current !== fetchId) return true; + if (!res.ok) return false; + let data: any = null; + try { + data = await res.json(); + } catch { + data = null; } - - // Check again after async operation - if (currentFetchIdRef.current !== fetchId) return; - - if (ok && data?.success && data?.artifact) { - setArtifact(data.artifact); - // Remember the last successful todo artifact id so we can fallback if a newer id 404s. - if (data.artifact?.artifact_type === 'todo_list') { - lastSuccessfulTodoArtifactIdRef.current = effectiveArtifactId; - } + if (currentFetchIdRef.current !== fetchId) return true; + if (data?.success && isDocumentArtifact(data.artifact)) { + setDocumentArtifact(data.artifact); + setArtifact(null); setLoading(false); - return; + return true; } - - const isTodoTool = (toolName ?? '').toLowerCase().includes('todo'); - - // If the latest todo artifact id is missing (404), fall back to the last known good one - // so the backend can still resolve `tool_id` for the todo list. - if ( - status === 404 && - isTodoTool && - lastSuccessfulTodoArtifactIdRef.current && - lastSuccessfulTodoArtifactIdRef.current !== effectiveArtifactId - ) { - // Update effectiveArtifactId to trigger a new fetch with the fallback id - setEffectiveArtifactId(lastSuccessfulTodoArtifactIdRef.current); - setLoading(false); - return; - } - - // Ensure we show a visible error state instead of rendering nothing. - const message = - data?.message || - (status === 404 ? 'Artifact not found' : null) || - 'Failed to load artifact'; - setError(message); - setLoading(false); + return false; }) - .catch((err) => { - // Ignore if this is not the current fetch + .catch(() => false) + .then((handled) => { + if (handled) return; if (currentFetchIdRef.current !== fetchId) return; - setError('Failed to fetch artifact'); - setLoading(false); + setDocumentArtifact(null); + fetchLegacyArtifact(fetchId); }); - }, [isOpen, effectiveArtifactId, token, toolName, conversationId]); + + // Legacy notes/todo fetch (unchanged behavior). + function fetchLegacyArtifact(fetchId: string) { + userService + .getArtifact(artifactIdToFetch, token) + .then(async (res: any) => { + // Ignore if this is not the current fetch + if (currentFetchIdRef.current !== fetchId) return; + + const isResponseLike = res && typeof res.json === 'function'; + const status = isResponseLike ? res.status : undefined; + const ok = isResponseLike ? Boolean(res.ok) : true; + + let data: any = res; + if (isResponseLike) { + try { + data = await res.json(); + } catch { + data = null; + } + } + + // Check again after async operation + if (currentFetchIdRef.current !== fetchId) return; + + if (ok && data?.success && data?.artifact) { + setArtifact(data.artifact); + // Remember the last successful todo artifact id so we can fallback if a newer id 404s. + if (data.artifact?.artifact_type === 'todo_list') { + lastSuccessfulTodoArtifactIdRef.current = effectiveArtifactId; + } + setLoading(false); + return; + } + + const isTodoTool = (toolName ?? '').toLowerCase().includes('todo'); + + // If the latest todo artifact id is missing (404), fall back to the last known good one + // so the backend can still resolve `tool_id` for the todo list. + if ( + status === 404 && + isTodoTool && + lastSuccessfulTodoArtifactIdRef.current && + lastSuccessfulTodoArtifactIdRef.current !== effectiveArtifactId + ) { + // Update effectiveArtifactId to trigger a new fetch with the fallback id + setEffectiveArtifactId(lastSuccessfulTodoArtifactIdRef.current); + setLoading(false); + return; + } + + // Ensure we show a visible error state instead of rendering nothing. + const message = + data?.message || + (status === 404 ? 'Artifact not found' : null) || + 'Failed to load artifact'; + setError(message); + setLoading(false); + }) + .catch(() => { + // Ignore if this is not the current fetch + if (currentFetchIdRef.current !== fetchId) return; + setError('Failed to fetch artifact'); + setLoading(false); + }); + } + }, [ + isOpen, + effectiveArtifactId, + token, + toolName, + conversationId, + refreshNonce, + ]); const renderContent = () => { if (loading) { @@ -401,6 +449,14 @@ export default function ArtifactSidebar({ ); } + if (documentArtifact) { + return ( + setRefreshNonce((n) => n + 1)} + /> + ); + } // Avoid rendering an empty panel if the artifact couldn't be loaded for any reason. if (!artifact) { return ( diff --git a/frontend/src/components/DocumentArtifactView.tsx b/frontend/src/components/DocumentArtifactView.tsx new file mode 100644 index 00000000..987d2f52 --- /dev/null +++ b/frontend/src/components/DocumentArtifactView.tsx @@ -0,0 +1,344 @@ +import mermaid from 'mermaid'; +import { Download, FileText, History, RotateCcw } from 'lucide-react'; +import { useEffect, useMemo, useRef, useState } from 'react'; +import { useSelector } from 'react-redux'; + +import userService from '../api/services/userService'; +import { useDarkTheme } from '../hooks'; +import { selectToken } from '../preferences/preferenceSlice'; +import Spinner from './Spinner'; +import { + buildPreviewDocument, + displayFilename, + findCurrentVersion, + formatBytes, + previewModeForKind, + sortVersionsDesc, + triggerResponseDownload, + type ArtifactVersion, + type DocumentArtifact, +} from './artifactViewUtils'; +import { Button } from './ui/button'; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from './ui/select'; + +function asPreviewString(spec: unknown): string { + if (typeof spec === 'string') return spec; + if (spec && typeof spec === 'object') { + const rec = spec as Record; + for (const key of ['html', 'svg', 'content', 'source', 'code']) { + if (typeof rec[key] === 'string') return rec[key] as string; + } + } + return ''; +} + +function FramePreview({ + artifactKind, + source, +}: { + artifactKind: string; + source: string; +}) { + const [isDarkTheme] = useDarkTheme(); + const [rendered, setRendered] = useState(null); + const [renderError, setRenderError] = useState(false); + const isMermaid = artifactKind.toLowerCase() === 'mermaid'; + + // `securityLevel: 'sandbox'` makes mermaid build the diagram's live DOM inside + // its OWN opaque-origin `sandbox=""` iframe (not the app document) and return + // an `', + ); + expect(doc).toContain(PREVIEW_CSP_MERMAID); + expect(PREVIEW_CSP_MERMAID).toContain('frame-src data:'); + // Scripts and origins stay forbidden on the outer document. + expect(PREVIEW_CSP_MERMAID).toContain("script-src 'none'"); + expect(PREVIEW_CSP_MERMAID).toContain("default-src 'none'"); + expect(PREVIEW_CSP_MERMAID).not.toContain('allow-same-origin'); + expect(PREVIEW_CSP_MERMAID).not.toContain('allow-scripts'); + }); +}); + +describe('filenameFromContentDisposition', () => { + it('reads the plain quoted filename', () => { + expect( + filenameFromContentDisposition('attachment; filename="report.pdf"'), + ).toBe('report.pdf'); + }); + + it('reads and decodes the RFC 5987 filename* form', () => { + expect( + filenameFromContentDisposition( + "attachment; filename*=UTF-8''my%20deck.pptx", + ), + ).toBe('my deck.pptx'); + }); + + it('returns null when absent', () => { + expect(filenameFromContentDisposition(null)).toBeNull(); + expect(filenameFromContentDisposition('attachment')).toBeNull(); + }); +}); diff --git a/frontend/src/components/artifactViewUtils.ts b/frontend/src/components/artifactViewUtils.ts new file mode 100644 index 00000000..d4c96069 --- /dev/null +++ b/frontend/src/components/artifactViewUtils.ts @@ -0,0 +1,227 @@ +/** + * Pure helpers for the document-artifact view: preview-mode mapping, version + * ordering, sandbox srcdoc/CSP construction, and download-trigger plumbing. + * Kept free of React so the logic is unit-testable in isolation. + */ + +export type ArtifactKind = + | 'document' + | 'spreadsheet' + | 'presentation' + | 'code' + | 'html' + | 'svg' + | 'mermaid' + | 'image' + | 'pdf' + | 'data' + | 'file'; + +export type ArtifactVersion = { + version: number; + mime_type: string | null; + filename: string | null; + size: number | null; + sha256: string | null; + preview_text: string | null; + produced_by: unknown; + created_at: string | null; +}; + +export type DocumentArtifact = { + id: string; + kind: string; + title: string | null; + current_version: number; + versions: ArtifactVersion[]; + spec: unknown; +}; + +/** How the sidebar should render an artifact's current version. */ +export type PreviewMode = 'frame' | 'image' | 'card' | 'text'; + +/** + * Discriminates the legacy notes/todo response (`artifact_type`) from the + * generalized document/file response (`kind`). + */ +export function isDocumentArtifact(value: unknown): value is DocumentArtifact { + if (!value || typeof value !== 'object') return false; + const rec = value as Record; + return ( + typeof rec.kind === 'string' && + typeof rec.id === 'string' && + Array.isArray(rec.versions) + ); +} + +const FRAME_KINDS = new Set(['html', 'svg', 'mermaid']); +const CARD_KINDS = new Set([ + 'document', + 'spreadsheet', + 'presentation', + 'pdf', + 'file', +]); +const TEXT_KINDS = new Set(['code', 'data']); + +/** + * Map a `kind` (with optional mime fallback) to a preview mode. SVG/PDF/image + * are also inferred from the mime type when the coarse `kind` is `html`/`file`, + * so a server that only emits the canonical kind set still previews correctly. + */ +export function previewModeForKind( + kind: string | null | undefined, + mimeType?: string | null, +): PreviewMode { + const mime = (mimeType ?? '').toLowerCase(); + + if (mime.startsWith('image/svg')) return 'frame'; + if (mime === 'application/pdf') return 'card'; + if (mime.startsWith('image/')) return 'image'; + + const k = (kind ?? '').toLowerCase(); + if (FRAME_KINDS.has(k)) return 'frame'; + if (k === 'image') return 'image'; + if (CARD_KINDS.has(k)) return 'card'; + if (TEXT_KINDS.has(k)) return 'text'; + return 'text'; +} + +/** Versions sorted newest-first by version number (does not mutate input). */ +export function sortVersionsDesc( + versions: ArtifactVersion[], +): ArtifactVersion[] { + return [...versions].sort((a, b) => b.version - a.version); +} + +/** The version row matching `current_version`, or the highest if absent. */ +export function findCurrentVersion( + artifact: Pick, +): ArtifactVersion | null { + if (!artifact.versions.length) return null; + const match = artifact.versions.find( + (v) => v.version === artifact.current_version, + ); + if (match) return match; + return sortVersionsDesc(artifact.versions)[0] ?? null; +} + +/** Human-readable byte size (e.g. 1536 -> "1.5 KB"). */ +export function formatBytes(size: number | null | undefined): string { + if (size == null || Number.isNaN(size) || size < 0) return ''; + if (size < 1024) return `${size} B`; + const units = ['KB', 'MB', 'GB', 'TB']; + let value = size / 1024; + let unit = 0; + while (value >= 1024 && unit < units.length - 1) { + value /= 1024; + unit += 1; + } + const rounded = value >= 10 ? Math.round(value) : Math.round(value * 10) / 10; + return `${rounded} ${units[unit]}`; +} + +/** Best display filename for a version, falling back to the artifact title/id. */ +export function displayFilename( + artifact: Pick, + version: Pick | null, +): string { + return version?.filename || artifact.title || `artifact-${artifact.id}`; +} + +/** + * Restrictive CSP for the preview iframe. No network, no scripts from origins, + * no plugins; inline styles/SVG allowed so generated markup renders. Used for + * html/svg markup that is embedded inline as static nodes. + */ +export const PREVIEW_CSP = + "default-src 'none'; img-src data: blob:; style-src 'unsafe-inline'; " + + "font-src data:; script-src 'none'; form-action 'none'; base-uri 'none'; " + + "frame-ancestors 'none'"; + +/** + * CSP for the mermaid preview. Mermaid's `securityLevel: 'sandbox'` output is an + * `