"""Code Executor tool: run sandboxed code in a semi-persistent session and capture produced files as artifacts.""" from __future__ import annotations import logging import re from typing import Any, Dict, List, Optional, Tuple from application.agents.tools.artifact_ref import resolve_artifact_id from application.agents.tools.attachment_bridge import ( AttachmentBridgeError, bridge_attachment, match_attachment, ) from application.agents.tools.base import Tool from application.core.settings import settings from application.sandbox.artifacts_capture import ( MAX_CAPTURED_FILES, capture_artifacts, snapshot_signatures, unique_input_path, ) from application.sandbox.artifacts_capture import ( infer_mime as _infer_mime, ) from application.sandbox.artifacts_capture import ( kind_for_mime as _kind_for_mime, ) from application.sandbox.base import ExecResult from application.sandbox.sandbox_creator import SandboxCreator from application.storage.db.repositories.artifacts import ArtifactsRepository from application.storage.db.session import db_readonly from application.storage.storage_creator import StorageCreator from application.utils import safe_filename logger = logging.getLogger(__name__) # Re-exported for back-compat: callers (and tests) import these mime helpers # from this module; they now live in the shared capture helper. __all__ = ["CodeExecutorTool", "_infer_mime", "_kind_for_mime", "_tail", "_OUTPUT_TAIL_BYTES"] # Maximum bytes of stdout/stderr returned to the LLM. The raw stream is never # forwarded; only this tail keeps binary/runaway output out of the context. _OUTPUT_TAIL_BYTES = 4000 # Session ids become a kernel workspace path component; the gateway only accepts # [A-Za-z0-9_-]+, so any disallowed character is stripped before binding. _SESSION_ID_RE = re.compile(r"[^A-Za-z0-9_-]+") def _tail(stream: Optional[str]) -> str: """Return the trailing slice of ``stream`` bounded by ``_OUTPUT_TAIL_BYTES``.""" if not stream: return "" if len(stream) <= _OUTPUT_TAIL_BYTES: return stream return stream[-_OUTPUT_TAIL_BYTES:] class CodeExecutorTool(Tool): """Code Executor Run code in a sandboxed session; files it writes become downloadable artifacts. """ def __init__(self, tool_config: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None) -> None: """Bind the tool to the invoker and its conversation/run-scoped sandbox session.""" self.config: Dict[str, Any] = tool_config or {} self.user_id: Optional[str] = user_id self.tool_id: Optional[str] = self.config.get("tool_id") self.conversation_id: Optional[str] = self.config.get("conversation_id") self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id") self.message_id: Optional[str] = self.config.get("message_id") # Static, deployment-level approval gate (mirrors the action metadata flag). self._require_approval: bool = bool(self.config.get("require_approval", False)) self._last_artifact_id: Optional[str] = None # ------------------------------------------------------------------ # Tool ABC # ------------------------------------------------------------------ @staticmethod def _environment_note() -> str: """Backend-specific note on what the sandbox has preinstalled. Without this the model discovers the environment by failing: importing pandas on a bare image, or pip-installing libraries that are already baked in. Keep the package lists in sync with deployment/sandbox/Dockerfile (jupyter) and scripts/build_daytona_snapshot.py (daytona snapshot). """ backend = str(getattr(settings, "SANDBOX_BACKEND", "jupyter") or "jupyter").lower() if backend == "daytona": if getattr(settings, "DAYTONA_SNAPSHOT", None): return ( "Preinstalled beyond the stdlib: python-pptx, python-docx, openpyxl, " "reportlab, lxml, pillow. pip install anything else from within the code " "before importing it." ) return ( "Only the Python stdlib is preinstalled. pip install any third-party " "package (pandas, python-docx, ...) from within the code before importing it." ) return ( "Preinstalled beyond the stdlib: pandas, matplotlib, python-pptx, python-docx, " "openpyxl, reportlab. pip install anything else from within the code before " "importing it." ) def get_actions_metadata(self) -> List[Dict[str, Any]]: """Return JSON metadata describing the ``run_code`` action for tool schemas.""" return [ { "name": "run_code", "description": ( "Execute Python in a sandboxed, stateful session bound to this conversation. " "Files written by the code are saved as downloadable artifacts (write throwaway " "files under `tmp/`, or pass `outputs` to save only specific files); only a compact " "summary (output tail + artifact references) is returned, never raw bytes. " "Each call is capped at ~60s of wall-clock; for longer work, start it in the " "background and poll with additional run_code calls (use persist=true to keep state). " + self._environment_note() ), "active": True, "require_approval": self._require_approval, "parameters": { "type": "object", "properties": { "code": { "type": "string", "description": "Python source to execute in the session. Install packages from " "within the code itself (e.g. subprocess pip install) if needed.", }, "inputs": { "type": "array", "items": {"type": "string"}, "description": "Files to materialize into the workspace; each accepts the short " "ref like `A1` returned by a previous artifact action, a full artifact id, or " "the name/id of a file the user attached to this conversation. Each is staged " "at `inputs/` before the code runs — read it from that path (the " "result's `inputs_loaded` echoes the exact staged paths).", }, "outputs": { "type": "array", "items": {"type": "string"}, "description": "Filenames or globs (e.g. `report.pdf`, `*.csv`) to save as " "downloadable artifacts. When set, only matching files are saved; when omitted, " "every produced file is saved except scratch paths under `tmp/`.", }, "ttl": { "type": "integer", "description": "Keep-alive lifetime (seconds) for the session; clamped by SANDBOX_MAX_TTL.", }, "persist": { "type": "boolean", "description": ( "Keep the session warm after the call (state survives the next run). " "The session is kept alive when this is true or a positive ttl is given " "(clamped by SANDBOX_MAX_TTL); otherwise it is closed after the run." ), }, "capture_artifacts": { "type": "boolean", "description": "Save produced workspace files as downloadable artifacts " "(default: true). Set false for setup or install-only steps that write nothing " "worth keeping.", }, }, "required": ["code"], }, } ] def get_config_requirements(self) -> Dict[str, Any]: """Return configuration requirements (none; approval is an action-level flag, and the sandbox backend is a deployment-level setting).""" return {} def get_artifact_id(self, action_name: str, **kwargs: Any) -> Optional[str]: """Return the primary produced artifact id so the UI artifact rail lights up.""" return self._last_artifact_id def preview_decision(self, action_name: str, params: dict) -> Tuple[bool, bool]: """Return ``(requires_approval, denylist_forced)`` for the approval gate; never denylist-forced here.""" if action_name != "run_code": return True, False return self._require_approval, False # ------------------------------------------------------------------ # Execution # ------------------------------------------------------------------ def execute_action(self, action_name: str, **kwargs: Any) -> Dict[str, Any]: """Dispatch a tool action; only ``run_code`` is supported.""" if action_name != "run_code": return {"status": "error", "error": f"unknown action: {action_name}"} self._last_artifact_id = None return self._run_code(**kwargs) def _run_code(self, **kwargs: Any) -> Dict[str, Any]: """Bind a session, materialize inputs, execute, and capture produced artifacts.""" if not self.user_id: return {"status": "error", "error": "code_executor requires a valid user_id."} session_id = self._resolve_session_id() if session_id is None: return {"status": "error", "error": "code_executor requires a conversation_id or workflow_run_id."} code = kwargs.get("code") if not isinstance(code, str) or not code.strip(): return {"status": "error", "error": "code is required."} should_capture = kwargs.get("capture_artifacts", True) outputs = self._normalize_outputs(kwargs.get("outputs")) ttl = self._coerce_int(kwargs.get("ttl")) timeout = self._exec_timeout() inputs = kwargs.get("inputs") or [] manager = SandboxCreator.get_manager() try: manager.open(session_id, ttl=ttl) except Exception as exc: logger.exception("code_executor: failed to open sandbox session") return {"status": "error", "error": f"sandbox unavailable: {type(exc).__name__}: {exc}"} try: materialized = self._materialize_inputs(manager, session_id, inputs) if materialized.get("error"): return {"status": "error", "error": materialized["error"]} pre_signatures: Dict[str, Tuple[int, Optional[str]]] = {} if should_capture: pre_signatures = self._snapshot_signatures(manager, session_id) try: result = manager.exec(session_id, code, timeout=timeout) except Exception as exc: logger.exception("code_executor: exec raised") return {"status": "error", "error": f"execution failed: {type(exc).__name__}: {exc}"} # Capture even on error/timeout while the runtime remains reachable # so partial outputs aren't lost; capture never masks the run status. artifacts: List[Dict[str, Any]] = [] if should_capture and not result.runtime_invalidated: try: artifacts = self._capture_artifacts(manager, session_id, pre_signatures, outputs) except Exception: logger.exception("code_executor: artifact capture failed") return self._shape_payload(result, artifacts, materialized.get("loaded", [])) finally: if not self._keep_alive(kwargs.get("persist"), ttl): try: manager.close(session_id) except Exception: logger.exception("code_executor: session close failed") # ------------------------------------------------------------------ # Inputs / outputs # ------------------------------------------------------------------ def _materialize_inputs(self, manager: Any, session_id: str, inputs: List[Any]) -> Dict[str, Any]: """Fetch parent-scoped input artifacts and copy their current-version bytes into the workspace.""" loaded: List[str] = [] if not inputs: return {"loaded": loaded} storage = StorageCreator.get_storage() # Two inputs whose current versions share a filename would clobber each other at # the same ``inputs/{name}`` path; track used paths and disambiguate deterministically. used_paths: set = set() for raw_id in inputs: raw = str(raw_id).strip() if not raw: continue artifact_id: Optional[str] = raw try: with db_readonly() as conn: repo = ArtifactsRepository(conn) # A short ref (A1/A2/...) resolves to an id within this parent # only; the resolved id still passes through the parent-scoped # gate so a ref can never reach another tenant. artifact_id = resolve_artifact_id( repo, raw, conversation_id=self.conversation_id, workflow_run_id=self.workflow_run_id, ) artifact = ( repo.get_artifact_in_parent( artifact_id, conversation_id=self.conversation_id, workflow_run_id=self.workflow_run_id, ) if artifact_id is not None else None ) if artifact is None: # Conversation scope only: a raw ref that is not an artifact # may name a chat attachment; bridge it on demand. Workflows # bridge attachments up front, so never double-bridge there. bridged_id = self._bridge_chat_attachment(raw) if isinstance(bridged_id, dict): return bridged_id # error payload if bridged_id is None: return {"error": f"input artifact {raw} not found in this conversation/run."} artifact_id = bridged_id artifact = repo.get_artifact_in_parent(artifact_id, conversation_id=self.conversation_id) if artifact is None: return {"error": f"input artifact {raw} not found in this conversation/run."} version = repo.get_version(artifact_id, artifact["current_version"]) except Exception: logger.exception("code_executor: failed to load input artifact") return {"error": f"failed to load input artifact {artifact_id}."} 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"]) 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."} rel_path = unique_input_path(f"inputs/{filename}", used_paths) try: manager.put_file(session_id, rel_path, data) except Exception: logger.exception("code_executor: put_file failed for input artifact") return {"error": f"failed to stage input artifact {artifact_id} into the workspace."} loaded.append(rel_path) return {"loaded": loaded} def _bridge_chat_attachment(self, raw: str) -> Any: """Bridge a referenced chat attachment to a conversation artifact id; None on miss, error dict on failure.""" if not self.conversation_id or not self.user_id: return None attachment = match_attachment(self.config.get("attachments"), raw, self.user_id) if attachment is None: return None try: return bridge_attachment(attachment, user_id=self.user_id, conversation_id=self.conversation_id) except AttachmentBridgeError as exc: return {"error": f"failed to attach {raw}: {exc}"} # Cap the per-run capture work so a workspace full of pre-existing files # can't turn one exec into an unbounded read+persist sweep. _MAX_CAPTURED_FILES = MAX_CAPTURED_FILES def _snapshot_signatures(self, manager: Any, session_id: str) -> Dict[str, Tuple[int, Optional[str]]]: """Map each non-input workspace file to a (size, sha256) signature for change detection.""" return snapshot_signatures(manager, session_id) @staticmethod def _normalize_outputs(raw: Any) -> Optional[List[str]]: """Coerce the ``outputs`` arg to a list of non-empty glob strings, or None. Tolerates a bare string (some models pass one instead of an array); an empty or non-list value means "no allow-list" (auto-capture). """ if isinstance(raw, str): raw = [raw] if not isinstance(raw, list): return None patterns = [str(p).strip() for p in raw if isinstance(p, str) and str(p).strip()] return patterns or None def _capture_artifacts( self, manager: Any, session_id: str, pre_signatures: Dict[str, Tuple[int, Optional[str]]], outputs: Optional[List[str]] = None, ) -> List[Dict[str, Any]]: """Persist produced workspace files (only ``outputs`` globs when given).""" captured = capture_artifacts( manager, session_id, pre_signatures, user_id=self.user_id, conversation_id=self.conversation_id, workflow_run_id=self.workflow_run_id, message_id=self.message_id, produced_by={ "tool": "code_executor", "action": "run_code", "session_id": session_id, }, outputs=outputs, ) if captured: self._last_artifact_id = captured[0]["artifact_id"] return captured def _shape_payload( self, result: ExecResult, artifacts: List[Dict[str, Any]], inputs_loaded: List[str] ) -> Dict[str, Any]: """Build the compact LLM-facing payload; raw bytes never appear here.""" status = "ok" if result.ok else "error" payload: Dict[str, Any] = { "status": status, "stdout_tail": _tail(result.stdout), "artifacts": artifacts, } stderr_tail = _tail(result.stderr) if stderr_tail: payload["stderr_tail"] = stderr_tail if not result.ok: if self._is_timeout(result): cap = int(self._exec_timeout()) payload["error"] = ( f"Execution timed out. Each run_code call is capped at {cap}s and the limit " "cannot be raised. For long-running work, start it in the background (e.g. launch a " "subprocess or `nohup ... &` and write progress to a file) and return immediately, " "then poll with additional run_code calls to check on it. Pass persist=true (or a " "ttl) so the background process and its files survive between calls." ) else: payload["error"] = ( f"{result.error_name}: {result.error_value}" if result.error_name else (result.error_value or "execution error") ) if inputs_loaded: payload["inputs_loaded"] = inputs_loaded return payload # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ def _resolve_session_id(self) -> Optional[str]: """Derive a sandbox session id from the bound conversation/run; sanitize to the gateway charset.""" raw = self.conversation_id or self.workflow_run_id if not raw: return None sanitized = _SESSION_ID_RE.sub("-", str(raw)) return sanitized or None @staticmethod def _coerce_int(value: Any) -> Optional[int]: """Coerce a value to a positive int, or None when absent/invalid.""" if value is None: return None try: parsed = int(value) except (TypeError, ValueError): return None return parsed if parsed > 0 else None @staticmethod def _exec_timeout() -> float: """Return the fixed per-run wall-clock cap (SANDBOX_EXEC_TIMEOUT; not caller-adjustable).""" return float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60)) @staticmethod def _is_timeout(result: ExecResult) -> bool: """True when a failed exec looks like a wall-clock timeout (any backend's naming/message).""" blob = f"{result.error_name or ''} {result.error_value or ''}".lower() return "timeout" in blob or "timed out" in blob @staticmethod def _keep_alive(persist: Any, ttl: Optional[int]) -> bool: """True when the agent asked to keep the session warm after the call.""" return bool(persist) or (ttl is not None and ttl > 0)