feat: more executor guards

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

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}

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:

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

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:

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:

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",
     ]
 

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:

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:

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

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

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.

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)."""

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

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:

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:

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
 

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

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)
 

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}

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.
 # ---------------------------------------------------------------------------

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
 # ---------------------------------------------------------------------------

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
 # ---------------------------------------------------------------------------

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

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 -----------------------------------------------------
 
 

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, "<kernel>", "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

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"]

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)

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