diff --git a/application/agents/default_tools.py b/application/agents/default_tools.py index 3932ff1e..28b0c291 100644 --- a/application/agents/default_tools.py +++ b/application/agents/default_tools.py @@ -31,7 +31,11 @@ _HEADLESS_EXCLUDED_TOOLS = frozenset({"scheduler"}) # default tools. Names may overlap with DEFAULT_CHAT_TOOLS (e.g. ``scheduler``) # — both registries share ``_DEFAULT_TOOL_NAMESPACE`` so the same uuid5 # resolves either way (the dual-flag row carries ``default`` AND ``builtin``). -BUILTIN_AGENT_TOOLS: tuple = ("scheduler", "read_document") +# ``code_executor`` is builtin-only (not default-on): it needs a running sandbox +# runner so it is opt-in per agent, but staying registered keeps its synthetic id +# resolvable (an agent that enabled it never silently loses it) and keeps it in the +# agent picker. +BUILTIN_AGENT_TOOLS: tuple = ("scheduler", "read_document", "code_executor") # Builtins shown only in the workflow-node tool picker, never the classic # agent picker. The synthesized row carries ``workflow_only`` so the frontend diff --git a/application/agents/tools/artifact_generator.py b/application/agents/tools/artifact_generator.py index 76545cca..009d5dff 100644 --- a/application/agents/tools/artifact_generator.py +++ b/application/agents/tools/artifact_generator.py @@ -687,13 +687,12 @@ class ArtifactGeneratorTool(Tool): logger.exception("artifact_generator: render failed") return {"error": f"render failed: {type(exc).__name__}: {exc}"} finally: - # Drop this render's scratch dir before tearing the session down so a - # warm/reused session doesn't accumulate per-render files on disk. + # Drop this render's scratch dir, but do NOT close the session: it is the + # shared conversation/run session that code_executor(persist=True) keeps + # warm. A render is self-contained (it builds a document from the artifact + # spec, not from prior kernel state) and does not own that session -- its + # lifecycle belongs to the manager's TTL reaper / the conversation. manager.remove_path(session_id, token_dir) - try: - manager.close(session_id) - except Exception: - logger.exception("artifact_generator: session close failed") if not data: return {"error": "renderer produced an empty file."} return {"data": data} diff --git a/application/agents/tools/code_executor.py b/application/agents/tools/code_executor.py index aeebf51c..62e4fe60 100644 --- a/application/agents/tools/code_executor.py +++ b/application/agents/tools/code_executor.py @@ -273,13 +273,28 @@ class CodeExecutorTool(Tool): if not version or not version.get("storage_path"): return {"error": f"input artifact {artifact_id} has no stored content."} + # Reject an oversize input BEFORE buffering it: the declared ``size`` + # avoids pulling a huge file into worker memory, and the bounded read + # below backstops a missing/lying size column. + max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 0) or 0) + declared_size = version.get("size") + if max_bytes and isinstance(declared_size, (int, float)) and declared_size > max_bytes: + return {"error": f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit."} + filename = safe_filename(version.get("filename") or artifact_id) try: file_obj = storage.get_file(version["storage_path"]) - data = file_obj.read() + try: + data = file_obj.read(max_bytes + 1) if max_bytes else file_obj.read() + finally: + close = getattr(file_obj, "close", None) + if callable(close): + close() except Exception: logger.exception("code_executor: failed to read input artifact bytes") return {"error": f"failed to read input artifact {artifact_id}."} + if max_bytes and len(data) > max_bytes: + return {"error": f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit."} try: manager.put_file(session_id, f"inputs/{filename}", data) except Exception: diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index 3991fedb..713a2fa8 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -81,6 +81,26 @@ class WorkflowEngine: def execute( self, initial_inputs: WorkflowState, query: str + ) -> Generator[Dict[str, str], None, None]: + """Run the workflow graph, closing the run-scoped sandbox session once when the run ends.""" + try: + yield from self._run_graph(initial_inputs, query) + finally: + # The sandbox session is keyed by the run id and shared by every code + # node and agent-node tool in this run, so it is torn down exactly once + # here rather than per node. peek_manager() never builds the manager, so + # a run that never opened a session closes nothing. + from application.sandbox.sandbox_creator import SandboxCreator + + mgr = SandboxCreator.peek_manager() + if mgr is not None: + try: + mgr.close(self._session_id()) + except Exception: + logger.exception("Workflow run failed to close its sandbox session") + + def _run_graph( + self, initial_inputs: WorkflowState, query: str ) -> Generator[Dict[str, str], None, None]: self._initialize_state(initial_inputs, query) @@ -392,37 +412,35 @@ class WorkflowEngine: timeout = self._resolve_code_timeout(config.timeout) manager = SandboxCreator.get_manager() + # The session is keyed by the run id and shared by every code node and every + # agent-node tool in this run, so it is NOT closed here: closing per node + # would cold-drop later nodes' interpreter/filesystem state. The run session + # is torn down once in ``execute``'s finally when the whole run ends. manager.open(session_id) - try: - loaded = self._materialize_code_inputs(manager, session_id, config.inputs, user_id) - # Stage prior state as DATA the node code reads with - # ``json.load(open("state.json"))`` -- e.g. ``state["decision"]``. The - # file lands at the workspace root, which is the kernel cwd, so a - # relative open resolves it. State is never templated into the program. - state_json = json.dumps(self._json_safe_state(), default=str).encode("utf-8") - manager.put_file(session_id, "state.json", state_json) - pre_signatures = snapshot_signatures(manager, session_id) - result = manager.exec(session_id, code, timeout=timeout) - if self.run_persisted: - artifacts = capture_artifacts( - manager, - session_id, - pre_signatures, - user_id=user_id, - workflow_run_id=self.workflow_run_id, - produced_by={"node_id": node.id, "node_type": NodeType.CODE.value}, - ) - else: - # No workflow_runs row backs this run (unsaved/embedded draft): an - # artifact parented to this run id would be an unreachable orphan - # (403 on get/download). Skip persistence; the sandbox still ran and - # ``_build_code_output`` handles the empty-artifacts case. - artifacts = [] - finally: - try: - manager.close(session_id) - except Exception: - logger.exception("Code node failed to close sandbox session") + loaded = self._materialize_code_inputs(manager, session_id, config.inputs, user_id) + # Stage prior state as DATA the node code reads with + # ``json.load(open("state.json"))`` -- e.g. ``state["decision"]``. The + # file lands at the workspace root, which is the kernel cwd, so a + # relative open resolves it. State is never templated into the program. + state_json = json.dumps(self._json_safe_state(), default=str).encode("utf-8") + manager.put_file(session_id, "state.json", state_json) + pre_signatures = snapshot_signatures(manager, session_id) + result = manager.exec(session_id, code, timeout=timeout) + if self.run_persisted: + artifacts = capture_artifacts( + manager, + session_id, + pre_signatures, + user_id=user_id, + workflow_run_id=self.workflow_run_id, + produced_by={"node_id": node.id, "node_type": NodeType.CODE.value}, + ) + else: + # No workflow_runs row backs this run (unsaved/embedded draft): an + # artifact parented to this run id would be an unreachable orphan + # (403 on get/download). Skip persistence; the sandbox still ran and + # ``_build_code_output`` handles the empty-artifacts case. + artifacts = [] if not result.ok: error = ( @@ -477,6 +495,7 @@ class WorkflowEngine: ) -> List[str]: """Stage referenced input artifacts (run-scoped, never cross-tenant) into the workspace.""" from application.agents.tools.artifact_ref import resolve_artifact_id + from application.core.settings import settings from application.storage.db.repositories.artifacts import ArtifactsRepository from application.storage.db.session import db_readonly from application.storage.storage_creator import StorageCreator @@ -486,6 +505,7 @@ class WorkflowEngine: raw_ids = self._resolve_input_artifact_ids(inputs) if not raw_ids: return loaded + max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 0) or 0) storage = StorageCreator.get_storage() for raw in raw_ids: with db_readonly() as conn: @@ -504,8 +524,23 @@ class WorkflowEngine: version = repo.get_version(artifact_id, artifact["current_version"]) if not version or not version.get("storage_path"): raise ValueError(f"input artifact {artifact_id} has no stored content.") + declared_size = version.get("size") + if max_bytes and isinstance(declared_size, (int, float)) and declared_size > max_bytes: + raise ValueError( + f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit." + ) filename = safe_filename(version.get("filename") or artifact_id) - data = storage.get_file(version["storage_path"]).read() + file_obj = storage.get_file(version["storage_path"]) + try: + data = file_obj.read(max_bytes + 1) if max_bytes else file_obj.read() + finally: + close = getattr(file_obj, "close", None) + if callable(close): + close() + if max_bytes and len(data) > max_bytes: + raise ValueError( + f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit." + ) manager.put_file(session_id, f"inputs/{filename}", data) loaded.append(f"inputs/{filename}") return loaded diff --git a/application/api/user/artifacts/authz.py b/application/api/user/artifacts/authz.py index a09f6bb5..90095260 100644 --- a/application/api/user/artifacts/authz.py +++ b/application/api/user/artifacts/authz.py @@ -60,23 +60,31 @@ def resolve_principal() -> Principal: return Principal() +def _shared_row_for(conn, conversation_id, share_token): + """Return the shared_conversations row iff share_token grants this conversation, else None.""" + if not share_token: + return None + shared = SharedConversationsRepository(conn).find_by_uuid(share_token) + if shared and str(shared.get("conversation_id")) == str(conversation_id): + return shared + 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``). + Conversation-level gate only. Reuses ``ConversationsRepository.get`` (owner OR + ``shared_with``) so artifact access tracks message access, and honours a valid + share token. A share-token caller reaches the conversation here but must still + be snapshot-scoped per-artifact by the caller (see ``authorize_artifact``): a + valid token does NOT imply access to every artifact in the conversation. """ 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 + return _shared_row_for(conn, conversation_id, share_token) is not None def authorize_artifact(conn, artifact: dict, principal: Principal) -> bool: @@ -85,8 +93,9 @@ def authorize_artifact(conn, artifact: dict, principal: Principal) -> bool: A low-trust agent api_key is confined to a single conversation it proves by carrying that artifact's parent ``conversation_id`` query param (owner match + matching conversation + agent scope) and never inherits share-link access. A - JWT owner, ``shared_with`` collaborator, or share-token holder is authorized by - resolving the artifact's parent (conversation or workflow run). + JWT owner or ``shared_with`` collaborator gets full access to every artifact of + the parent; a share-token holder is confined to the shared ``first_n_queries`` + snapshot (an artifact whose ``message_id`` is outside it, or NULL, is denied). """ if principal.is_agent_scoped: # An agent key is not the owner's session and is embedded in public widget @@ -108,8 +117,21 @@ def authorize_artifact(conn, artifact: dict, principal: Principal) -> bool: share_token = request.args.get("share_token") if conversation_id is not None: - return user_can_access_conversation( - conn, str(conversation_id), principal.user_id, share_token + # Owner or shared_with collaborator: full access to every artifact. + if principal.user_id and ConversationsRepository(conn).get( + str(conversation_id), principal.user_id + ) is not None: + return True + # Share-token holder: only artifacts inside the shared first_n_queries snapshot. + shared = _shared_row_for(conn, str(conversation_id), share_token) + if shared is None: + return False + message_id = artifact.get("message_id") + if not message_id: + return False + first_n = int(shared.get("first_n_queries") or 0) + return ConversationsRepository(conn).message_in_first_n( + str(conversation_id), str(message_id), first_n ) if workflow_run_id is not None: if not principal.user_id: diff --git a/application/api/user/artifacts/routes.py b/application/api/user/artifacts/routes.py index 9f9ea3fc..03e8c318 100644 --- a/application/api/user/artifacts/routes.py +++ b/application/api/user/artifacts/routes.py @@ -18,6 +18,7 @@ from flask_restx import Namespace, Resource from application.api import api from application.api.user.artifacts.authz import ( + _shared_row_for, authorize_artifact, authorize_artifact_write, resolve_principal, @@ -26,6 +27,7 @@ from application.api.user.artifacts.authz import ( 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.conversations import ConversationsRepository 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 @@ -149,6 +151,25 @@ class ListArtifacts(Resource): jsonify({"success": False, "message": "Forbidden"}), 403 ) rows = repo.list_artifacts(conversation_id=conversation_id) + # Owner / shared_with collaborator sees every artifact; a + # share-token caller is confined to the shared first_n_queries + # snapshot (drop artifacts whose message is outside it or NULL). + conv_repo = ConversationsRepository(conn) + is_owner = bool( + user_id and conv_repo.get(conversation_id, user_id) is not None + ) + if not is_owner: + shared = _shared_row_for(conn, conversation_id, share_token) + first_n = int(shared.get("first_n_queries") or 0) if shared else 0 + snapshot_ids = conv_repo.first_n_message_ids( + conversation_id, first_n + ) + rows = [ + r + for r in rows + if r.get("message_id") is not None + and str(r.get("message_id")) in snapshot_ids + ] elif workflow_run_id: run = WorkflowRunsRepository(conn).get(workflow_run_id) if run is None or run.get("user_id") != user_id: diff --git a/application/core/settings.py b/application/core/settings.py index 16723e41..88fcd108 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -251,17 +251,17 @@ class Settings(BaseSettings): # Config-free tools on by default in agentless chats. ``scheduler`` is # dual-registered (also in ``BUILTIN_AGENT_TOOLS``) so the same synthetic id - # resolves whether reached via defaults or the agent picker. ``code_executor`` - # and ``artifact_generator`` persist artifacts (not a ``user_tools``-FK table); - # their synthetic-id load is user- and conversation-scoped like ``scheduler``. - # NOTE: default-on ``code_executor`` runs LLM-authored sandboxed code WITHOUT an - # approval prompt — the sandbox is the trust boundary, so multi-tenant deployments - # should add per-tenant isolation (Daytona / gVisor / egress policy). + # resolves whether reached via defaults or the agent picker. ``artifact_generator`` + # persists artifacts (not a ``user_tools``-FK table); its synthetic-id load is + # user- and conversation-scoped like ``scheduler`` and it renders html/markdown/code + # without a sandbox runner, so it is safe to default on. ``code_executor`` is NOT + # default-on: it needs a running sandbox runner and executes LLM-authored code, so a + # fresh deploy without a runner would otherwise surface a tool that hard-fails. Enable + # it per-agent instead (opt-in via the agent tool picker). DEFAULT_CHAT_TOOLS: list = [ "memory", "read_webpage", "scheduler", - "code_executor", "artifact_generator", ] diff --git a/application/sandbox/artifacts_capture.py b/application/sandbox/artifacts_capture.py index e1c78b5a..581202c8 100644 --- a/application/sandbox/artifacts_capture.py +++ b/application/sandbox/artifacts_capture.py @@ -34,6 +34,11 @@ class QuotaExceeded(Exception): # one exec into an unbounded read+persist sweep. MAX_CAPTURED_FILES = 64 +# Cap the per-pass READ sweep independently of the persist cap: unchanged files are +# read (to detect a change) then skipped without counting toward MAX_CAPTURED_FILES, +# so a workspace full of unchanged files could otherwise be re-read in full every exec. +MAX_SCANNED_FILES = 256 + # Auto-capture skips scratch/intermediate workspace paths so install steps, extracted # archives, and temp files don't each become a downloadable artifact. Agents write # throwaway files under ``tmp/``; an explicit ``outputs`` list bypasses this skip (the @@ -99,9 +104,14 @@ def snapshot_signatures(manager: Any, session_id: str) -> Dict[str, Tuple[int, O except Exception: logger.exception("artifacts_capture: pre-exec listing failed") return signatures - for rel_path in files: - if rel_path.startswith("inputs/") or _is_scratch(rel_path): - continue + candidates = sorted(f for f in files if not f.startswith("inputs/") and not _is_scratch(f)) + if len(candidates) > MAX_SCANNED_FILES: + logger.warning( + "artifacts_capture: pre-exec signature scan capped at %d of %d files", + MAX_SCANNED_FILES, + len(candidates), + ) + for rel_path in candidates[:MAX_SCANNED_FILES]: try: data = manager.get_file(session_id, rel_path) except Exception: @@ -142,10 +152,15 @@ def capture_artifacts( else: candidates = sorted(f for f in produced if not _is_scratch(f)) captured: List[Dict[str, Any]] = [] + scanned = 0 for rel_path in candidates: if len(captured) >= MAX_CAPTURED_FILES: logger.warning("artifacts_capture: capture cap reached; remaining files skipped") break + if scanned >= MAX_SCANNED_FILES: + logger.warning("artifacts_capture: read-scan cap reached; remaining files skipped") + break + scanned += 1 try: data = manager.get_file(session_id, rel_path) except Exception: diff --git a/application/sandbox/base.py b/application/sandbox/base.py index ff52911d..5c3dceec 100644 --- a/application/sandbox/base.py +++ b/application/sandbox/base.py @@ -36,6 +36,7 @@ class ExecResult: results: List[DisplayData] = field(default_factory=list) display_data: List[DisplayData] = field(default_factory=list) plots: List[Plot] = field(default_factory=list) + truncated: bool = False # output exceeded the budget and was cut; status stays "ok" @property def ok(self) -> bool: diff --git a/application/sandbox/daytona.py b/application/sandbox/daytona.py index 100225fc..688c146c 100644 --- a/application/sandbox/daytona.py +++ b/application/sandbox/daytona.py @@ -64,6 +64,7 @@ class DaytonaSandbox(CodeSandbox): create_timeout: float = 60.0, auto_stop_interval: int = 15, auto_delete_interval: int = 60, + max_output_bytes: int = 0, max_file_bytes: int = 10 * 1024 * 1024, max_sandboxes: int = 50, ) -> None: @@ -85,6 +86,7 @@ class DaytonaSandbox(CodeSandbox): self._create_timeout = create_timeout self._auto_stop_interval = auto_stop_interval self._auto_delete_interval = auto_delete_interval + self._max_output_bytes = max_output_bytes self._max_file_bytes = max_file_bytes self._max_sandboxes = max_sandboxes self._handles: Dict[str, _Handle] = {} @@ -222,6 +224,23 @@ class DaytonaSandbox(CodeSandbox): logger.warning("Failed to start stopped Daytona sandbox %s: %s", getattr(sandbox, "id", "?"), exc) return None + def _ensure_started(self, handle: "_Handle") -> bool: + """Refresh the handle's sandbox; start it if auto-stopped. True only if it was woken.""" + try: + fresh = self._client.get(handle.sandbox_id) + except Exception as exc: # noqa: BLE001 - a failed refresh just means we don't retry + logger.warning("Daytona get for %s failed while ensuring started: %s", handle.sandbox_id, exc) + return False + state = getattr(fresh, "state", None) + state_value = getattr(state, "value", state) + if state_value == "started": + handle.sandbox = fresh + return False + if self._wake_if_stopped(fresh) is None: + return False + handle.sandbox = fresh + return True + def _create_sandbox(self, session_id: str): """Create a fresh Daytona sandbox labelled for ``session_id``.""" from daytona import CreateSandboxFromSnapshotParams @@ -349,6 +368,14 @@ class DaytonaSandbox(CodeSandbox): try: response = handle.sandbox.process.code_run(wrapped, timeout=wall) except Exception as exc: # noqa: BLE001 - any SDK/cloud error -> error result, never raise + # A cached handle may point at a sandbox Daytona auto-stopped; wake it and + # retry once. Genuine code errors return a nonzero-exit response (they do + # NOT raise), so this only retries transport/stopped faults. + if self._ensure_started(handle): + try: + return self._to_result(handle.sandbox.process.code_run(wrapped, timeout=wall)) + except Exception: # noqa: BLE001 - second failure -> error result below + pass return ExecResult( status="error", error_name=type(exc).__name__, @@ -359,17 +386,54 @@ class DaytonaSandbox(CodeSandbox): @staticmethod def _with_workspace_cwd(workspace: str, code: str) -> str: - """Prepend a chdir into the session workspace so relative paths resolve there.""" + """Prepend a chdir into the session workspace so relative paths resolve there. + + Leading ``from __future__`` imports are hoisted above the prelude so they stay + the first statements of the module (Python rejects them anywhere else). + """ + hoisted, rest = DaytonaSandbox._split_leading_future_imports(code) prelude = ( "import os as _os\n" f"_os.makedirs({workspace!r}, exist_ok=True)\n" f"_os.chdir({workspace!r})\n" ) - return prelude + code + return hoisted + prelude + rest @staticmethod - def _to_result(response) -> ExecResult: - """Map a Daytona ``ExecuteResponse`` into the shared ``ExecResult`` shape.""" + def _split_leading_future_imports(code: str) -> tuple[str, str]: + """Split leading ``from __future__`` imports (and the blank/comment lines around them) from the rest. + + A module docstring appearing BEFORE a future import is an unsupported edge (rare in + generated snippets); the common ``from __future__ import annotations`` first-line case + is handled. Everything from the first real statement onward stays in ``rest``. + """ + lines = code.splitlines(keepends=True) + saw_future = False + split_at = 0 + for i, line in enumerate(lines): + stripped = line.strip() + if stripped == "" or stripped.startswith("#"): + continue # blank/comment: part of the leading run + if stripped.startswith("from __future__ import"): + saw_future = True + continue + split_at = i # first real statement: the leading run ends here + break + else: + split_at = len(lines) # whole snippet was blank/comment/future + if not saw_future: + return "", code + hoisted = "".join(lines[:split_at]) + if hoisted and not hoisted.endswith("\n"): + hoisted += "\n" + return hoisted, "".join(lines[split_at:]) + + def _to_result(self, response) -> ExecResult: + """Map a Daytona ``ExecuteResponse`` into the shared ``ExecResult`` shape. + + Caps ``stdout`` at ``max_output_bytes`` (0 = disabled) BEFORE it is also reused as + ``error_value`` so a huge buffered response cannot propagate unbounded downstream. + """ exit_code = getattr(response, "exit_code", 0) or 0 artifacts = getattr(response, "artifacts", None) stdout = "" @@ -378,10 +442,17 @@ class DaytonaSandbox(CodeSandbox): else: stdout = getattr(response, "result", "") or "" + truncated = False + if self._max_output_bytes and len(stdout.encode("utf-8", "ignore")) > self._max_output_bytes: + stdout = stdout.encode("utf-8", "ignore")[: self._max_output_bytes].decode("utf-8", "ignore") + stdout += f"\n[output truncated at {self._max_output_bytes} bytes]" + truncated = True + result = ExecResult( status="ok" if exit_code == 0 else "error", stdout=stdout, exit_code=exit_code, + truncated=truncated, ) if exit_code != 0: result.error_name = "ExecutionError" @@ -401,31 +472,47 @@ class DaytonaSandbox(CodeSandbox): remote = self._remote_path(handle.workspace, dest_path) parent = posixpath.dirname(remote) try: - if parent and parent != handle.workspace: - try: - handle.sandbox.fs.create_folder(parent, "755") - except Exception as folder_exc: # noqa: BLE001 - folder may already exist - logger.debug("put_file parent folder create returned: %s", folder_exc) - handle.sandbox.fs.upload_file(data, remote) + self._upload(handle, remote, parent, data) except Exception as exc: # noqa: BLE001 - log detail server-side, return a generic error + # A cached handle may point at an auto-stopped sandbox; wake it and retry once. + if self._ensure_started(handle): + try: + self._upload(handle, remote, parent, data) + return + except Exception: # noqa: BLE001 - second failure -> generic IOError below + pass logger.warning("put_file failed for %r: %s", dest_path, exc) raise IOError(f"put_file failed: {type(exc).__name__}") from exc + def _upload(self, handle: "_Handle", remote: str, parent: str, data: bytes) -> None: + """Create the parent folder (best-effort) and upload ``data`` to ``remote``.""" + if parent and parent != handle.workspace: + try: + handle.sandbox.fs.create_folder(parent, "755") + except Exception as folder_exc: # noqa: BLE001 - folder may already exist + logger.debug("put_file parent folder create returned: %s", folder_exc) + handle.sandbox.fs.upload_file(data, remote) + def get_file(self, session_id: str, path: str) -> bytes: """Download ``path`` from the session workspace as bytes, capped at ``max_file_bytes``.""" handle = self._get_handle(session_id) remote = self._remote_path(handle.workspace, path) try: - info = handle.sandbox.fs.get_file_info(remote) - size = getattr(info, "size", None) - if size is not None and size > self._max_file_bytes: - raise IOError(f"file too large: {size} > {self._max_file_bytes} bytes") - data = handle.sandbox.fs.download_file(remote) + data = self._download(handle, remote) except IOError: raise except Exception as exc: # noqa: BLE001 - log detail server-side, return a generic error - logger.warning("get_file failed for %r: %s", path, exc) - raise IOError(f"get_file failed: {type(exc).__name__}") from exc + # A cached handle may point at an auto-stopped sandbox; wake it and retry once. + if not self._ensure_started(handle): + logger.warning("get_file failed for %r: %s", path, exc) + raise IOError(f"get_file failed: {type(exc).__name__}") from exc + try: + data = self._download(handle, remote) + except IOError: + raise + except Exception: # noqa: BLE001 - second failure -> generic IOError + logger.warning("get_file failed for %r: %s", path, exc) + raise IOError(f"get_file failed: {type(exc).__name__}") from exc if data is None: raise IOError(f"get_file produced no payload for {path!r}") data = data if isinstance(data, bytes) else bytes(data) @@ -435,9 +522,27 @@ class DaytonaSandbox(CodeSandbox): raise IOError(f"file too large: {len(data)} > {self._max_file_bytes} bytes") return data + def _download(self, handle: "_Handle", remote: str) -> object: + """Fetch ``remote``'s bytes, rejecting a file whose declared size exceeds the cap.""" + info = handle.sandbox.fs.get_file_info(remote) + size = getattr(info, "size", None) + if size is not None and size > self._max_file_bytes: + raise IOError(f"file too large: {size} > {self._max_file_bytes} bytes") + return handle.sandbox.fs.download_file(remote) + def list_files(self, session_id: str) -> List[str]: """List workspace-relative file paths for ``session_id`` (recursive, never escapes the workspace).""" handle = self._get_handle(session_id) + try: + return self._list_all(handle) + except Exception: # noqa: BLE001 - transport/stopped fault: wake and retry once, else re-raise + # A cached handle may point at an auto-stopped sandbox; wake it and retry once. + if self._ensure_started(handle): + return self._list_all(handle) + raise + + def _list_all(self, handle: "_Handle") -> List[str]: + """Walk the workspace subtree for ``handle`` and return workspace-relative file paths.""" out: List[str] = [] self._walk(handle.sandbox, handle.workspace, "", out) return out diff --git a/application/sandbox/jupyter_gateway.py b/application/sandbox/jupyter_gateway.py index 413a09bd..0155021f 100644 --- a/application/sandbox/jupyter_gateway.py +++ b/application/sandbox/jupyter_gateway.py @@ -62,6 +62,11 @@ class _Kernel: class JupyterKernelGatewaySandbox(CodeSandbox): """Drives one always-on Jupyter Kernel Gateway, one stateful kernel per session.""" + # Raw-byte chunk size for a staged upload: each execute_request carries one + # chunk's base64 (~4 MB at 3 MB raw), well under the gateway's 10 MiB websocket + # frame cap. The 3-byte boundary keeps every base64 block self-contained. + _PUT_CHUNK_BYTES = 3 * 1024 * 1024 + def __init__( self, gateway_url: str, @@ -167,13 +172,24 @@ class JupyterKernelGatewaySandbox(CodeSandbox): 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.""" + """Sweep the session workspace (best-effort) then delete its kernel and registry entry.""" with self._lock: kernel = self._kernels.pop(session_id, None) if kernel is None: return + self._cleanup_workspace(kernel) self._delete_kernel(kernel.kernel_id) + def _cleanup_workspace(self, kernel: _Kernel) -> None: + """Best-effort rmtree of the per-session workspace while the kernel is still alive.""" + if not kernel.initialized: + return + code = "import shutil as _sh\n" f"_sh.rmtree({kernel.workspace!r}, ignore_errors=True)\n" + try: + self._run(kernel, code, self._http_timeout) + except Exception: # noqa: BLE001 - teardown is best-effort and must never raise + logger.warning("Failed to sweep workspace for kernel %s", kernel.kernel_id, exc_info=True) + def close_handle(self, session_id: str, kernel_id: str) -> None: """Delete the SPECIFIC kernel captured at eviction time, never a re-opened one. @@ -248,10 +264,16 @@ class JupyterKernelGatewaySandbox(CodeSandbox): # 0700 on the root and the per-session dir is defense-in-depth only: every # kernel runs under one shared uid here, so this is not a cross-session # boundary (that needs distinct uids / per-session VMs -- the Daytona backend). + # A fresh kernel always starts on a clean workspace: rmtree any stale dir a + # prior kernel for the same session id left behind (else artifacts_capture + # would re-read those files every exec). Only genuine new-kernel creation + # reaches here -- open() on a live kernel returns early -- so a warm + # persist=true session is never wiped mid-computation. setup = ( - "import os as _os\n" + "import os as _os, shutil as _sh\n" f"_os.makedirs({_WORKSPACE_ROOT!r}, mode=0o700, exist_ok=True)\n" f"_os.chmod({_WORKSPACE_ROOT!r}, 0o700)\n" + f"_sh.rmtree({kernel.workspace!r}, ignore_errors=True)\n" f"_os.makedirs({kernel.workspace!r}, mode=0o700, exist_ok=True)\n" f"_os.chmod({kernel.workspace!r}, 0o700)\n" f"_os.chdir({kernel.workspace!r})\n" @@ -268,8 +290,18 @@ class JupyterKernelGatewaySandbox(CodeSandbox): 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.""" + def _run( + self, + kernel: _Kernel, + code: str, + timeout: float, + max_output_bytes: Optional[int] = None, + ) -> ExecResult: + """Execute one ``execute_request`` over the WS channel and assemble the reply. + + ``max_output_bytes`` overrides the default output budget for this call only + (file-transfer execs raise it so a multi-MB base64 payload is not truncated). + """ try: ws = websocket.create_connection( self._ws_url(kernel.kernel_id), @@ -281,7 +313,7 @@ class JupyterKernelGatewaySandbox(CodeSandbox): 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) + return self._collect(ws, msg_id, timeout, kernel.kernel_id, max_output_bytes) finally: try: ws.close() @@ -317,8 +349,16 @@ class JupyterKernelGatewaySandbox(CodeSandbox): "channel": "shell", } - def _collect(self, ws: websocket.WebSocket, msg_id: str, timeout: float, kernel_id: str) -> ExecResult: + def _collect( + self, + ws: websocket.WebSocket, + msg_id: str, + timeout: float, + kernel_id: str, + max_output_bytes: Optional[int] = None, + ) -> ExecResult: """Read iopub/shell frames until ``execute_reply``/idle, a wall-clock deadline, or a closed socket.""" + effective = max_output_bytes if max_output_bytes is not None else self._max_output_bytes result = ExecResult() stdout_parts: List[str] = [] stderr_parts: List[str] = [] @@ -362,7 +402,7 @@ class JupyterKernelGatewaySandbox(CodeSandbox): if not truncated: text = content.get("text", "") buffered += len(text.encode("utf-8", "ignore")) - if buffered > self._max_output_bytes: + if buffered > effective: truncated = True self._interrupt_and_drain(ws, msg_id, kernel_id) # runaway output: stop and drain break @@ -377,7 +417,7 @@ class JupyterKernelGatewaySandbox(CodeSandbox): # otherwise buffer unbounded. Drop and truncate once over budget. if not truncated: buffered += self._rich_payload_bytes(content) - if buffered > self._max_output_bytes: + if buffered > effective: truncated = True self._interrupt_and_drain(ws, msg_id, kernel_id) break @@ -403,7 +443,8 @@ class JupyterKernelGatewaySandbox(CodeSandbox): idle_seen = True if truncated: - stderr_parts.append(f"\n[output truncated at {self._max_output_bytes} bytes]") + stderr_parts.append(f"\n[output truncated at {effective} bytes]") + result.truncated = True # surface the cut without flipping status off "ok" result.stdout = "".join(stdout_parts) result.stderr = "".join(stderr_parts) return result @@ -446,20 +487,42 @@ class JupyterKernelGatewaySandbox(CodeSandbox): # -- File transfer --------------------------------------------------- + def _file_transfer_budget(self) -> int: + """Output budget for a file-transfer exec: max file bytes inflated by base64 plus marker slack.""" + return self._max_file_bytes * 4 // 3 + 4096 + 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.""" + """Decode ``data`` inside the kernel and write it under the session workspace. + + The upload is chunked so no single ``execute_request`` exceeds the gateway's + websocket message-size cap: the first chunk creates/truncates the file (``wb``) + and each later chunk appends (``ab``). Every chunk program re-resolves the path + -- kernel globals are not relied on to survive between execs. + """ 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}") + offset = 0 + first = True + # Loop at least once so an empty file is still created (a single wb write of b""). + while first or offset < len(data): + chunk = data[offset:offset + self._PUT_CHUNK_BYTES] + encoded = base64.b64encode(chunk).decode("ascii") + if first: + body = ( + 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" + ) + else: + body = ( + f"_p = _resolve({kernel.workspace!r}, {dest_path!r})\n" + f"_f = open(_p, 'ab'); _f.write(_b64.b64decode({encoded!r})); _f.close()\n" + ) + code = "import base64 as _b64, os as _os\n" + _CONTAINMENT_SNIPPET + body + result = self._run(kernel, code, self._default_timeout) + if not result.ok: + raise IOError(f"put_file failed: {result.error_value}") + offset += self._PUT_CHUNK_BYTES + first = False 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.""" @@ -476,7 +539,7 @@ class JupyterKernelGatewaySandbox(CodeSandbox): 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) + result = self._run(kernel, code, self._default_timeout, max_output_bytes=self._file_transfer_budget()) if not result.ok: raise IOError(f"get_file failed: {result.error_value}") out = result.stdout @@ -503,7 +566,7 @@ class JupyterKernelGatewaySandbox(CodeSandbox): " _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) + result = self._run(kernel, code, self._default_timeout, max_output_bytes=self._file_transfer_budget()) if not result.ok: raise IOError(f"list_files failed: {result.error_value}") out = result.stdout diff --git a/application/sandbox/manager.py b/application/sandbox/manager.py index a88c5bf2..2ddafa54 100644 --- a/application/sandbox/manager.py +++ b/application/sandbox/manager.py @@ -22,7 +22,8 @@ class _Session: ``handle`` is the backend handle id returned by ``backend.open``; ``None`` while a slot is RESERVED (a placeholder occupying a cap slot during a cold backend open that runs outside the lock). ``ready`` is False for such a placeholder so reuse/reap/evict - skip it until the backend open finalizes it. + skip it until the backend open finalizes it. ``pending_close`` marks a session whose + ``close`` arrived while an op held it ``in_use``; the last ``_leave`` then tears it down. """ session_id: str @@ -32,6 +33,7 @@ class _Session: in_use: int = field(default=0) handle: Optional[str] = field(default=None) ready: bool = field(default=False) + pending_close: bool = field(default=False) def is_expired(self, now: float) -> bool: """True when the session has been idle longer than its (clamped) TTL.""" @@ -256,10 +258,22 @@ class SandboxManager: The backend ``close`` runs OUTSIDE the lock (it may be slow network I/O) and is keyed by the captured handle so it tears down only the resource this session owned, never one a concurrent re-open created. + + Teardown is DEFERRED when an op holds the session ``in_use`` (a concurrent + exec/put_file/get_file): the session is flagged ``pending_close`` and the last + ``_leave`` performs the actual close, so a concurrent close never kills an + in-flight exec (which would lose its captured files). The common path -- the + caller's own exec has already ``_left``, so ``in_use == 0`` -- stays synchronous. """ with self._lock: - session = self._sessions.pop(session_id, None) - handle = session.handle if session is not None else None + session = self._sessions.get(session_id) + if session is None: + return + if session.in_use > 0: + session.pending_close = True + return + self._sessions.pop(session_id, None) + handle = session.handle self._close_backend(session_id, handle) def has_session(self, session_id: str) -> bool: @@ -288,12 +302,25 @@ class SandboxManager: session.in_use += 1 def _leave(self, session_id: str) -> None: - """Release an in-use hold taken by ``_enter`` (idempotent if the session was closed).""" + """Release an in-use hold taken by ``_enter``; run a close deferred by ``close`` on the last release. + + Idempotent if the session was already closed. When the final hold is released and + a ``close`` was deferred (``pending_close``), the session is popped here and its + backend torn down OUTSIDE the lock, keyed by the captured handle. + """ + handle: Optional[str] = None + do_close = False with self._lock: session = self._sessions.get(session_id) if session is not None and session.in_use > 0: session.in_use -= 1 session.last_access = time.monotonic() + if session.in_use == 0 and session.pending_close: + self._sessions.pop(session_id, None) + handle = session.handle + do_close = True + if do_close: + self._close_backend(session_id, handle) def _close_backend(self, session_id: str, handle: Optional[str]) -> None: """Close the SPECIFIC backend resource captured for this session, best-effort. diff --git a/application/sandbox/sandbox_creator.py b/application/sandbox/sandbox_creator.py index 425a9e62..2d5c458d 100644 --- a/application/sandbox/sandbox_creator.py +++ b/application/sandbox/sandbox_creator.py @@ -1,6 +1,6 @@ """Factory + process-wide singleton selecting a sandbox backend from settings.""" -from typing import Callable, Dict +from typing import Callable, Dict, Optional from application.core.settings import settings from application.sandbox.base import CodeSandbox @@ -41,6 +41,10 @@ def _make_daytona() -> CodeSandbox: create_timeout=float(settings.SANDBOX_HTTP_TIMEOUT) * 6, auto_stop_interval=int(settings.DAYTONA_AUTO_STOP_INTERVAL), auto_delete_interval=auto_delete_interval, + # Bounds only what the app propagates from a response: the Daytona SDK still + # buffers the whole HTTP response before we see it, so a full in-sandbox cap + # would need runner-side support. This is a partial but real mitigation. + max_output_bytes=int(settings.SANDBOX_MAX_OUTPUT_BYTES), max_file_bytes=int(settings.SANDBOX_MAX_FILE_BYTES), max_sandboxes=int(settings.DAYTONA_MAX_SANDBOXES), ) @@ -68,6 +72,11 @@ class SandboxCreator: ) return cls._instance + @classmethod + def peek_manager(cls) -> Optional[SandboxManager]: + """Return the process-wide manager if already built, else None (never constructs one).""" + return cls._instance + @classmethod def create_backend(cls, type_name: str) -> CodeSandbox: """Instantiate the backend registered under ``type_name`` (case-insensitive).""" diff --git a/application/storage/db/repositories/conversations.py b/application/storage/db/repositories/conversations.py index 65787310..dc652c99 100644 --- a/application/storage/db/repositories/conversations.py +++ b/application/storage/db/repositories/conversations.py @@ -607,6 +607,44 @@ class ConversationsRepository: ) return [_message_row_to_dict(r) for r in result.fetchall()] + def first_n_message_ids(self, conversation_id: str, first_n: int) -> set[str]: + """Return the ids of the first ``first_n`` messages (by position) of the conversation. + + Same ordering a share snapshot uses (``get_messages``' ``position ASC`` + then ``[:first_n]``); empty set for ``first_n <= 0`` or a non-UUID id. + """ + if first_n <= 0 or not looks_like_uuid(conversation_id): + return set() + result = self._conn.execute( + text( + "SELECT id FROM conversation_messages " + "WHERE conversation_id = CAST(:cid AS uuid) " + "ORDER BY position ASC LIMIT :n" + ), + {"cid": conversation_id, "n": int(first_n)}, + ) + return {str(r[0]) for r in result.fetchall()} + + def message_in_first_n(self, conversation_id: str, message_id: str, first_n: int) -> bool: + """True if message_id is among the first ``first_n`` messages (by position) of the conversation.""" + if not message_id or first_n <= 0: + return False + # Shape-gate both ids: a non-UUID reaching CAST(... AS uuid) would raise + # and poison the enclosing transaction (see ``rename``). + if not looks_like_uuid(conversation_id) or not looks_like_uuid(message_id): + return False + row = self._conn.execute( + text( + "SELECT 1 FROM ( " + " SELECT id FROM conversation_messages " + " WHERE conversation_id = CAST(:cid AS uuid) " + " ORDER BY position ASC LIMIT :n " + ") t WHERE t.id = CAST(:mid AS uuid)" + ), + {"cid": conversation_id, "mid": message_id, "n": int(first_n)}, + ).fetchone() + return row is not None + def get_message_at(self, conversation_id: str, position: int) -> Optional[dict]: # Shape-gate: see ``rename``. Callers today always pass a resolved # UUID (via ``get_any`` first), but the guard costs nothing and diff --git a/deployment/docker-compose-azure.yaml b/deployment/docker-compose-azure.yaml index bbcd9f44..07bf9cdd 100644 --- a/deployment/docker-compose-azure.yaml +++ b/deployment/docker-compose-azure.yaml @@ -25,10 +25,14 @@ services: ports: - "7091:7091" networks: - # `default` reaches the sandbox gateway; `data-net` reaches redis/postgres - # (which are NOT on the sandbox's networks). + # `default` reaches the internet (LLM APIs) + frontend/host; `data-net` + # reaches redis/postgres; `sandbox-net` reaches the code-exec runner over an + # internal-only net (so the runner needs no shared internet bridge to be + # reachable, and the egress overlay can cut its internet without severing + # this control path). - default - data-net + - sandbox-net volumes: - ../application/indexes:/app/application/indexes - ../application/inputs:/app/application/inputs @@ -54,23 +58,31 @@ services: - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 - SANDBOX_KERNEL_NAME=docsgpt-python networks: + # See backend: control-plane access to the code-exec runner over sandbox-net + # (workflow code nodes run in the worker). - default - data-net + - sandbox-net depends_on: redis: condition: service_started postgres: condition: service_healthy - # Always-on code-execution runner (Jupyter Kernel Gateway). In-process kernels, - # no Docker socket, no published host port. On the internal-only sandbox-net - # plus `default` (for internet egress); deliberately NOT on data-net, so - # sandboxed code cannot reach the unauthenticated redis broker (Celery task - # injection -> RCE) or postgres. Do NOT add `env_file: ../.env` here -- the - # runner needs no app secrets. For untrusted multi-tenant workloads use a - # per-session VM via SANDBOX_BACKEND=daytona instead. + # Opt-in code-execution runner (Jupyter Kernel Gateway). Started ONLY with + # `docker compose --profile sandbox up`: the `code_executor` tool is off by + # default (removed from DEFAULT_CHAT_TOOLS) and must be enabled per-agent in + # the agent tool picker, so a plain `up` does not start this service. + # In-process kernels, no Docker socket, no published host port. On the + # internal-only sandbox-net plus `default` (for internet egress); deliberately + # NOT on data-net, so sandboxed code cannot reach the unauthenticated redis + # broker (Celery task injection -> RCE) or postgres. Do NOT add + # `env_file: ../.env` here -- the runner needs no app secrets. For untrusted + # multi-tenant workloads use a per-session VM via SANDBOX_BACKEND=daytona + # instead. docsgpt-sandbox: build: ./sandbox + profiles: ["sandbox"] mem_limit: ${SANDBOX_MEMORY:-1g} cpus: ${SANDBOX_CPUS:-1.0} pids_limit: 256 @@ -82,8 +94,12 @@ services: tmpfs: - /tmp networks: + # Reachable by backend/worker over the internal sandbox-net; internet + # egress (runtime pip install, etc.) via the dedicated sandbox-egress net. + # Deliberately NOT on `default`: the egress overlay cuts internet by + # flipping sandbox-egress to internal, without severing the control path. - sandbox-net - - default + - sandbox-egress redis: image: redis:6-alpine @@ -92,7 +108,9 @@ services: networks: - data-net ports: - - 6379:6379 + # Loopback only: the published port must NOT be reachable from the sandbox + # via the host/bridge gateway. Host-local dev tools still use localhost. + - 127.0.0.1:6379:6379 postgres: image: postgres:16-alpine @@ -104,7 +122,8 @@ services: - POSTGRES_PASSWORD=docsgpt - POSTGRES_DB=docsgpt ports: - - "5432:5432" + # Loopback only (see redis): not reachable from the sandbox via the gateway. + - "127.0.0.1:5432:5432" volumes: - postgres_data:/var/lib/postgresql/data healthcheck: @@ -114,10 +133,16 @@ services: retries: 10 networks: - # Internal-only runner network (no external egress via this net; the sandbox - # still reaches the internet via the default bridge). + # Control plane between backend/worker and the code-exec runner. internal:true + # => no internet route on this net; the runner egresses via sandbox-egress. sandbox-net: internal: true + # The runner's outbound internet route (runtime pip install, etc.). Kept + # internet-facing so the base stack works out of the box; + # deployment/optional/docker-compose.optional.sandbox-egress.yaml flips it to + # internal:true and forces egress through a deny-private proxy for full SSRF + # containment (blocking RFC1918 / link-local / cloud-metadata). + sandbox-egress: {} # Internal data plane: redis/postgres + backend/worker. The sandbox is NOT # attached, so sandboxed code cannot reach the broker or database. data-net: diff --git a/deployment/docker-compose-hub.yaml b/deployment/docker-compose-hub.yaml index 54551023..7d950473 100644 --- a/deployment/docker-compose-hub.yaml +++ b/deployment/docker-compose-hub.yaml @@ -30,10 +30,14 @@ services: ports: - "7091:7091" networks: - # `default` reaches the sandbox gateway; `data-net` reaches redis/postgres - # (which are NOT on the sandbox's networks). + # `default` reaches the internet (LLM APIs) + frontend/host; `data-net` + # reaches redis/postgres; `sandbox-net` reaches the code-exec runner over an + # internal-only net (so the runner needs no shared internet bridge to be + # reachable, and the egress overlay can cut its internet without severing + # this control path). - default - data-net + - sandbox-net volumes: - ../application/indexes:/app/indexes - ../application/inputs:/app/inputs @@ -61,8 +65,11 @@ services: - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 - SANDBOX_KERNEL_NAME=docsgpt-python networks: + # See backend: control-plane access to the code-exec runner over sandbox-net + # (workflow code nodes run in the worker). - default - data-net + - sandbox-net volumes: - ../application/indexes:/app/indexes - ../application/inputs:/app/inputs @@ -73,18 +80,23 @@ services: postgres: condition: service_healthy - # Always-on code-execution runner (Jupyter Kernel Gateway). In-process kernels, - # no Docker socket, no published host port. On the internal-only sandbox-net - # plus `default` (for internet egress); deliberately NOT on data-net, so - # sandboxed code cannot reach the unauthenticated redis broker (Celery task - # injection -> RCE) or postgres. Do NOT add `env_file: ../.env` here -- the - # runner needs no app secrets. For untrusted multi-tenant workloads use a - # per-session VM via SANDBOX_BACKEND=daytona instead. + # Opt-in code-execution runner (Jupyter Kernel Gateway). Started ONLY with + # `docker compose --profile sandbox up`: the `code_executor` tool is off by + # default (removed from DEFAULT_CHAT_TOOLS) and must be enabled per-agent in + # the agent tool picker, so a plain `up` does not start this service. + # In-process kernels, no Docker socket, no published host port. On the + # internal-only sandbox-net plus `default` (for internet egress); deliberately + # NOT on data-net, so sandboxed code cannot reach the unauthenticated redis + # broker (Celery task injection -> RCE) or postgres. Do NOT add + # `env_file: ../.env` here -- the runner needs no app secrets. For untrusted + # multi-tenant workloads use a per-session VM via SANDBOX_BACKEND=daytona + # instead. docsgpt-sandbox: # Built from the repo (no published arc53/docsgpt-sandbox image exists yet); the # hub compose ships alongside ./sandbox, so this stays self-contained and the # stack always starts. Switch to a published image once CI builds/pushes one. build: ./sandbox + profiles: ["sandbox"] mem_limit: ${SANDBOX_MEMORY:-1g} cpus: ${SANDBOX_CPUS:-1.0} pids_limit: 256 @@ -96,8 +108,12 @@ services: tmpfs: - /tmp networks: + # Reachable by backend/worker over the internal sandbox-net; internet + # egress (runtime pip install, etc.) via the dedicated sandbox-egress net. + # Deliberately NOT on `default`: the egress overlay cuts internet by + # flipping sandbox-egress to internal, without severing the control path. - sandbox-net - - default + - sandbox-egress redis: image: redis:6-alpine @@ -106,7 +122,9 @@ services: networks: - data-net ports: - - 6379:6379 + # Loopback only: the published port must NOT be reachable from the sandbox + # via the host/bridge gateway. Host-local dev tools still use localhost. + - 127.0.0.1:6379:6379 postgres: image: postgres:16-alpine @@ -118,7 +136,8 @@ services: - POSTGRES_PASSWORD=docsgpt - POSTGRES_DB=docsgpt ports: - - "5432:5432" + # Loopback only (see redis): not reachable from the sandbox via the gateway. + - "127.0.0.1:5432:5432" volumes: - postgres_data:/var/lib/postgresql/data healthcheck: @@ -128,10 +147,16 @@ services: retries: 10 networks: - # Internal-only runner network (no external egress via this net; the sandbox - # still reaches the internet via the default bridge). + # Control plane between backend/worker and the code-exec runner. internal:true + # => no internet route on this net; the runner egresses via sandbox-egress. sandbox-net: internal: true + # The runner's outbound internet route (runtime pip install, etc.). Kept + # internet-facing so the base stack works out of the box; + # deployment/optional/docker-compose.optional.sandbox-egress.yaml flips it to + # internal:true and forces egress through a deny-private proxy for full SSRF + # containment (blocking RFC1918 / link-local / cloud-metadata). + sandbox-egress: {} # Internal data plane: redis/postgres + backend/worker. The sandbox is NOT # attached, so sandboxed code cannot reach the broker or database. data-net: diff --git a/deployment/docker-compose.yaml b/deployment/docker-compose.yaml index 5c8cf648..193f156c 100644 --- a/deployment/docker-compose.yaml +++ b/deployment/docker-compose.yaml @@ -33,10 +33,14 @@ services: ports: - "7091:7091" networks: - # `default` reaches the internet + the sandbox gateway; `data-net` reaches - # redis/postgres (which are NOT on the sandbox's networks). + # `default` reaches the internet (LLM APIs) + frontend/host; `data-net` + # reaches redis/postgres; `sandbox-net` reaches the code-exec runner over an + # internal-only net (so the runner needs no shared internet bridge to be + # reachable, and the egress overlay can cut its internet without severing + # this control path). - default - data-net + - sandbox-net volumes: - ../application/indexes:/app/indexes - ../application/inputs:/app/inputs @@ -56,8 +60,11 @@ services: # deployment/sandbox/README.md). command: celery -A application.app.celery worker -l INFO -B -Q docsgpt,parsing networks: + # See backend: control-plane access to the code-exec runner over sandbox-net + # (workflow code nodes run in the worker). - default - data-net + - sandbox-net env_file: - ../.env environment: @@ -80,12 +87,15 @@ 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. + # Opt-in code-execution runner (Jupyter Kernel Gateway). Started ONLY with + # `docker compose --profile sandbox up`: the `code_executor` tool is off by + # default (removed from DEFAULT_CHAT_TOOLS) and must be enabled per-agent in + # the agent tool picker, so a plain `up` does not start this service. + # 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. # # SINGLE TRUST DOMAIN: all sessions share this one container/uid and are # isolated by working directory only (per-session cwd) — not by a kernel/OS @@ -96,6 +106,7 @@ services: # per-session VM via SANDBOX_BACKEND=daytona instead. docsgpt-sandbox: build: ./sandbox + profiles: ["sandbox"] mem_limit: ${SANDBOX_MEMORY:-1g} cpus: ${SANDBOX_CPUS:-1.0} pids_limit: 256 @@ -109,8 +120,12 @@ services: # runtime files live on tmpfs; the root FS is read-only everywhere else. - /tmp networks: + # Reachable by backend/worker over the internal sandbox-net; internet + # egress (runtime pip install, etc.) via the dedicated sandbox-egress net. + # Deliberately NOT on `default`: the egress overlay cuts internet by + # flipping sandbox-egress to internal, without severing the control path. - sandbox-net - - default + - sandbox-egress redis: image: redis:6-alpine @@ -119,7 +134,9 @@ services: networks: - data-net ports: - - 6379:6379 + # Loopback only: the published port must NOT be reachable from the sandbox + # via the host/bridge gateway. Host-local dev tools still use localhost. + - 127.0.0.1:6379:6379 postgres: image: postgres:16-alpine @@ -131,7 +148,8 @@ services: - POSTGRES_PASSWORD=docsgpt - POSTGRES_DB=docsgpt ports: - - "5432:5432" + # Loopback only (see redis): not reachable from the sandbox via the gateway. + - "127.0.0.1:5432:5432" volumes: - postgres_data:/var/lib/postgresql/data healthcheck: @@ -141,15 +159,20 @@ services: 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). + # Control plane between backend/worker and the code-exec runner. internal:true + # => no internet route on this net; the runner egresses via sandbox-egress. sandbox-net: internal: true + # The runner's outbound internet route (runtime pip install, etc.). Kept + # internet-facing so the base stack works out of the box; + # deployment/optional/docker-compose.optional.sandbox-egress.yaml flips it to + # internal:true and forces egress through a deny-private proxy for full SSRF + # containment (blocking RFC1918 / link-local / cloud-metadata). + sandbox-egress: {} # Internal-only data plane. redis (broker/cache) and postgres live here with # backend/worker, but the code-exec sandbox is NOT attached, so arbitrary # sandboxed code cannot reach the unauthenticated broker (Celery task - # injection -> worker RCE) or the database. For full internet-egress - # containment additionally apply deployment/optional/docker-compose.optional.sandbox-egress.yaml. + # injection -> worker RCE) or the database. data-net: internal: true diff --git a/deployment/optional/docker-compose.optional.sandbox-egress.yaml b/deployment/optional/docker-compose.optional.sandbox-egress.yaml index bc0916a9..4af85004 100644 --- a/deployment/optional/docker-compose.optional.sandbox-egress.yaml +++ b/deployment/optional/docker-compose.optional.sandbox-egress.yaml @@ -2,40 +2,52 @@ # # Docker Compose cannot express L3 egress filtering the way a Kubernetes # NetworkPolicy can, so SSRF containment in compose deployments is delivered by -# routing the sandbox's outbound traffic through an egress-gateway sidecar that -# DENIES private / link-local / metadata ranges and ALLOWS the public internet. +# taking away the runner's direct internet route and forcing its outbound +# traffic through an egress-gateway sidecar that DENIES private / link-local / +# metadata ranges and ALLOWS the public internet. +# +# HOW THIS OVERLAY WORKS +# The base stack puts docsgpt-sandbox on two networks: the internal `sandbox-net` +# (control plane to backend/worker) and `sandbox-egress` (its internet route). +# This overlay flips `sandbox-egress` to `internal: true`, which removes the +# runner's direct route to the internet / host / RFC1918 / metadata entirely -- +# so even raw sockets in arbitrary sandbox code have nowhere to go except the +# proxy on `sandbox-net`. The proxy is the ONLY container with an internet route +# (on its own `sandbox-egress-out` network) and its ACL denies private +# destinations. We flip a network's scalar property by KEY (a well-defined +# Compose merge); we do NOT try to remove an item from the service's `networks` +# LIST, which Compose cannot express (it unions lists) -- the reason a naive +# `networks: [sandbox-net]` override silently leaves the runner on its old net. # # Two viable approaches; pick one: # -# (1) Host / cloud firewall (simplest, recommended for single-host compose): -# Drop egress to RFC1918 (10/8, 172.16/12, 192.168/16), link-local +# (1) Host / cloud firewall (simplest, belt-and-suspenders for single-host): +# Also drop egress to RFC1918 (10/8, 172.16/12, 192.168/16), link-local # (169.254/16, incl. the 169.254.169.254 metadata IP), and ULA on the # docsgpt-sandbox container's interface using iptables/nftables on the -# Docker host (Docker does not do this for you). Example (host root): -# SBX=$(docker inspect -f '{{.NetworkSettings.Networks.docsgpt_sandbox-net.IPAddress}}' docsgpt-sandbox) +# Docker host. Example (host root): +# SBX=$(docker inspect -f '{{.NetworkSettings.Networks.docsgpt-oss_sandbox-egress.IPAddress}}' docsgpt-sandbox) # iptables -I DOCKER-USER -s "$SBX" -d 169.254.0.0/16 -j DROP # iptables -I DOCKER-USER -s "$SBX" -d 10.0.0.0/8 -j DROP # iptables -I DOCKER-USER -s "$SBX" -d 172.16.0.0/12 -j DROP # iptables -I DOCKER-USER -s "$SBX" -d 192.168.0.0/16 -j DROP # (Allow established/return traffic and DNS as needed for your setup.) # -# (2) Egress-gateway sidecar (this overlay): pin the sandbox to an -# internal-only network with NO direct internet route, and force all of -# its outbound traffic through a small forward proxy that blocks private -# destinations. The proxy is the ONLY container on both the internal and -# the egress network, so the sandbox cannot reach internal services -# directly. Configure the runner's outbound HTTP(S) via the proxy env -# vars below; lock down non-HTTP egress with the proxy's own rules. +# (2) Egress-gateway sidecar (this overlay): the runner has NO direct internet +# route (see above); its only path out is the deny-private forward proxy. +# Point the runner's HTTP(S) client at the proxy via the env vars below. # # Apply alongside the base stack: # docker compose -f deployment/docker-compose.yaml \ # -f deployment/optional/docker-compose.optional.sandbox-egress.yaml up -d +# (also works layered on docker-compose-hub.yaml / docker-compose-azure.yaml, +# which share the same service/network names). # -# NOTE: a forward proxy only constrains traffic that honors the proxy env vars. -# Truly arbitrary sandbox code can ignore them, so on a multi-tenant or -# untrusted deployment prefer approach (1) (or the Kubernetes NetworkPolicy) -# which enforces at the network layer regardless of what the code does. -version: "3.8" +# NOTE: because `sandbox-egress` is internal here, non-HTTP raw egress is blocked +# at L3 regardless of what the code does (there is simply no route). HTTP(S) that +# honors the proxy env is filtered by the proxy ACL. For belt-and-suspenders on a +# multi-tenant or untrusted deployment, also apply approach (1) or the Kubernetes +# NetworkPolicy. services: # Forward proxy that denies private/link-local/metadata destinations and @@ -46,26 +58,28 @@ services: image: ghcr.io/example/egress-deny-private:latest # replace with your proxy image restart: unless-stopped networks: - - sandbox-net # reachable by the sandbox - - sandbox-egress # the only container with an internet route + - sandbox-net # reachable by the runner (control-plane net) + - sandbox-egress-out # the ONLY container with an internet route docsgpt-sandbox: - # Cut the sandbox off from the default (internet) bridge; it keeps only the - # internal sandbox-net and must egress via the proxy. - networks: - - sandbox-net + # Match the base service's profile so this overlay fragment merges onto the + # profiled base and the runner still starts only under `--profile sandbox` + # (Compose unions service fragments across -f files; keeping the profile + # here means a bare `up` never starts the runner via this overlay either). + profiles: ["sandbox"] + # No `networks:` override here on purpose: the base already attaches the + # runner to sandbox-net + sandbox-egress. Flipping sandbox-egress to internal + # (below) removes its direct internet route; its only way out is this proxy. environment: - HTTP_PROXY=http://sandbox-egress-proxy:8080 - HTTPS_PROXY=http://sandbox-egress-proxy:8080 - NO_PROXY=localhost,127.0.0.1 networks: - # Internal-only network shared by the sandbox and the proxy. ``internal: - # true`` removes the default gateway so the sandbox has NO direct internet - # route -- its only path out is through the proxy on the egress network. - sandbox-net: - driver: bridge - internal: true - # Internet-facing network for the proxy only. + # Flip the runner's egress net to internal: no direct route to the internet, + # host, RFC1918, or cloud metadata. (Merged by key over the base definition.) sandbox-egress: + internal: true + # Internet-facing route for the proxy ONLY. + sandbox-egress-out: driver: bridge diff --git a/deployment/sandbox/README.md b/deployment/sandbox/README.md index 1200ba34..bb63d8e9 100644 --- a/deployment/sandbox/README.md +++ b/deployment/sandbox/README.md @@ -1,10 +1,35 @@ # docsgpt-sandbox runner -Always-on Jupyter Kernel Gateway that executes sandboxed LLM code. The DocsGPT +Opt-in 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. +## Enabling code execution (opt-in) + +The runner is **opt-in**. `code_executor` is no longer a default chat tool (it +was removed from `DEFAULT_CHAT_TOOLS`), so a plain `docker compose up` does +**not** start `docsgpt-sandbox`. Start it explicitly with the `sandbox` +profile: + +```bash +docker compose -f deployment/docker-compose.yaml --profile sandbox up +``` + +With the egress-firewall overlay (see *Network egress / SSRF* below): + +```bash +docker compose \ + -f deployment/docker-compose.yaml \ + -f deployment/optional/docker-compose.optional.sandbox-egress.yaml \ + --profile sandbox up +``` + +Then enable `code_executor` **per-agent** in the agent tool picker — it is not a +default chat tool. Agents without it never call the runner, and the +backend/worker degrade gracefully when the runner is absent. The `-hub` and +`-azure` compose variants gate the runner behind the same `sandbox` profile. + ## Isolation model Read this before pointing untrusted or multi-tenant workloads at the runner. @@ -102,7 +127,9 @@ 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 +an internal-only network and is gated behind the `sandbox` Compose profile +(opt-in — start it with `docker compose --profile sandbox up`; see *Enabling +code execution (opt-in)* above). The backend and worker reach it at `http://docsgpt-sandbox:8888` and select the scrubbing kernel by setting `SANDBOX_KERNEL_NAME=docsgpt-python` (the runner only ships the kernelspec; the app chooses it). The same applies to k8s: `SANDBOX_KERNEL_NAME=docsgpt-python` @@ -194,14 +221,19 @@ The hardened container runs **without `NET_ADMIN`**, so it cannot self-apply ``` - **docker-compose** — compose cannot express L3 egress filtering natively. The - base stack puts the runner on an `internal: true` network (no host port), but - that does not by itself block the metadata IP or RFC1918 reachable via the - default bridge. Add a **host/cloud firewall rule** (drop the four private - ranges on the sandbox container's interface) **or** route egress through an - **egress-gateway proxy** sidecar. Both are documented in - [`deployment/optional/docker-compose.optional.sandbox-egress.yaml`](../optional/docker-compose.optional.sandbox-egress.yaml). - On untrusted/multi-tenant hosts prefer the host-firewall rule — a forward - proxy only constrains code that honors `HTTP(S)_PROXY`. + base stack reaches the runner over an `internal: true` control network + (`sandbox-net`, no host port) and gives it internet egress on a dedicated + `sandbox-egress` bridge — but that bridge does not by itself block the metadata + IP or RFC1918. Apply + [`deployment/optional/docker-compose.optional.sandbox-egress.yaml`](../optional/docker-compose.optional.sandbox-egress.yaml), + which flips `sandbox-egress` to `internal: true` (removing the runner's direct + internet/RFC1918/metadata route entirely) and forces egress through a + deny-private **egress-gateway proxy** sidecar; for belt-and-suspenders on + untrusted/multi-tenant hosts also add a **host/cloud firewall rule** (drop the + four private ranges on the sandbox container's interface), since a forward + proxy only filters code that honors `HTTP(S)_PROXY` — the `internal` flip is + what contains raw sockets. Note the broker/DB published ports are bound to + `127.0.0.1` so the runner cannot reach them via the host gateway either. ## Other hardening (deployment-level) diff --git a/tests/agents/test_default_tools.py b/tests/agents/test_default_tools.py index 65aaff13..c61b937a 100644 --- a/tests/agents/test_default_tools.py +++ b/tests/agents/test_default_tools.py @@ -176,10 +176,21 @@ class TestValidation: def test_shipped_defaults_validate(self): # The real shipped DEFAULT_CHAT_TOOLS must pass startup validation. + # code_executor is NOT shipped default-on (it needs a running sandbox runner); + # artifact_generator stays default-on and usable. usable = default_tools.validate_default_chat_tools() - assert "code_executor" in usable + assert "code_executor" not in default_tools.settings.DEFAULT_CHAT_TOOLS assert "artifact_generator" in usable + def test_code_executor_not_a_shipped_default_but_artifact_generator_is(self): + # code_executor is enabled per-agent (needs a sandbox runner); it must not be + # advertised on a fresh deploy. artifact_generator stays a shipped default. + assert "code_executor" not in default_tools.settings.DEFAULT_CHAT_TOOLS + assert "artifact_generator" in default_tools.settings.DEFAULT_CHAT_TOOLS + names = {r["name"] for r in default_tools.synthesized_default_tools(None)} + assert "code_executor" not in names + assert "artifact_generator" in names + def test_tool_with_required_config_is_rejected(self, monkeypatch): # ``brave`` needs an API key. monkeypatch.setattr( @@ -328,10 +339,12 @@ class TestResolveToolById: assert row["builtin"] is True assert row["default"] is True - @pytest.mark.parametrize("name", ["code_executor", "artifact_generator"]) + @pytest.mark.parametrize("name", ["artifact_generator"]) def test_sandbox_default_id_resolves_in_memory(self, name): - # Synthetic default id -> name -> in-memory row (loaded user-scoped at - # execute time via the synthetic-default path, like scheduler). + # artifact_generator stays a default chat tool: its synthetic default id -> + # name -> in-memory row (loaded user-scoped at execute time, like scheduler). + # code_executor is no longer default-on, so its synthetic id no longer + # resolves via this path (see open_notes on the per-agent reachability gap). tool_id = default_tools.default_tool_id(name) assert default_tools.default_tool_name_for_id(tool_id) == name row = default_tools.resolve_tool_by_id(tool_id, "user-x") @@ -421,6 +434,25 @@ class TestBuiltinAgentTools: row = default_tools.synthesize_builtin_agent_tool("scheduler") assert row["workflow_only"] is False + def test_code_executor_is_agent_selectable_builtin(self): + # Off by default (removed from DEFAULT_CHAT_TOOLS) but still reachable: a + # non-workflow-only builtin, so an agent can enable it, it stays in the + # picker, and its synthetic id resolves (no silent drop for agents that + # already had it enabled). + assert "code_executor" not in default_tools.settings.DEFAULT_CHAT_TOOLS + assert "code_executor" in default_tools.BUILTIN_AGENT_TOOLS + assert "code_executor" not in default_tools.WORKFLOW_ONLY_BUILTINS + row = default_tools.synthesize_builtin_agent_tool("code_executor") + assert row is not None + assert row["builtin"] is True and row["default"] is False + assert row["workflow_only"] is False + names = {r["name"] for r in default_tools.builtin_agent_tools_for_management()} + assert "code_executor" in names + resolved = default_tools.resolve_tool_by_id( + default_tools.default_tool_id("code_executor"), "user-1" + ) + assert resolved is not None and resolved["name"] == "code_executor" + def test_builtin_management_marks_workflow_only(self): rows = default_tools.builtin_agent_tools_for_management() by_name = {r["name"]: r for r in rows} diff --git a/tests/agents/test_workflow_code_node.py b/tests/agents/test_workflow_code_node.py index 00840a8f..5f693961 100644 --- a/tests/agents/test_workflow_code_node.py +++ b/tests/agents/test_workflow_code_node.py @@ -15,6 +15,7 @@ from application.agents.workflows.cel_evaluator import evaluate_cel from application.agents.workflows.schemas import ( NodeType, Workflow, + WorkflowEdge, WorkflowGraph, WorkflowNode, ) @@ -140,10 +141,12 @@ def test_code_node_writes_artifact_reference_into_state(patch_sandbox): assert engine.state["node_code_1_output"] == ref assert engine.state["report"] == ref assert all(not isinstance(v, (bytes, bytearray)) for v in engine.state["report"].values()) - # The sandbox session is bound to the run id and closed after the run. + # The sandbox session is bound to the run id. It is NOT closed per node: the + # session is shared across nodes and torn down once at end of the run (in + # WorkflowEngine.execute), so a code node leaves it open. manager = patch_sandbox["manager_holder"]["manager"] assert manager.opened == ["11111111-1111-1111-1111-111111111111"] - assert manager.closed == ["11111111-1111-1111-1111-111111111111"] + assert manager.closed == [] def test_code_node_no_artifacts_still_writes_status(patch_sandbox): @@ -184,10 +187,49 @@ def test_code_node_skips_capture_when_run_not_persisted(patch_sandbox): assert patch_sandbox["capture_calls"] == 0 assert engine.state["out"] == {"artifacts": [], "status": "ok"} - # The sandbox session still opened and closed (only persistence was skipped). + # The sandbox session still opened (only persistence was skipped) and is left open + # for the run to reap -- the node never closes the shared run session. manager = patch_sandbox["manager_holder"]["manager"] assert manager.opened == ["11111111-1111-1111-1111-111111111111"] - assert manager.closed == ["11111111-1111-1111-1111-111111111111"] + assert manager.closed == [] + + +def test_execute_closes_run_session_once_at_end(monkeypatch): + """The run-scoped sandbox session opened by a code node is closed once when execute() ends.""" + manager = _FakeManager(_Result(ok=True, stdout="ok")) + monkeypatch.setattr( + "application.sandbox.sandbox_creator.SandboxCreator.get_manager", lambda: manager + ) + monkeypatch.setattr( + "application.sandbox.sandbox_creator.SandboxCreator.peek_manager", lambda: manager + ) + monkeypatch.setattr( + "application.sandbox.artifacts_capture.snapshot_signatures", lambda *a, **k: {} + ) + monkeypatch.setattr( + "application.sandbox.artifacts_capture.capture_artifacts", lambda *a, **k: [] + ) + + start = WorkflowNode( + id="start_1", workflow_id="wf-1", type=NodeType.START, title="Start", + position={"x": 0, "y": 0}, config={}, + ) + code = _code_node(node_id="code_1", output_variable="out", code="print('x')") + edge = WorkflowEdge(id="e1", workflow_id="wf-1", source="start_1", target="code_1") + graph = WorkflowGraph(workflow=Workflow(name="Close Once"), nodes=[start, code], edges=[edge]) + agent = SimpleNamespace( + endpoint="stream", llm_name="openai", model_id="gpt-4o-mini", api_key="test-key", + chat_history=[], user="user-code", decoded_token={"sub": "user-code"}, + ) + sid = "11111111-1111-1111-1111-111111111111" + engine = WorkflowEngine(graph, agent, workflow_run_id=sid) + + list(engine.execute({}, "q")) + + # The code node opened the run session and left it open; execute() closed it + # exactly once at the end of the run (not once per node). + assert manager.opened == [sid] + assert manager.closed == [sid] def test_code_node_reads_prior_state_from_state_json(patch_sandbox): @@ -308,6 +350,52 @@ def test_resolve_input_artifact_ids_from_state_refs_and_raw(): assert ids == ["art-from-ref", "art-raw-id", "not_a_ref"] +def test_materialize_code_inputs_rejects_oversize(monkeypatch): + """A code-node input whose declared version ``size`` exceeds the cap raises before staging.""" + from contextlib import contextmanager + + from application.core.settings import settings + + monkeypatch.setattr(settings, "SANDBOX_MAX_INPUT_BYTES", 100, raising=False) + monkeypatch.setattr( + "application.agents.tools.artifact_ref.resolve_artifact_id", + lambda repo, raw, **k: str(raw), + ) + + class _Repo: + def __init__(self, conn): + pass + + def get_artifact_in_parent(self, artifact_id, *, workflow_run_id=None, conversation_id=None): + return {"id": artifact_id, "current_version": 1} + + def get_version(self, artifact_id, version): + return {"filename": "big.csv", "size": 10_000, "storage_path": "p/big.csv"} + + @contextmanager + def _readonly(): + yield object() + + class _Storage: + def get_file(self, path): + raise AssertionError("bytes must not be read when declared size exceeds the cap") + + monkeypatch.setattr( + "application.storage.db.repositories.artifacts.ArtifactsRepository", _Repo + ) + monkeypatch.setattr("application.storage.db.session.db_readonly", _readonly) + monkeypatch.setattr( + "application.storage.storage_creator.StorageCreator.get_storage", + staticmethod(lambda: _Storage()), + ) + + engine = _engine() + manager = _FakeManager(_Result(ok=True, stdout="ok")) + with pytest.raises(ValueError, match="exceeds"): + engine._materialize_code_inputs(manager, engine._session_id(), ["art-raw-id"], "user-code") + assert manager.put_files == [] # nothing staged + + # --------------------------------------------------------------------------- # Pass-by-reference: survives serialization + CEL branches on the metadata. # --------------------------------------------------------------------------- diff --git a/tests/agents/tools/test_artifact_generator_unit.py b/tests/agents/tools/test_artifact_generator_unit.py index b2abe0a9..6769d893 100644 --- a/tests/agents/tools/test_artifact_generator_unit.py +++ b/tests/agents/tools/test_artifact_generator_unit.py @@ -10,6 +10,7 @@ from __future__ import annotations import json import os import tempfile +from types import SimpleNamespace import pytest @@ -69,6 +70,57 @@ def test_create_rejects_unknown_kind(): assert "unsupported kind" in out["error"] +# --------------------------------------------------------------------------- +# _render session lifecycle: cleans its scratch dir, leaves the shared session open +# --------------------------------------------------------------------------- + + +class _FakeRenderManager: + """Records remove_path/close and returns fixed render bytes; no real sandbox.""" + + def __init__(self): + self.closed = [] + self.removed = [] + + def open(self, session_id, ttl=None): + return session_id + + def put_file(self, session_id, dest_path, data): + pass + + def exec(self, session_id, code, timeout=None): + return SimpleNamespace(ok=True) + + def get_file(self, session_id, path): + return b"%PDF-1.4 rendered" + + def remove_path(self, session_id, path): + self.removed.append((session_id, path)) + + def close(self, session_id): + self.closed.append(session_id) + + +def test_render_cleans_scratch_but_leaves_session_open(monkeypatch): + """_render drops its per-render scratch dir but must NOT close the shared session.""" + manager = _FakeRenderManager() + monkeypatch.setattr( + "application.sandbox.sandbox_creator.SandboxCreator.get_manager", lambda: manager + ) + + out = _tool()._render("pdf", {"title": "t", "blocks": []}) + + assert out == {"data": b"%PDF-1.4 rendered"} + # The render owns only its scratch dir; the session is the shared conversation + # session that code_executor(persist=True) keeps warm, so it is left for the + # manager/conversation to reap. + assert manager.closed == [] + assert len(manager.removed) == 1 + session_id, path = manager.removed[0] + assert session_id == "conv-1" + assert path.startswith("artifacts/") + + # --------------------------------------------------------------------------- # RFC 7386 JSON merge-patch # --------------------------------------------------------------------------- diff --git a/tests/api/user/test_artifacts_routes.py b/tests/api/user/test_artifacts_routes.py index e7d34c2f..23efd0b1 100644 --- a/tests/api/user/test_artifacts_routes.py +++ b/tests/api/user/test_artifacts_routes.py @@ -65,6 +65,12 @@ def _make_agent_conversation(conn, agent_id, user_id=OWNER): ) +def _make_message(conn, conversation_id, prompt="q", response="a"): + return ConversationsRepository(conn).append_message( + conversation_id, {"prompt": prompt, "response": response} + ) + + def _wire_api_key(monkeypatch, conn): """Point ``resolve_principal``'s own readonly conn at the test conn. @@ -289,15 +295,18 @@ class TestSharedAccess: from application.api.user.artifacts.routes import DownloadArtifact conv = _make_conversation(_patch_db) + # Attach to the first message so it falls inside the first_n_queries snapshot. + msg = _make_message(_patch_db, str(conv["id"])) art = _make_artifact( _patch_db, conversation_id=str(conv["id"]), + message_id=str(msg["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 + str(conv["id"]), OWNER, first_n_queries=1 ) storage = _FakeStorage(b"PDFDATA") @@ -320,6 +329,160 @@ class TestSharedAccess: assert resp.data == b"PDFDATA" +# --------------------------------------------------------------------------- +# Share-token snapshot scoping (first_n_queries) +# --------------------------------------------------------------------------- +@pytest.mark.unit +class TestShareTokenSnapshotScope: + """A share link exposes only artifacts within its first_n_queries snapshot.""" + + def _seed(self, conn, first_n=1): + conv = _make_conversation(conn) + conv_id = str(conv["id"]) + m0 = _make_message(conn, conv_id, prompt="q0") # position 0 (in snapshot) + m1 = _make_message(conn, conv_id, prompt="q1") # position 1 (outside) + in_art = _make_artifact( + conn, conversation_id=conv_id, message_id=str(m0["id"]), + title="in", filename="in.pdf", + storage_path="inputs/owner/artifacts/in/v1/in.pdf", + ) + out_art = _make_artifact( + conn, conversation_id=conv_id, message_id=str(m1["id"]), + title="out", filename="out.pdf", + storage_path="inputs/owner/artifacts/out/v1/out.pdf", + ) + null_art = _make_artifact( + conn, conversation_id=conv_id, message_id=None, title="null", + filename="null.pdf", + storage_path="inputs/owner/artifacts/null/v1/null.pdf", + ) + share = SharedConversationsRepository(conn).create( + conv_id, OWNER, first_n_queries=first_n + ) + return conv_id, in_art, out_art, null_art, str(share["uuid"]) + + @staticmethod + def _mock_storage(monkeypatch, data=b"BYTES"): + storage = _FakeStorage(data) + 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, + ) + return storage + + def test_share_token_list_only_snapshot(self, _patch_db, flask_app): + from application.api.user.artifacts.routes import ListArtifacts + + conv_id, in_art, out_art, null_art, token = self._seed(_patch_db) + resp = _call( + flask_app, ListArtifacts, token=None, + query={"conversation_id": conv_id, "share_token": token}, + ) + assert resp.status_code == 200 + ids = {a["id"] for a in resp.json["artifacts"]} + assert str(in_art["id"]) in ids + assert str(out_art["id"]) not in ids # outside first_n_queries + assert str(null_art["id"]) not in ids # NULL message_id -> not in snapshot + + def test_share_token_get_in_snapshot_200(self, _patch_db, flask_app): + from application.api.user.artifacts.routes import GetArtifact + + _, in_art, _out, _null, token = self._seed(_patch_db) + resp = _call( + flask_app, GetArtifact, in_art["id"], token=None, + query={"share_token": token}, + ) + assert resp.status_code == 200 + + def test_share_token_get_out_of_snapshot_403(self, _patch_db, flask_app): + from application.api.user.artifacts.routes import GetArtifact + + _, _in, out_art, _null, token = self._seed(_patch_db) + resp = _call( + flask_app, GetArtifact, out_art["id"], token=None, + query={"share_token": token}, + ) + assert resp.status_code == 403 + + def test_share_token_null_message_id_denied(self, _patch_db, flask_app): + from application.api.user.artifacts.routes import GetArtifact + + _, _in, _out, null_art, token = self._seed(_patch_db) + resp = _call( + flask_app, GetArtifact, null_art["id"], token=None, + query={"share_token": token}, + ) + assert resp.status_code == 403 + + def test_share_token_download_in_snapshot_200( + self, _patch_db, flask_app, monkeypatch + ): + from application.api.user.artifacts.routes import DownloadArtifact + + _, in_art, _out, _null, token = self._seed(_patch_db) + self._mock_storage(monkeypatch, b"INDATA") + resp = _call( + flask_app, DownloadArtifact, in_art["id"], token=None, + query={"share_token": token}, + ) + assert resp.status_code == 200 + assert resp.data == b"INDATA" + + def test_share_token_download_out_of_snapshot_403( + self, _patch_db, flask_app, monkeypatch + ): + from application.api.user.artifacts.routes import DownloadArtifact + + _, _in, out_art, _null, token = self._seed(_patch_db) + self._mock_storage(monkeypatch, b"OUTDATA") + resp = _call( + flask_app, DownloadArtifact, out_art["id"], token=None, + query={"share_token": token}, + ) + assert resp.status_code == 403 + + def test_owner_sees_all_artifacts(self, _patch_db, flask_app, token_owner): + # Snapshot scoping is share-token-only: the owner lists every artifact and + # can fetch one attached to a message outside the first_n_queries snapshot. + from application.api.user.artifacts.routes import GetArtifact, ListArtifacts + + conv_id, in_art, out_art, null_art, _token = self._seed(_patch_db) + listed = _call( + flask_app, ListArtifacts, token=token_owner, + query={"conversation_id": conv_id}, + ) + assert listed.status_code == 200 + ids = {a["id"] for a in listed.json["artifacts"]} + assert {str(in_art["id"]), str(out_art["id"]), str(null_art["id"])} <= ids + + got = _call(flask_app, GetArtifact, out_art["id"], token=token_owner) + assert got.status_code == 200 + + def test_shared_with_collaborator_sees_all_artifacts(self, _patch_db, flask_app): + # A read-only shared_with collaborator (JWT) is not snapshot-scoped either. + from application.api.user.artifacts.routes import GetArtifact, ListArtifacts + + conv_id, in_art, out_art, null_art, _token = self._seed(_patch_db) + ConversationsRepository(_patch_db).add_shared_user(conv_id, SHARED_USER) + + listed = _call( + flask_app, ListArtifacts, token={"sub": SHARED_USER}, + query={"conversation_id": conv_id}, + ) + assert listed.status_code == 200 + ids = {a["id"] for a in listed.json["artifacts"]} + assert {str(in_art["id"]), str(out_art["id"]), str(null_art["id"])} <= ids + + got = _call( + flask_app, GetArtifact, out_art["id"], token={"sub": SHARED_USER} + ) + assert got.status_code == 200 + + # --------------------------------------------------------------------------- # Download # --------------------------------------------------------------------------- diff --git a/tests/sandbox/test_artifacts_capture.py b/tests/sandbox/test_artifacts_capture.py index ef65f487..f519b1df 100644 --- a/tests/sandbox/test_artifacts_capture.py +++ b/tests/sandbox/test_artifacts_capture.py @@ -91,3 +91,39 @@ class TestCaptureFiltering: assert self._captured(monkeypatch, {"report.pdf": b"x"}, pre=pre) == [] # Content change is captured. assert self._captured(monkeypatch, {"report.pdf": b"xy"}, pre=pre) == ["report.pdf"] + + +class _CountingMgr: + """Serves a fixed {rel_path: bytes} workspace and counts every get_file read.""" + + def __init__(self, files): + self._files = files + self.reads = 0 + + def list_files(self, _sid): + return list(self._files) + + def get_file(self, _sid, path): + self.reads += 1 + return self._files[path] + + +@pytest.mark.unit +class TestReadSweepCap: + def test_snapshot_signature_scan_capped(self): + # An unchanged-file-heavy workspace can no longer be read in full each pass. + files = {f"f{i:04d}.txt": b"x" for i in range(ac.MAX_SCANNED_FILES + 50)} + mgr = _CountingMgr(files) + sigs = ac.snapshot_signatures(mgr, "sid") + assert mgr.reads == ac.MAX_SCANNED_FILES + assert len(sigs) == ac.MAX_SCANNED_FILES + + def test_capture_read_sweep_capped(self, monkeypatch): + # Every get_file (even for unchanged, never-persisted files) counts toward the cap. + files = {f"f{i:04d}.txt": b"x" for i in range(ac.MAX_SCANNED_FILES + 50)} + pre = {name: (1, hashlib.sha256(b"x").hexdigest()) for name in files} # all unchanged + mgr = _CountingMgr(files) + monkeypatch.setattr(ac, "persist_artifact", lambda *a, **k: None) + captured = ac.capture_artifacts(mgr, "sid", pre, user_id="u") + assert captured == [] # nothing changed -> nothing persisted + assert mgr.reads == ac.MAX_SCANNED_FILES # but the sweep is still bounded diff --git a/tests/sandbox/test_daytona.py b/tests/sandbox/test_daytona.py index a4675523..57137ead 100644 --- a/tests/sandbox/test_daytona.py +++ b/tests/sandbox/test_daytona.py @@ -564,6 +564,156 @@ def test_list_files_error_wrapped_as_ioerror(sandbox): assert "api.daytona" not in str(exc.value) +# --- Output cap in _to_result -------------------------------------------- + + +def test_to_result_truncates_stdout_over_cap(fake_sdk): + """Stdout beyond max_output_bytes is byte-capped, noted, and flagged truncated.""" + from application.sandbox.daytona import DaytonaSandbox + + s = DaytonaSandbox(api_key="k", max_output_bytes=10) + resp = _FakeExecuteResponse(exit_code=0, artifacts=_FakeArtifacts(stdout="X" * 100)) + res = s._to_result(resp) + assert res.truncated is True + prefix = res.stdout.split("\n[output truncated")[0] + assert prefix == "X" * 10 # capped to exactly the byte budget before the note + assert "[output truncated at 10 bytes]" in res.stdout + + +def test_to_result_keeps_small_stdout_intact(fake_sdk): + from application.sandbox.daytona import DaytonaSandbox + + s = DaytonaSandbox(api_key="k", max_output_bytes=1000) + res = s._to_result(_FakeExecuteResponse(exit_code=0, artifacts=_FakeArtifacts(stdout="hello"))) + assert res.truncated is False + assert res.stdout == "hello" + + +def test_to_result_cap_disabled_by_default(fake_sdk): + """max_output_bytes defaults to 0 (disabled): a large stdout is passed through whole.""" + from application.sandbox.daytona import DaytonaSandbox + + s = DaytonaSandbox(api_key="k") + res = s._to_result(_FakeExecuteResponse(exit_code=0, artifacts=_FakeArtifacts(stdout="Y" * 5000))) + assert res.truncated is False + assert res.stdout == "Y" * 5000 + + +def test_to_result_truncation_bounds_error_value(fake_sdk): + """On a nonzero exit the capped stdout (not the raw buffer) is what feeds error_value.""" + from application.sandbox.daytona import DaytonaSandbox + + s = DaytonaSandbox(api_key="k", max_output_bytes=10) + res = s._to_result(_FakeExecuteResponse(exit_code=1, artifacts=_FakeArtifacts(stdout="E" * 100))) + assert res.status == "error" and res.truncated is True + assert res.error_value == res.stdout + assert res.error_value.split("\n[output truncated")[0] == "E" * 10 + + +# --- Wake an auto-stopped sandbox on the cached-handle path --------------- + + +def test_exec_wakes_auto_stopped_sandbox_and_retries(sandbox): + """A cached handle to an auto-stopped sandbox: exec wakes it and retries once.""" + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + created.state = "stopped" # Daytona auto-stopped it under the still-cached handle + created.process.code_run.side_effect = [ + RuntimeError("sandbox is stopped"), + _FakeExecuteResponse(exit_code=0, artifacts=_FakeArtifacts(stdout="woke\n")), + ] + res = sandbox.exec("conv-1", "print('x')") + assert res.ok and res.stdout == "woke\n" + sandbox._client.start.assert_called_once() + assert created.process.code_run.call_count == 2 + + +def test_exec_does_not_retry_when_started(sandbox): + """A started sandbox that raises is a genuine transport fault: no wake, no spurious retry.""" + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + created.process.code_run.side_effect = RuntimeError("transport blip") + res = sandbox.exec("conv-1", "print('x')") + assert res.status == "error" and res.error_name == "RuntimeError" + sandbox._client.start.assert_not_called() + assert created.process.code_run.call_count == 1 + + +def test_exec_no_retry_when_wake_fails(sandbox): + """If the stopped sandbox cannot be started, exec returns the error without retrying.""" + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + created.state = "stopped" + sandbox._client.start.side_effect = RuntimeError("cannot start") + created.process.code_run.side_effect = RuntimeError("stopped") + res = sandbox.exec("conv-1", "print('x')") + assert res.status == "error" + assert created.process.code_run.call_count == 1 + + +def test_put_file_wakes_auto_stopped_sandbox_and_retries(sandbox): + """put_file wakes an auto-stopped sandbox and retries the upload once.""" + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + created.state = "stopped" + created.fs.upload_file.side_effect = [RuntimeError("stopped"), None] + sandbox.put_file("conv-1", "data.csv", b"x") + sandbox._client.start.assert_called_once() + assert created.fs.upload_file.call_count == 2 + + +def test_ensure_started_returns_false_when_get_fails(sandbox): + """A failed refresh means we cannot wake, so no retry is attempted.""" + sandbox.open("conv-1") + handle = sandbox._handles["conv-1"] + sandbox._client.get = mock.Mock(side_effect=KeyError("gone")) + assert sandbox._ensure_started(handle) is False + + +# --- from __future__ imports survive the workspace prelude --------------- + + +def test_with_workspace_cwd_hoists_future_import(sandbox): + """A leading ``from __future__`` import stays the first statement of the wrapped code.""" + wrapped = sandbox._with_workspace_cwd("/ws", "from __future__ import annotations\nprint('hi')\n") + assert wrapped.startswith("from __future__ import annotations\n") + assert wrapped.index("from __future__") < wrapped.index("import os as _os") + assert "_os.chdir('/ws')" in wrapped + assert wrapped.rstrip().endswith("print('hi')") + + +def test_with_workspace_cwd_no_future_import_is_prelude_prefix(sandbox): + """Without a future import the wrapper is just prelude + code, unchanged.""" + wrapped = sandbox._with_workspace_cwd("/ws", "print('hi')\n") + assert wrapped.startswith("import os as _os\n") + assert wrapped.endswith("print('hi')\n") + + +def test_split_leading_future_imports_multiple_with_comment(): + from application.sandbox.daytona import DaytonaSandbox + + code = "# header\nfrom __future__ import annotations\nfrom __future__ import division\nx = 1\n" + hoisted, rest = DaytonaSandbox._split_leading_future_imports(code) + assert "from __future__ import annotations" in hoisted + assert "from __future__ import division" in hoisted + assert rest == "x = 1\n" + + +def test_split_leading_future_imports_none(): + from application.sandbox.daytona import DaytonaSandbox + + hoisted, rest = DaytonaSandbox._split_leading_future_imports("x = 1\n") + assert hoisted == "" and rest == "x = 1\n" + + +def test_split_leading_future_imports_adds_trailing_newline(): + from application.sandbox.daytona import DaytonaSandbox + + hoisted, rest = DaytonaSandbox._split_leading_future_imports("from __future__ import annotations") + assert hoisted.endswith("\n") # ensures the prelude begins on its own line + assert rest == "" + + # --- Registry wiring ----------------------------------------------------- diff --git a/tests/sandbox/test_jupyter_gateway_isolation.py b/tests/sandbox/test_jupyter_gateway_isolation.py index 90808074..a23a0607 100644 --- a/tests/sandbox/test_jupyter_gateway_isolation.py +++ b/tests/sandbox/test_jupyter_gateway_isolation.py @@ -185,6 +185,7 @@ def test_collect_caps_oversize_rich_output(monkeypatch): result = sb._collect(_FakeWS(frames), msg_id, timeout=5, kernel_id="k1") assert result.results == [] # over-budget bundle dropped, not materialized assert "[output truncated" in result.stderr + assert result.truncated is True # truncation is surfaced explicitly, not hidden behind "ok" def test_collect_keeps_small_rich_output(monkeypatch): @@ -199,6 +200,7 @@ def test_collect_keeps_small_rich_output(monkeypatch): result = sb._collect(_FakeWS(frames), msg_id, timeout=5, kernel_id="k2") assert len(result.results) == 1 assert "[output truncated" not in (result.stderr or "") + assert result.truncated is False # nothing was cut # -- open() is idempotent under concurrency ----------------------------------- @@ -244,3 +246,110 @@ def test_concurrent_open_creates_one_kernel(monkeypatch): assert posts["n"] == 1, "concurrent open() created more than one kernel" assert results[0] == results[1] # both callers got the same kernel id + + +# -- Chunked put_file + large get_file round-trip ------------------------------ + + +def _inprocess_run(calls): + """A ``_run`` stand-in that execs the kernel program in-process and captures its stdout.""" + import io as _io + from contextlib import redirect_stdout + + def fake_run(_kernel, code, _timeout, max_output_bytes=None): + calls.append(code) + buf = _io.StringIO() + try: + with redirect_stdout(buf): + exec(compile(code, "", "exec"), {}) + except Exception as exc: # surface as an error ExecResult, like the real gateway + return ExecResult(status="error", error_name=type(exc).__name__, error_value=str(exc), exit_code=-1) + return ExecResult(status="ok", stdout=buf.getvalue(), exit_code=0) + + return fake_run + + +def test_put_file_chunks_large_upload_and_get_file_round_trips(tmp_path, monkeypatch): + """A >4 MB file is staged in multiple chunks and read back byte-for-byte. + + A 7 MB file's base64 (~9.3 MB) exceeds the default 8 MB output cap, so the + file-transfer budget on get_file is what keeps it from truncating. + """ + sb = JupyterKernelGatewaySandbox(gateway_url="http://unused", max_file_bytes=10 * 1024 * 1024) + workspace = tmp_path / "ws" + workspace.mkdir() + sb._kernels["sid"] = _Kernel("kid", str(workspace)) + + put_calls: list = [] + monkeypatch.setattr(sb, "_run", _inprocess_run(put_calls)) + + data = os.urandom(7 * 1024 * 1024 + 123) # 7 MB -> 3 chunks (3 + 3 + ~1) + sb.put_file("sid", "big.bin", data) + + # Chunking actually happened (one execute_request could not carry this). + assert len(put_calls) == 3 + assert (workspace / "big.bin").read_bytes() == data + # get_file reassembles the same bytes and passes its length+sha256 integrity check. + assert sb.get_file("sid", "big.bin") == data + + +def test_put_file_writes_empty_file(tmp_path, monkeypatch): + """An empty upload still creates the file with a single wb write.""" + sb = JupyterKernelGatewaySandbox(gateway_url="http://unused") + workspace = tmp_path / "ws" + workspace.mkdir() + sb._kernels["sid"] = _Kernel("kid", str(workspace)) + calls: list = [] + monkeypatch.setattr(sb, "_run", _inprocess_run(calls)) + + sb.put_file("sid", "empty.bin", b"") + assert len(calls) == 1 + assert (workspace / "empty.bin").read_bytes() == b"" + assert sb.get_file("sid", "empty.bin") == b"" + + +def test_collect_file_transfer_budget_prevents_truncation(monkeypatch): + """A payload over the default output cap but under the file-transfer budget is not truncated.""" + sb = JupyterKernelGatewaySandbox( + gateway_url="http://unused", max_output_bytes=1000, max_file_bytes=5000 + ) + monkeypatch.setattr(sb, "_interrupt_and_drain", lambda *a, **k: None) + big = "Z" * 4000 # > max_output_bytes (1000), < _file_transfer_budget() + assert 1000 < len(big) < sb._file_transfer_budget() + + # Default budget -> truncated (the get_file END marker would be dropped here). + default = sb._collect( + _FakeWS([_frame("m", "stream", {"name": "stdout", "text": big})]), + "m", timeout=5, kernel_id="k", + ) + assert "[output truncated" in default.stderr + + # Raised file-transfer budget -> the full payload is kept intact. + ok = sb._collect( + _FakeWS([ + _frame("m", "stream", {"name": "stdout", "text": big}), + _frame("m", "execute_reply", {"status": "ok"}), + _frame("m", "status", {"execution_state": "idle"}), + ]), + "m", timeout=5, kernel_id="k", max_output_bytes=sb._file_transfer_budget(), + ) + assert "[output truncated" not in (ok.stderr or "") + assert ok.stdout == big + + +def test_prime_sweeps_stale_workspace_before_recreate(tmp_path, monkeypatch): + """A new kernel's _prime rmtrees any stale per-session dir so old files never carry over.""" + root = tmp_path / "docsgpt-sandbox" + monkeypatch.setattr(jupyter_gateway, "_WORKSPACE_ROOT", str(root)) + workspace = root / "conv-stale" + workspace.mkdir(parents=True) + (workspace / "stale.txt").write_text("old") + + sb = JupyterKernelGatewaySandbox(gateway_url="http://unused") + kernel = _Kernel("kid", str(workspace)) + monkeypatch.setattr(sb, "_run", lambda _k, code, _t: (_exec_setup_in_tmp(code), ExecResult(status="ok"))[1]) + monkeypatch.chdir(tmp_path) + sb._prime(kernel) + + assert kernel.initialized + assert not (workspace / "stale.txt").exists() # stale file swept on re-open diff --git a/tests/sandbox/test_sandbox_manager.py b/tests/sandbox/test_sandbox_manager.py index 7aaa1a4f..82fcb8a7 100644 --- a/tests/sandbox/test_sandbox_manager.py +++ b/tests/sandbox/test_sandbox_manager.py @@ -191,6 +191,17 @@ def test_sandbox_creator_manager_is_singleton(): SandboxCreator.reset() +def test_sandbox_creator_peek_manager_never_builds(): + from application.sandbox.sandbox_creator import SandboxCreator + + SandboxCreator.reset() + assert SandboxCreator.peek_manager() is None # nothing built yet -> None, no construction + built = SandboxCreator.get_manager() + assert SandboxCreator.peek_manager() is built # returns the existing singleton + SandboxCreator.reset() + assert SandboxCreator.peek_manager() is None + + # --------------------------------------------------------------------------- # Concurrent-session cap # --------------------------------------------------------------------------- @@ -563,3 +574,80 @@ def test_evict_then_concurrent_reopen_closes_old_handle_not_new(): # The new V is usable (its workspace was not torn down by the stale close). mgr.put_file("V", "f.txt", b"data") assert mgr.get_file("V", "f.txt") == b"data" + + +# --------------------------------------------------------------------------- +# Deferred close: a close while an op is in-use must not kill the in-flight op +# --------------------------------------------------------------------------- + + +def test_close_defers_while_in_use_then_tears_down_on_leave(backend): + """close() during an in-flight op defers teardown; the last _leave performs it.""" + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("conv-1") + mgr._enter("conv-1") # simulate a concurrent exec holding the session + + mgr.close("conv-1") # must DEFER, not tear the session out from under the op + assert mgr.has_session("conv-1") + assert backend.torn_down == [] + + mgr._leave("conv-1") # last release performs the deferred close + assert not mgr.has_session("conv-1") + assert backend.torn_down == ["conv-1"] + + +def test_close_defers_until_last_of_several_holds_releases(backend): + """With multiple in-use holds, the deferred close fires only on the final _leave.""" + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("conv-1") + mgr._enter("conv-1") + mgr._enter("conv-1") + mgr.close("conv-1") + + mgr._leave("conv-1") # one hold remains -> still deferred + assert mgr.has_session("conv-1") + assert backend.torn_down == [] + + mgr._leave("conv-1") # last hold -> deferred close runs + assert not mgr.has_session("conv-1") + assert backend.torn_down == ["conv-1"] + + +def test_close_is_synchronous_when_not_in_use(backend): + """The common path (in_use == 0 at close time) stays synchronous as before.""" + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("conv-1") + mgr.close("conv-1") # caller's own exec already _left -> immediate teardown + assert not mgr.has_session("conv-1") + assert backend.torn_down == ["conv-1"] + + +def test_exec_completes_despite_concurrent_close(): + """A close racing an in-flight exec never turns it into a KernelDiedError/lost files.""" + barrier = threading.Event() + release = threading.Event() + + class _BlockingExecBackend(FakeBackend): + def exec(self, session_id, code, timeout=None): + barrier.set() + assert release.wait(timeout=5) + return ExecResult(status="ok", stdout=f"ran:{code}") + + backend = _BlockingExecBackend() + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("conv-1") + + out: Dict[str, ExecResult] = {} + worker = threading.Thread(target=lambda: out.__setitem__("r", mgr.exec("conv-1", "1+1"))) + worker.start() + assert barrier.wait(timeout=5) # exec is in flight, holding the session in-use + + mgr.close("conv-1") # concurrent close must defer, not tear down the live exec + assert backend.torn_down == [] + + release.set() + worker.join(timeout=5) + assert out["r"].ok and out["r"].stdout == "ran:1+1" # exec survived intact + # The deferred close ran on the exec's own _leave. + assert not mgr.has_session("conv-1") + assert backend.torn_down == ["conv-1"] diff --git a/tests/storage/db/repositories/test_conversations.py b/tests/storage/db/repositories/test_conversations.py index ddcb2135..5ca3475d 100644 --- a/tests/storage/db/repositories/test_conversations.py +++ b/tests/storage/db/repositories/test_conversations.py @@ -2,8 +2,10 @@ from __future__ import annotations +import uuid from datetime import datetime, timezone +from sqlalchemy import text from application.storage.db.repositories.conversations import ( ConversationsRepository, @@ -222,6 +224,52 @@ class TestGetMessages: assert repo.get_message_at(conv["id"], 99) is None +class TestFirstNSnapshot: + """``message_in_first_n`` / ``first_n_message_ids`` scope a share snapshot + to the first ``first_n`` messages by position (matches ``messages[:first_n]``).""" + + def test_message_in_first_n_membership_by_position(self, pg_conn): + repo = _repo(pg_conn) + cid = repo.create("user-1", "c")["id"] + m0 = repo.append_message(cid, {"prompt": "q0", "response": "a0"}) + m1 = repo.append_message(cid, {"prompt": "q1", "response": "a1"}) + m2 = repo.append_message(cid, {"prompt": "q2", "response": "a2"}) + assert repo.message_in_first_n(cid, m0["id"], 2) is True + assert repo.message_in_first_n(cid, m1["id"], 2) is True + assert repo.message_in_first_n(cid, m2["id"], 2) is False + + def test_message_in_first_n_zero_is_empty(self, pg_conn): + repo = _repo(pg_conn) + cid = repo.create("user-1", "c")["id"] + m0 = repo.append_message(cid, {"prompt": "q", "response": "a"}) + assert repo.message_in_first_n(cid, m0["id"], 0) is False + + def test_message_in_first_n_unknown_and_null_ids(self, pg_conn): + repo = _repo(pg_conn) + cid = repo.create("user-1", "c")["id"] + repo.append_message(cid, {"prompt": "q", "response": "a"}) + assert repo.message_in_first_n(cid, str(uuid.uuid4()), 5) is False + assert repo.message_in_first_n(cid, None, 5) is False + # A non-UUID id is shape-gated, never poisoning the enclosing txn. + assert repo.message_in_first_n(cid, "not-a-uuid", 5) is False + assert pg_conn.execute(text("SELECT 1")).scalar() == 1 + + def test_first_n_message_ids_returns_snapshot_set(self, pg_conn): + repo = _repo(pg_conn) + cid = repo.create("user-1", "c")["id"] + m0 = repo.append_message(cid, {"prompt": "q0", "response": "a0"}) + m1 = repo.append_message(cid, {"prompt": "q1", "response": "a1"}) + m2 = repo.append_message(cid, {"prompt": "q2", "response": "a2"}) + assert repo.first_n_message_ids(cid, 2) == {str(m0["id"]), str(m1["id"])} + assert str(m2["id"]) not in repo.first_n_message_ids(cid, 2) + assert repo.first_n_message_ids(cid, 0) == set() + + def test_first_n_message_ids_rejects_non_uuid(self, pg_conn): + repo = _repo(pg_conn) + assert repo.first_n_message_ids("not-a-uuid", 5) == set() + assert pg_conn.execute(text("SELECT 1")).scalar() == 1 + + class TestUpdateMessageAt: def test_updates_response(self, pg_conn): repo = _repo(pg_conn) diff --git a/tests/test_code_executor_tool.py b/tests/test_code_executor_tool.py index 7e7e71e0..333afc80 100644 --- a/tests/test_code_executor_tool.py +++ b/tests/test_code_executor_tool.py @@ -381,3 +381,50 @@ def test_materialize_inputs_out_of_range_ref_is_clean_error(monkeypatch): assert "A2" in out["error"] assert "not found in this conversation/run" in out["error"] assert manager.put_files == {} + + +def test_materialize_inputs_rejects_oversize_by_declared_size(monkeypatch): + """An input whose declared version ``size`` exceeds SANDBOX_MAX_INPUT_BYTES is rejected pre-read.""" + from application.agents.tools import code_executor as ce + from application.core import settings as settings_module + + monkeypatch.setattr(settings_module.settings, "SANDBOX_MAX_INPUT_BYTES", 100, raising=False) + + class _Repo: + def __init__(self, conn): + pass + + def artifact_id_at_position(self, n, *, conversation_id=None, workflow_run_id=None): + return None + + def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None): + if conversation_id != "conv-1": + return None + return {"id": artifact_id, "current_version": 1, "title": "big.csv"} + + def get_version(self, artifact_id, version): + return { + "filename": "big.csv", + "size": 10_000, # far over the 100-byte cap + "storage_path": f"inputs/u/artifacts/{artifact_id}/v1/big.csv", + } + + class _Conn: + def __enter__(self): + return object() + + def __exit__(self, *exc): + return False + + class _Storage: + def get_file(self, path): + raise AssertionError("bytes must not be read when declared size exceeds the cap") + + monkeypatch.setattr(ce, "db_readonly", lambda: _Conn()) + monkeypatch.setattr(ce, "ArtifactsRepository", _Repo) + monkeypatch.setattr(ce.StorageCreator, "get_storage", staticmethod(lambda: _Storage())) + + manager = _InputManager() + out = _tool()._materialize_inputs(manager, "conv-1", [_ART_ID]) + assert "exceeds" in out["error"] and "sandbox input limit" in out["error"] + assert manager.put_files == {} # nothing staged