From 66786c276048397f44106bc4587909169ac3a1c8 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 8 Jul 2026 18:50:10 +0100 Subject: [PATCH] fix: remove code exec as default and sec improvements --- .github/THREAT_MODEL.md | 6 +- application/agents/default_tools.py | 15 +- application/agents/tool_executor.py | 233 ++++++------------ .../agents/tools/artifact_generator.py | 8 +- application/agents/tools/code_executor.py | 18 +- application/agents/tools/read_document.py | 14 +- application/agents/workflow_agent.py | 77 +++--- application/api/user/artifacts/routes.py | 212 +++++----------- application/api/user/tasks.py | 34 +++ application/core/settings.py | 15 +- application/parser/document_reader.py | 45 +++- application/prompts/agentic/creative.txt | 4 +- application/prompts/agentic/default.txt | 4 +- application/prompts/agentic/strict.txt | 4 +- application/prompts/chat_combine_creative.txt | 4 +- application/prompts/chat_combine_default.txt | 4 +- application/prompts/chat_combine_strict.txt | 4 +- application/sandbox/artifacts_capture.py | 6 + .../storage/db/repositories/artifacts.py | 26 +- .../storage/db/repositories/workflow_runs.py | 24 +- application/worker.py | 1 + deployment/docker-compose-azure.yaml | 88 +------ deployment/docker-compose-hub.yaml | 91 +------ deployment/docker-compose.yaml | 100 +------- .../k8s/deployments/docsgpt-deploy.yaml | 16 -- .../k8s/deployments/sandbox-deploy.yaml | 33 ++- ...ocker-compose.optional.sandbox-egress.yaml | 26 +- .../docker-compose.optional.sandbox.yaml | 91 +++++++ deployment/sandbox/Dockerfile | 16 +- deployment/sandbox/README.md | 99 +++++--- deployment/sandbox/gateway-launch.sh | 33 +++ deployment/sandbox/kernel-launch.sh | 0 .../Tools/artifacts-and-code-execution.mdx | 2 +- .../src/components/artifactViewUtils.test.ts | 29 ++- frontend/src/components/artifactViewUtils.ts | 26 +- tests/agents/test_default_tools.py | 35 +-- tests/api/user/test_artifacts_routes.py | 3 + tests/api/user/test_tasks.py | 5 +- tests/parser/test_document_reader.py | 38 +++ .../storage/db/repositories/test_artifacts.py | 48 ++++ .../db/repositories/test_workflow_runs.py | 35 ++- 41 files changed, 800 insertions(+), 772 deletions(-) create mode 100644 deployment/optional/docker-compose.optional.sandbox.yaml create mode 100755 deployment/sandbox/gateway-launch.sh mode change 100644 => 100755 deployment/sandbox/kernel-launch.sh diff --git a/.github/THREAT_MODEL.md b/.github/THREAT_MODEL.md index babeb54a..498e3a8b 100644 --- a/.github/THREAT_MODEL.md +++ b/.github/THREAT_MODEL.md @@ -107,9 +107,9 @@ Untrusted input includes API payloads, file uploads, remote URLs, OAuth/webhook ### I. Sandboxed code execution and tenant isolation - Threat: LLM-authored code (the code-execution tool, document/artifact generation, and workflow code nodes) runs attacker-influenceable Python; a poisoned document or prompt can shape what executes. -- Threat: on the self-hosted Jupyter Kernel Gateway runner, all sessions run as kernels inside one shared container and uid — a kernel can read sibling sessions' workspaces and reach the network. Treat a single runner as one trust domain, not a per-tenant boundary. -- Threat: `code_executor` is on by default and runs sandboxed code WITHOUT an approval prompt, so any chat agent (including one steered by prompt injection) can execute code unprompted; the sandbox is the only boundary. -- Mitigations: code-exec approval exists per tool but is off for the default tool (enable it or rely on the sandbox boundary); pass workflow state to code nodes as data (a `state.json` file), never templated into the executed program; scrub secrets from the kernel environment (no provider keys/tokens/DB URL reach kernel code); path-traversal-safe file I/O with output/time/size caps and per-session `0700` workspaces; block egress at the network layer (NetworkPolicy/host firewall). For per-tenant isolation use the Daytona per-session-VM backend (`SANDBOX_BACKEND=daytona`); run the self-hosted runner under gVisor for host protection. Artifacts are access-controlled by their parent (conversation or workflow run). +- Threat: on the self-hosted Jupyter Kernel Gateway runner, all sessions run as kernels inside one shared container and uid — a kernel can read sibling sessions' workspaces and reach the network. Treat a single runner as one trust domain, not a per-tenant boundary. The gateway's control API is reachable from kernel code over loopback, so it is authenticated (a required, env-scrubbed token) to stop a kernel from driving sibling kernels or bypassing the session cap. +- Threat: an agent with `code_executor` / `artifact_generator` enabled runs sandboxed code that a poisoned document or prompt can shape; a prompt-injected agent can execute code within the sandbox boundary. Both tools are opt-in (off by default, not in `DEFAULT_CHAT_TOOLS`, and gated behind a per-agent enable plus a running runner), which limits exposure to agents an operator deliberately configured for code execution. +- Mitigations: code-exec approval is available per tool; the runner and both tools are opt-in (a fresh deploy runs no sandbox); the gateway requires an auth token (fails closed) and scrubs it plus all secrets from the kernel environment; pass workflow state to code nodes as data (a `state.json` file), never templated into the executed program; path-traversal-safe file I/O with output/time/size caps and per-session `0700` workspaces; block egress at the network layer (NetworkPolicy/host firewall). For per-tenant isolation use the Daytona per-session-VM backend (`SANDBOX_BACKEND=daytona`); run the self-hosted runner under gVisor for host protection. Artifacts are access-controlled by their parent (conversation or workflow run). ## 8) Example attacker stories diff --git a/application/agents/default_tools.py b/application/agents/default_tools.py index 28b0c291..b29f6220 100644 --- a/application/agents/default_tools.py +++ b/application/agents/default_tools.py @@ -31,11 +31,16 @@ _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``). -# ``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") +# ``code_executor`` and ``artifact_generator`` are builtin-only (not default-on): +# both render/execute through a running sandbox runner, so they are opt-in per +# agent, but staying registered keeps their synthetic ids resolvable (an agent +# that enabled one never silently loses it) and keeps them in the agent picker. +BUILTIN_AGENT_TOOLS: tuple = ( + "scheduler", + "read_document", + "code_executor", + "artifact_generator", +) # Builtins shown only in the workflow-node tool picker, never the classic # agent picker. The synthesized row carries ``workflow_only`` so the frontend diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py index 545bdd9f..be392b30 100644 --- a/application/agents/tool_executor.py +++ b/application/agents/tool_executor.py @@ -85,11 +85,7 @@ def _record_proposed( tool_id=tool_id if tool_id and looks_like_uuid(tool_id) else None, message_id=message_id, user_id=user_id, - agent_id=( - str(agent_id) - if agent_id and looks_like_uuid(str(agent_id)) - else None - ), + agent_id=(str(agent_id) if agent_id and looks_like_uuid(str(agent_id)) else None), ) if not inserted: logger.warning( @@ -147,11 +143,7 @@ def _mark_executed( message_id=message_id, artifact_id=artifact_id, user_id=user_id, - agent_id=( - str(agent_id) - if agent_id and looks_like_uuid(str(agent_id)) - else None - ), + agent_id=(str(agent_id) if agent_id and looks_like_uuid(str(agent_id)) else None), ) except Exception: logger.exception("tool_call_attempts executed write failed for %s", call_id) @@ -160,9 +152,7 @@ def _mark_executed( def _mark_failed(call_id: str, error: str, *, user_id: Optional[str] = None) -> None: try: with db_session() as conn: - ToolCallAttemptsRepository(conn).mark_failed( - call_id, error, user_id=user_id - ) + ToolCallAttemptsRepository(conn).mark_failed(call_id, error, user_id=user_id) except Exception: logger.exception("tool_call_attempts failed-write failed for %s", call_id) @@ -191,9 +181,7 @@ class ToolExecutor: # so check_pause returns headless_denied sentinels instead. self.headless = bool(headless) # Tool-instance ids pre-authorized for headless approval-gated execution. - self.tool_allowlist: set = ( - {str(x) for x in tool_allowlist} if tool_allowlist else set() - ) + self.tool_allowlist: set = {str(x) for x in tool_allowlist} if tool_allowlist else set() self.tool_calls: List[Dict] = [] self._loaded_tools: Dict[str, object] = {} self.conversation_id: Optional[str] = None @@ -234,11 +222,7 @@ class ToolExecutor: tools plus the synthesized defaults. Used to gate tool-specific prompt sections via the ``tools.enabled`` template namespace. """ - return { - str(tool["name"]) - for tool in self.get_tools().values() - if isinstance(tool, dict) and tool.get("name") - } + return {str(tool["name"]) for tool in self.get_tools().values() if isinstance(tool, dict) and tool.get("name")} def _get_tools_by_api_key(self, api_key: str) -> Dict[str, Dict]: """Resolve an agent's toolset — exactly ``agents.tools``, no defaults.""" @@ -249,11 +233,7 @@ class ToolExecutor: agent_data = AgentsRepository(conn).find_by_key(api_key) tool_ids = agent_data.get("tools", []) if agent_data else [] tools_repo = UserToolsRepository(conn) - owner = ( - (agent_data.get("user_id") or agent_data.get("user")) - if agent_data - else None - ) + owner = (agent_data.get("user_id") or agent_data.get("user")) if agent_data else None tools: List[Dict] = [] for tid in tool_ids: row = resolve_tool_by_id(tid, owner, user_tools_repo=tools_repo) @@ -270,29 +250,23 @@ class ToolExecutor: """Resolve an agentless chat's toolset: explicit user tools plus defaults.""" with db_readonly() as conn: user_tools = UserToolsRepository(conn).list_active_for_user(user) - user_doc = ( - UsersRepository(conn).get(user) if self.agent_id is None else None - ) + user_doc = UsersRepository(conn).get(user) if self.agent_id is None else None # Headless agentless runs (e.g. scheduled fire) drop chat-only # tools (``scheduler``) from explicit user_tools too. filtered_user_tools = [ - t for t in user_tools - if not (self.headless and is_headless_excluded_tool(t.get("name"))) + t for t in user_tools if not (self.headless and is_headless_excluded_tool(t.get("name"))) ] # Index keys (ints) and synthetic uuid5 keys can't collide. - tools: Dict[str, Dict] = { - str(i): tool for i, tool in enumerate(filtered_user_tools) - } + tools: Dict[str, Dict] = {str(i): tool for i, tool in enumerate(filtered_user_tools)} if self.agent_id is None: for default_row in synthesized_default_tools( - user_doc, headless=self.headless, + user_doc, + headless=self.headless, ): tools[str(default_row["id"])] = default_row return tools - def merge_client_tools( - self, tools_dict: Dict, client_tools: List[Dict] - ) -> Dict: + def merge_client_tools(self, tools_dict: Dict, client_tools: List[Dict]) -> Dict: """Merge client-provided tool definitions into tools_dict. Client tools use the standard function-calling format:: @@ -358,18 +332,12 @@ class ToolExecutor: if not is_api and "actions" not in tool: continue - actions = ( - tool["config"]["actions"].values() - if is_api - else tool["actions"] - ) + actions = tool["config"]["actions"].values() if is_api else tool["actions"] for action in actions: if not action.get("active", True): continue - entries.append( - (tool_id, tool.get("name", ""), action["name"], action, is_client) - ) + entries.append((tool_id, tool.get("name", ""), action["name"], action, is_client)) name_counts[action["name"]] += 1 # Pass 2: assign LLM-visible names and build mappings @@ -379,34 +347,19 @@ class ToolExecutor: result = [] for tool_id, tool_name, action_name, action, is_client in entries: - if ( - name_counts[action_name] == 1 - and len(action_name) <= _MAX_LLM_NAME_LEN - ): + if name_counts[action_name] == 1 and len(action_name) <= _MAX_LLM_NAME_LEN: llm_name = action_name else: # An over-long unique name skips the prefix — it needs # truncation, not disambiguation. - prefix = ( - _sanitize_tool_prefix(tool_name) - if name_counts[action_name] > 1 - else "" - ) - base = ( - f"{prefix}_{action_name}" - if prefix and not action_name.startswith(f"{prefix}_") - else action_name - ) + prefix = _sanitize_tool_prefix(tool_name) if name_counts[action_name] > 1 else "" + base = f"{prefix}_{action_name}" if prefix and not action_name.startswith(f"{prefix}_") else action_name base = base[:_MAX_LLM_NAME_LEN] # A duplicated bare name stays ambiguous, and a candidate # must not steal a unique action's name or one already taken. candidate = base counter = 1 - while ( - candidate == action_name - or candidate in all_llm_names - or name_counts.get(candidate, 0) == 1 - ): + while candidate == action_name or candidate in all_llm_names or name_counts.get(candidate, 0) == 1: suffix = f"_{counter}" candidate = base[: _MAX_LLM_NAME_LEN - len(suffix)] + suffix counter += 1 @@ -421,14 +374,16 @@ class ToolExecutor: else: params = self._build_tool_parameters(action) - result.append({ - "type": "function", - "function": { - "name": llm_name, - "description": action.get("description", ""), - "parameters": params, - }, - }) + result.append( + { + "type": "function", + "function": { + "name": llm_name, + "description": action.get("description", ""), + "parameters": params, + }, + } + ) return result def _build_tool_parameters(self, action: Dict) -> Dict: @@ -438,17 +393,13 @@ class ToolExecutor: for k, v in action[param_type]["properties"].items(): if v.get("filled_by_llm", True): params["properties"][k] = { - key: value - for key, value in v.items() - if key not in ("filled_by_llm", "value", "required") + key: value for key, value in v.items() if key not in ("filled_by_llm", "value", "required") } if v.get("required", False): params["required"].append(k) return params - def check_pause( - self, tools_dict: Dict, call, llm_class_name: str - ) -> Optional[Dict]: + def check_pause(self, tools_dict: Dict, call, llm_class_name: str) -> Optional[Dict]: """Return a pending-action dict (approval / client / headless_denied) or None. In headless mode the dict's pause_type is ``headless_denied`` so the @@ -478,9 +429,7 @@ class ToolExecutor: "llm_name": llm_name, "arguments": arguments, "pause_type": "headless_denied", - "deny_reason": ( - "Client-side tools cannot run in headless / scheduled runs." - ), + "deny_reason": ("Client-side tools cannot run in headless / scheduled runs."), "error_type": "tool_not_allowed", "thought_signature": getattr(call, "thought_signature", None), } @@ -498,9 +447,7 @@ class ToolExecutor: # Approval required if tool_data["name"] == "api_tool": - action_data = tool_data.get("config", {}).get("actions", {}).get( - action_name, {} - ) + action_data = tool_data.get("config", {}).get("actions", {}).get(action_name, {}) else: action_data = next( (a for a in tool_data.get("actions", []) if a["name"] == action_name), @@ -517,17 +464,22 @@ class ToolExecutor: # not reflect later approval-mode changes nor command-level # heuristics, so consult the tool directly. if tool_data.get("name") == "remote_device": - require_approval, denylist_forced = ( - self._remote_device_requires_approval( - tool_data, action_name, arguments, - ) + require_approval, denylist_forced = self._remote_device_requires_approval( + tool_data, + action_name, + arguments, ) elif tool_data.get("name") == "code_executor": # The deployment-level ``config.require_approval`` is authoritative # over the cached action snapshot, so consult the tool directly. - require_approval = self._code_executor_requires_approval( - tool_data, action_name, arguments, - ) or require_approval + require_approval = ( + self._code_executor_requires_approval( + tool_data, + action_name, + arguments, + ) + or require_approval + ) if require_approval: if self.headless: @@ -548,10 +500,7 @@ class ToolExecutor: "llm_name": llm_name, "arguments": arguments, "pause_type": "headless_denied", - "deny_reason": ( - "This tool requires approval and is not in the run's " - "tool_allowlist." - ), + "deny_reason": ("This tool requires approval and is not in the run's tool_allowlist."), "error_type": "tool_not_allowed", "thought_signature": getattr(call, "thought_signature", None), } @@ -577,7 +526,10 @@ class ToolExecutor: return None def _remote_device_requires_approval( - self, tool_data: Dict, action_name: str, arguments: Dict, + self, + tool_data: Dict, + action_name: str, + arguments: Dict, ) -> tuple[bool, bool]: """Live approval decision for a ``remote_device`` invocation. @@ -597,13 +549,15 @@ class ToolExecutor: return tool.preview_decision(action_name, arguments) except Exception: logger.exception( - "remote_device preview_decision failed; defaulting to a " - "forced prompt", + "remote_device preview_decision failed; defaulting to a forced prompt", ) return True, True def _code_executor_requires_approval( - self, tool_data: Dict, action_name: str, arguments: Dict, + self, + tool_data: Dict, + action_name: str, + arguments: Dict, ) -> bool: """Live approval decision for a ``code_executor`` invocation. @@ -662,9 +616,7 @@ class ToolExecutor: user_id=self.user, agent_id=self.agent_id, ): - _mark_failed( - call_id, tool_call_data["result"], user_id=self.user - ) + _mark_failed(call_id, tool_call_data["result"], user_id=self.user) yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}} self.tool_calls.append(tool_call_data) return "Failed to parse tool call.", call_id @@ -738,10 +690,7 @@ class ToolExecutor: # event and flip the journal row to ``failed`` instead of # killing the stream. if not isinstance(call_args, dict): - error_message = ( - f"Tool call arguments must be a JSON object, got " - f"{type(call_args).__name__}." - ) + error_message = f"Tool call arguments must be a JSON object, got {type(call_args).__name__}." tool_call_data["result"] = error_message tool_call_data["arguments"] = {} if proposed_ok: @@ -756,11 +705,7 @@ class ToolExecutor: action_data = ( tool_data["config"]["actions"][action_name] if tool_data["name"] == "api_tool" - else next( - action - for action in tool_data["actions"] - if action["name"] == action_name - ) + else next(action for action in tool_data["actions"] if action["name"] == action_name) ) query_params, headers, body, parameters = {}, {}, {}, {} @@ -774,29 +719,25 @@ class ToolExecutor: for param_type, target_dict in param_types.items(): if param_type in action_data and action_data[param_type].get("properties"): for param, details in action_data[param_type]["properties"].items(): - if ( - param not in call_args - and "value" in details - and details["value"] - ): + if param not in call_args and "value" in details and details["value"]: target_dict[param] = details["value"] for param, value in call_args.items(): for param_type, target_dict in param_types.items(): - if param_type in action_data and param in action_data[param_type].get( - "properties", {} - ): + if param_type in action_data and param in action_data[param_type].get("properties", {}): target_dict[param] = value # Load tool (with caching) tool = self._get_or_load_tool( - tool_data, tool_id, action_name, - headers=headers, query_params=query_params, + tool_data, + tool_id, + action_name, + headers=headers, + query_params=query_params, ) if tool is None: error_message = ( - f"Failed to load tool '{tool_data.get('name')}' (tool_id key={tool_id}): " - "missing 'id' on tool row." + f"Failed to load tool '{tool_data.get('name')}' (tool_id key={tool_id}): missing 'id' on tool row." ) logger.error( "tool_load_failed", @@ -841,11 +782,7 @@ class ToolExecutor: _mark_failed(call_id, str(exc), user_id=self.user) raise - get_artifact_id = ( - getattr(tool, "get_artifact_id", None) - if tool_data["name"] != "api_tool" - else None - ) + get_artifact_id = getattr(tool, "get_artifact_id", None) if tool_data["name"] != "api_tool" else None artifact_id = None if callable(get_artifact_id): @@ -864,9 +801,7 @@ class ToolExecutor: result_full = str(result) tool_call_data["resolved_arguments"] = resolved_arguments tool_call_data["result_full"] = result_full - tool_call_data["result"] = ( - f"{result_full[:50]}..." if len(result_full) > 50 else result_full - ) + tool_call_data["result"] = f"{result_full[:50]}..." if len(result_full) > 50 else result_full # Tool side effect has run; flip the journal row so the # message-finalize path can later confirm it. If the proposed @@ -887,9 +822,7 @@ class ToolExecutor: ) stream_tool_call_data = { - key: value - for key, value in tool_call_data.items() - if key not in {"result_full", "resolved_arguments"} + key: value for key, value in tool_call_data.items() if key not in {"result_full", "resolved_arguments"} } yield {"type": "tool_call", "data": {**stream_tool_call_data, "status": "completed"}} self.tool_calls.append(tool_call_data) @@ -897,8 +830,12 @@ class ToolExecutor: return result, call_id def _get_or_load_tool( - self, tool_data: Dict, tool_id: str, action_name: str, - headers: Optional[Dict] = None, query_params: Optional[Dict] = None, + self, + tool_data: Dict, + tool_id: str, + action_name: str, + headers: Optional[Dict] = None, + query_params: Optional[Dict] = None, ): """Load a tool, using cache when possible.""" cache_key = f"{tool_data['name']}:{tool_id}:{self.user or ''}" @@ -924,12 +861,8 @@ class ToolExecutor: "query_params": query_params or {}, } if "body_content_type" in action_config: - tool_config["body_content_type"] = action_config.get( - "body_content_type", "application/json" - ) - tool_config["body_encoding_rules"] = action_config.get( - "body_encoding_rules", {} - ) + tool_config["body_content_type"] = action_config.get("body_content_type", "application/json") + tool_config["body_encoding_rules"] = action_config.get("body_encoding_rules", {}) else: tool_config = tool_data["config"].copy() if tool_data["config"] else {} # Credentials are PBKDF2-bound to the tool OWNER's sub, not the @@ -955,9 +888,7 @@ class ToolExecutor: "agent_id": self.agent_id, }, ) - decrypted = decrypt_credentials( - tool_config["encrypted_credentials"], tool_owner - ) + decrypted = decrypt_credentials(tool_config["encrypted_credentials"], tool_owner) tool_config.update(decrypted) tool_config["auth_credentials"] = decrypted tool_config.pop("encrypted_credentials", None) @@ -975,6 +906,8 @@ class ToolExecutor: tool_config["tool_id"] = str(row_id) if self.conversation_id: tool_config["conversation_id"] = self.conversation_id + if self.message_id: + tool_config["message_id"] = self.message_id # Carry the request's own attachments so sandbox tools can # lazily bridge a referenced chat attachment (conversation # scope only; workflow nodes bridge attachments up front). @@ -988,9 +921,7 @@ class ToolExecutor: # Agent-bound: stamp schedules.agent_id. Agentless: the tool # falls back to ``origin_conversation_id`` as the schedule's # conversation home. - tool_config["agent_id"] = ( - str(self.agent_id) if self.agent_id else None - ) + tool_config["agent_id"] = str(self.agent_id) if self.agent_id else None if tool_data["name"] == "mcp_tool": tool_config["query_mode"] = True @@ -1015,9 +946,7 @@ class ToolExecutor: "arguments": tool_call.get("arguments"), "artifact_id": tool_call.get("artifact_id"), "result": ( - f"{str(tool_call['result'])[:50]}..." - if len(str(tool_call["result"])) > 50 - else tool_call["result"] + f"{str(tool_call['result'])[:50]}..." if len(str(tool_call["result"])) > 50 else tool_call["result"] ), "status": "completed", } diff --git a/application/agents/tools/artifact_generator.py b/application/agents/tools/artifact_generator.py index 009d5dff..e1457e53 100644 --- a/application/agents/tools/artifact_generator.py +++ b/application/agents/tools/artifact_generator.py @@ -275,7 +275,7 @@ _RENDERERS: Dict[str, str] = { " if c is None:\n" " return ''\n" " if isinstance(c, str) and c[:1] in ('=', '+', '-', '@', chr(9), chr(13), chr(10)):\n" - " return \"'\" + c\n" + ' return "\'" + c\n' " return c\n" "wb = Workbook()\n" "wb.remove(wb.active)\n" @@ -359,8 +359,8 @@ _RENDERERS: Dict[str, str] = { " 'code{{font-family:ui-monospace,SFMono-Regular,Menlo,monospace}}'\n" ")\n" "doc = (\n" - " ''\n" - " ''\n" + ' \'\'\n' + ' \'\'\n' " '' + esc(title or 'Report') + ''\n" " ''\n" " + ''.join(parts) + ''\n" @@ -397,6 +397,7 @@ class ArtifactGeneratorTool(Tool): 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") self._last_artifact_id: Optional[str] = None # ------------------------------------------------------------------ @@ -530,6 +531,7 @@ class ArtifactGeneratorTool(Tool): title=title, conversation_id=self.conversation_id, workflow_run_id=self.workflow_run_id, + message_id=self.message_id, spec=spec, produced_by=self._produced_by("create_artifact", kind), ) diff --git a/application/agents/tools/code_executor.py b/application/agents/tools/code_executor.py index e71c18d7..260128b9 100644 --- a/application/agents/tools/code_executor.py +++ b/application/agents/tools/code_executor.py @@ -17,11 +17,15 @@ from application.core.settings import settings from application.sandbox.artifacts_capture import ( MAX_CAPTURED_FILES, capture_artifacts, - infer_mime as _infer_mime, - kind_for_mime as _kind_for_mime, 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 @@ -65,6 +69,7 @@ class CodeExecutorTool(Tool): 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 @@ -264,9 +269,7 @@ class CodeExecutorTool(Tool): 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 - ) + 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"]) @@ -316,9 +319,7 @@ class CodeExecutorTool(Tool): if attachment is None: return None try: - return bridge_attachment( - attachment, user_id=self.user_id, conversation_id=self.conversation_id - ) + 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}"} @@ -359,6 +360,7 @@ class CodeExecutorTool(Tool): 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", diff --git a/application/agents/tools/read_document.py b/application/agents/tools/read_document.py index 7acdf60b..15a8b13e 100644 --- a/application/agents/tools/read_document.py +++ b/application/agents/tools/read_document.py @@ -56,6 +56,7 @@ class ReadDocumentTool(Tool): 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") self._last_artifact_id: Optional[str] = None # ------------------------------------------------------------------ @@ -222,9 +223,7 @@ class ReadDocumentTool(Tool): timeout = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120)) queue = getattr(settings, "DOCUMENT_PARSE_QUEUE", "parsing") try: - async_result = parse_document.apply_async( - args=[artifact_id, parent, self.user_id, options], queue=queue - ) + async_result = parse_document.apply_async(args=[artifact_id, parent, self.user_id, options], queue=queue) # The web process (not a worker) awaits here; ``disable_sync_subtasks=False`` keeps # the call correct if invoked from a non-prefork (eventlet/gevent) worker where the # inline branch above still ran but the blanket guard would otherwise raise. @@ -283,16 +282,17 @@ class ReadDocumentTool(Tool): if attachment is None: return None try: - return bridge_attachment( - attachment, user_id=self.user_id, conversation_id=self.conversation_id - ) + return bridge_attachment(attachment, user_id=self.user_id, conversation_id=self.conversation_id) except AttachmentBridgeError as exc: return {"status": "error", "error": f"failed to attach {raw_id}: {exc}"} def _parent(self) -> Dict[str, Any]: """Build the run-scoped parent dict passed to the worker for its independent re-resolve.""" if self.conversation_id is not None: - return {"conversation_id": self.conversation_id} + parent: Dict[str, Any] = {"conversation_id": self.conversation_id} + if self.message_id: + parent["message_id"] = self.message_id + return parent return {"workflow_run_id": self.workflow_run_id} # ------------------------------------------------------------------ diff --git a/application/agents/workflow_agent.py b/application/agents/workflow_agent.py index 5a644097..b75049a0 100644 --- a/application/agents/workflow_agent.py +++ b/application/agents/workflow_agent.py @@ -13,7 +13,7 @@ from application.agents.workflows.schemas import ( ) from application.agents.workflows.workflow_engine import WorkflowEngine from application.core.settings import settings -from application.logging import log_activity, LogContext +from application.logging import LogContext, log_activity from application.sandbox.artifacts_capture import QuotaExceeded from application.storage.db.base_repository import looks_like_uuid from application.storage.db.repositories.workflow_edges import WorkflowEdgesRepository @@ -51,14 +51,10 @@ class WorkflowAgent(BaseAgent): self._bridge_error: Optional[str] = None @log_activity() - def gen( - self, query: str, log_context: LogContext = None - ) -> Generator[Dict[str, str], None, None]: + def gen(self, query: str, log_context: LogContext = None) -> Generator[Dict[str, str], None, None]: yield from self._gen_inner(query, log_context) - def _gen_inner( - self, query: str, log_context: LogContext - ) -> Generator[Dict[str, str], None, None]: + def _gen_inner(self, query: str, log_context: LogContext) -> Generator[Dict[str, str], None, None]: graph = self._load_workflow_graph() if not graph: yield {"type": "error", "error": "Failed to load workflow configuration."} @@ -73,15 +69,11 @@ class WorkflowAgent(BaseAgent): # the outputs of the run they triggered (authz is run.user_id == caller). workflow_owner_id = self._resolve_owner_id() run_user_id = self._resolve_run_user_id(workflow_owner_id) - pg_workflow_id = self._precreate_workflow_run( - workflow_owner_id, run_user_id, query - ) + pg_workflow_id = self._precreate_workflow_run(workflow_owner_id, run_user_id, query) self._run_persisted = pg_workflow_id is not None try: - input_documents, dropped = self._bridge_attachments( - run_user_id, persisted=self._run_persisted - ) + input_documents, dropped = self._bridge_attachments(run_user_id, persisted=self._run_persisted) except QuotaExceeded as exc: # The run's input documents exceed the uploader's artifact quota. Surface # a clean error and finalize the pre-created RUNNING row as FAILED rather @@ -94,9 +86,7 @@ class WorkflowAgent(BaseAgent): "Delete some artifacts and try again." ), } - self._finalize_workflow_run( - workflow_owner_id, run_user_id, pg_workflow_id, query - ) + self._finalize_workflow_run(workflow_owner_id, run_user_id, pg_workflow_id, query) return # Non-fatal: some attachments were dropped (oversize / unreadable). Tell the @@ -105,10 +95,18 @@ class WorkflowAgent(BaseAgent): yield {"type": "error", "error": " ".join(dropped)} self._engine.run_persisted = self._run_persisted - yield from self._engine.execute({"input_documents": input_documents}, query) - self._finalize_workflow_run( - workflow_owner_id, run_user_id, pg_workflow_id, query - ) + interrupted = True + try: + yield from self._engine.execute({"input_documents": input_documents}, query) + interrupted = False + finally: + self._finalize_workflow_run( + workflow_owner_id, + run_user_id, + pg_workflow_id, + query, + interrupted=interrupted, + ) def _load_workflow_graph(self) -> Optional[WorkflowGraph]: if self._workflow_data: @@ -167,9 +165,7 @@ class WorkflowAgent(BaseAgent): if not owner_id and isinstance(self.decoded_token, dict): owner_id = self.decoded_token.get("sub") if not owner_id: - logger.error( - f"Workflow owner not available for workflow load: {self.workflow_id}" - ) + logger.error(f"Workflow owner not available for workflow load: {self.workflow_id}") return None with db_readonly() as conn: @@ -179,10 +175,7 @@ class WorkflowAgent(BaseAgent): else: workflow_row = wf_repo.get_by_legacy_id(self.workflow_id, owner_id) if workflow_row is None: - logger.error( - f"Workflow {self.workflow_id} not found or inaccessible " - f"for user {owner_id}" - ) + logger.error(f"Workflow {self.workflow_id} not found or inaccessible for user {owner_id}") return None pg_workflow_id = str(workflow_row["id"]) graph_version = workflow_row.get("current_graph_version", 1) @@ -194,10 +187,12 @@ class WorkflowAgent(BaseAgent): graph_version = 1 node_rows = WorkflowNodesRepository(conn).find_by_version( - pg_workflow_id, graph_version, + pg_workflow_id, + graph_version, ) edge_rows = WorkflowEdgesRepository(conn).find_by_version( - pg_workflow_id, graph_version, + pg_workflow_id, + graph_version, ) workflow = Workflow( @@ -250,9 +245,7 @@ class WorkflowAgent(BaseAgent): """ return getattr(self, "initial_user_id", None) or getattr(self, "user", None) or workflow_owner_id - def _resolve_owned_workflow_pg_id( - self, conn: Any, owner_id: Optional[str] - ) -> Optional[str]: + def _resolve_owned_workflow_pg_id(self, conn: Any, owner_id: Optional[str]) -> Optional[str]: """Return the owned workflow's PG id, or None for an unowned/draft id.""" if not self.workflow_id or not owner_id: return None @@ -278,9 +271,7 @@ class WorkflowAgent(BaseAgent): return None try: with db_session() as conn: - pg_workflow_id = self._resolve_owned_workflow_pg_id( - conn, workflow_owner_id - ) + pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, workflow_owner_id) if pg_workflow_id is None: return None WorkflowRunsRepository(conn).create( @@ -404,19 +395,24 @@ class WorkflowAgent(BaseAgent): run_user_id: Optional[str], pg_workflow_id: Optional[str], query: str, + interrupted: bool = False, ) -> None: """Write the run's terminal status/result; upsert the row if pre-creation was skipped. The run is owned by the *runner* (so it stays readable to the caller and matches the pre-created row); the workflow row is resolved by its *owner*. + When ``interrupted`` is set (client disconnect / mid-run error), the run is + recorded as FAILED regardless of the per-node log, so a partial run is never + left looking complete. """ if not self._engine: return try: + status = ExecutionStatus.FAILED if interrupted else self._determine_run_status() run = WorkflowRun( workflow_id=self.workflow_id or "unknown", user=run_user_id, - status=self._determine_run_status(), + status=status, inputs={"query": query}, outputs=self._serialize_state(self._engine.state), steps=self._engine.get_execution_summary(), @@ -429,9 +425,7 @@ class WorkflowAgent(BaseAgent): return with db_session() as conn: if pg_workflow_id is None: - pg_workflow_id = self._resolve_owned_workflow_pg_id( - conn, workflow_owner_id - ) + pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, workflow_owner_id) if pg_workflow_id is None: return runs_repo = WorkflowRunsRepository(conn) @@ -486,10 +480,7 @@ class WorkflowAgent(BaseAgent): def _serialize_state_value(self, value: Any) -> Any: if isinstance(value, dict): - return { - str(dict_key): self._serialize_state_value(dict_value) - for dict_key, dict_value in value.items() - } + return {str(dict_key): self._serialize_state_value(dict_value) for dict_key, dict_value in value.items()} if isinstance(value, list): return [self._serialize_state_value(item) for item in value] if isinstance(value, tuple): diff --git a/application/api/user/artifacts/routes.py b/application/api/user/artifacts/routes.py index 96911794..d4b8a350 100644 --- a/application/api/user/artifacts/routes.py +++ b/application/api/user/artifacts/routes.py @@ -38,6 +38,9 @@ artifacts_ns = Namespace("artifacts", description="Artifact operations", path="/ _PRESIGNED_URL_TTL = 300 +_ARTIFACT_URL_ENVELOPE_MIME = "application/vnd.docsgpt.artifact-url+json" + + def _sanitize_header_filename(filename: Optional[str], fallback: str) -> str: """Strip CRLF / quotes from a display filename for a Content-Disposition header.""" if not filename: @@ -50,21 +53,9 @@ def _artifact_summary(artifact: dict) -> dict: """Project an artifact identity row to its API metadata shape (owner id withheld).""" return { "id": str(artifact.get("id")), - "conversation_id": ( - str(artifact["conversation_id"]) - if artifact.get("conversation_id") is not None - else None - ), - "workflow_run_id": ( - str(artifact["workflow_run_id"]) - if artifact.get("workflow_run_id") is not None - else None - ), - "message_id": ( - str(artifact["message_id"]) - if artifact.get("message_id") is not None - else None - ), + "conversation_id": (str(artifact["conversation_id"]) if artifact.get("conversation_id") is not None else None), + "workflow_run_id": (str(artifact["workflow_run_id"]) if artifact.get("workflow_run_id") is not None else None), + "message_id": (str(artifact["message_id"]) if artifact.get("message_id") is not None else None), "kind": artifact.get("kind"), "title": artifact.get("title"), "metadata": artifact.get("metadata"), @@ -107,20 +98,14 @@ class ListArtifacts(Resource): share_token = request.args.get("share_token") if not user_id and not share_token: - return make_response( - jsonify({"success": False, "message": "Authentication required"}), 401 - ) + return make_response(jsonify({"success": False, "message": "Authentication required"}), 401) # Gate UUID-shape before any CAST(:id AS uuid) reaches the repo, so a # malformed id is rejected cleanly instead of poisoning the transaction. if conversation_id and not looks_like_uuid(conversation_id): - return make_response( - jsonify({"success": False, "message": "Invalid conversation_id"}), 400 - ) + return make_response(jsonify({"success": False, "message": "Invalid conversation_id"}), 400) if workflow_run_id and not looks_like_uuid(workflow_run_id): - return make_response( - jsonify({"success": False, "message": "Invalid workflow_run_id"}), 400 - ) + return make_response(jsonify({"success": False, "message": "Invalid workflow_run_id"}), 400) try: with db_readonly() as conn: @@ -133,49 +118,32 @@ class ListArtifacts(Resource): # and scope the query to it. Workflow runs are not agent-scoped, # so reject that filter. if workflow_run_id: - return make_response( - jsonify({"success": False, "message": "Forbidden"}), 403 - ) + return make_response(jsonify({"success": False, "message": "Forbidden"}), 403) if not conversation_id: - return make_response( - jsonify({"success": False, "message": "Forbidden"}), 403 - ) - rows = repo.list_artifacts_for_agent( - principal.agent_id, user_id, conversation_id=conversation_id - ) + return make_response(jsonify({"success": False, "message": "Forbidden"}), 403) + rows = repo.list_artifacts_for_agent(principal.agent_id, user_id, conversation_id=conversation_id) elif conversation_id: - if not user_can_access_conversation( - conn, conversation_id, user_id, share_token - ): - return make_response( - jsonify({"success": False, "message": "Forbidden"}), 403 - ) + if not user_can_access_conversation(conn, conversation_id, user_id, share_token): + return make_response(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 - ) + 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 - ) + 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 + 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: - return make_response( - jsonify({"success": False, "message": "Forbidden"}), 403 - ) + return make_response(jsonify({"success": False, "message": "Forbidden"}), 403) rows = repo.list_artifacts(workflow_run_id=workflow_run_id) else: if not user_id: @@ -186,9 +154,7 @@ class ListArtifacts(Resource): rows = repo.list_artifacts(user_id=user_id) return make_response( - jsonify( - {"success": True, "artifacts": [_artifact_summary(r) for r in rows]} - ), + jsonify({"success": True, "artifacts": [_artifact_summary(r) for r in rows]}), 200, ) except Exception as err: @@ -201,22 +167,16 @@ class GetArtifact(Resource): @api.doc(description="Get an artifact's metadata, version list, and current spec") def get(self, artifact_id: str): if not looks_like_uuid(artifact_id): - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) principal = resolve_principal() try: with db_readonly() as conn: repo = ArtifactsRepository(conn) artifact = repo.get_artifact(artifact_id) if artifact is None: - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) if not authorize_artifact(conn, artifact, principal): - return make_response( - jsonify({"success": False, "message": "Forbidden"}), 403 - ) + return make_response(jsonify({"success": False, "message": "Forbidden"}), 403) versions = repo.list_versions(artifact_id) current = repo.get_version(artifact_id, artifact.get("current_version")) @@ -231,24 +191,18 @@ class GetArtifact(Resource): @api.doc(description="Delete an artifact and all its versions (owner only)") def delete(self, artifact_id: str): if not looks_like_uuid(artifact_id): - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) principal = resolve_principal() try: with db_session() as conn: repo = ArtifactsRepository(conn) artifact = repo.get_artifact(artifact_id) if artifact is None: - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) # Delete is a WRITE: only the parent owner may delete; share # links / read-only collaborators are denied (read access only). if not authorize_artifact_write(conn, artifact, principal): - return make_response( - jsonify({"success": False, "message": "Forbidden"}), 403 - ) + return make_response(jsonify({"success": False, "message": "Forbidden"}), 403) storage_paths = repo.delete_artifact(artifact_id) # Reap the bytes best-effort AFTER the row delete commits. _reap_storage(storage_paths) @@ -271,9 +225,7 @@ def _reap_storage(paths: list) -> None: try: storage.delete_file(path) except Exception: - current_app.logger.warning( - "artifact delete: failed to delete bytes %s", path, exc_info=True - ) + current_app.logger.warning("artifact delete: failed to delete bytes %s", path, exc_info=True) @artifacts_ns.route("/artifacts//versions/") @@ -281,37 +233,25 @@ class GetArtifactVersion(Resource): @api.doc(description="Get a single artifact version's metadata and spec") def get(self, artifact_id: str, version: int): if not looks_like_uuid(artifact_id): - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) principal = resolve_principal() try: with db_readonly() as conn: repo = ArtifactsRepository(conn) artifact = repo.get_artifact(artifact_id) if artifact is None: - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) if not authorize_artifact(conn, artifact, principal): - return make_response( - jsonify({"success": False, "message": "Forbidden"}), 403 - ) + return make_response(jsonify({"success": False, "message": "Forbidden"}), 403) version_row = repo.get_version(artifact_id, version) if version_row is None: - return make_response( - jsonify({"success": False, "message": "Version not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Version not found"}), 404) return make_response( - jsonify( - {"success": True, "version": _version_summary(version_row, include_spec=True)} - ), + jsonify({"success": True, "version": _version_summary(version_row, include_spec=True)}), 200, ) except Exception as err: - current_app.logger.error( - f"Error retrieving artifact version: {err}", exc_info=True - ) + current_app.logger.error(f"Error retrieving artifact version: {err}", exc_info=True) return make_response(jsonify({"success": False}), 400) @@ -320,9 +260,7 @@ class DownloadArtifact(Resource): @api.doc(description="Download an artifact's bytes (302 to a presigned URL on S3)") def get(self, artifact_id: str): if not looks_like_uuid(artifact_id): - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) principal = resolve_principal() version_arg = request.args.get("version") try: @@ -330,37 +268,25 @@ class DownloadArtifact(Resource): repo = ArtifactsRepository(conn) artifact = repo.get_artifact(artifact_id) if artifact is None: - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) if not authorize_artifact(conn, artifact, principal): - return make_response( - jsonify({"success": False, "message": "Forbidden"}), 403 - ) + return make_response(jsonify({"success": False, "message": "Forbidden"}), 403) version = artifact.get("current_version") if version_arg is not None: try: version = int(version_arg) except ValueError: - return make_response( - jsonify({"success": False, "message": "Invalid version"}), 400 - ) + return make_response(jsonify({"success": False, "message": "Invalid version"}), 400) version_row = repo.get_version(artifact_id, version) if version_row is None: - return make_response( - jsonify({"success": False, "message": "Version not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Version not found"}), 404) # The object key is derived only from the stored path, never client input. storage_path = version_row.get("storage_path") if not storage_path: - return make_response( - jsonify({"success": False, "message": "No file for this version"}), 404 - ) + return make_response(jsonify({"success": False, "message": "No file for this version"}), 404) - filename = _sanitize_header_filename( - version_row.get("filename"), f"artifact-{artifact_id}" - ) + filename = _sanitize_header_filename(version_row.get("filename"), f"artifact-{artifact_id}") mime_type = version_row.get("mime_type") or "application/octet-stream" storage = StorageCreator.get_storage() @@ -370,18 +296,14 @@ class DownloadArtifact(Resource): # the operator expected to be off the hot path. if getattr(settings, "URL_STRATEGY", "backend") == "s3": try: - url = storage.generate_presigned_url( - storage_path, expires_in=_PRESIGNED_URL_TTL - ) + url = storage.generate_presigned_url(storage_path, expires_in=_PRESIGNED_URL_TTL) except NotImplementedError: current_app.logger.error( "URL_STRATEGY=s3 but %s cannot mint presigned URLs", type(storage).__name__, ) return make_response( - jsonify( - {"success": False, "message": "Storage misconfigured"} - ), + jsonify({"success": False, "message": "Storage misconfigured"}), 500, ) # A 302 to a cross-origin S3 URL can't be read by the app's authed @@ -392,7 +314,17 @@ class DownloadArtifact(Resource): if request.args.get("disposition") == "url" or ( "application/json" in request.headers.get("Accept", "") ): - return make_response(jsonify({"success": True, "url": url}), 200) + resp = make_response(jsonify({"success": True, "url": url}), 200) + # Tag the envelope with a distinctive media type so the client + # keys off a server-set signal, not the JSON body shape. A + # ``data`` artifact whose bytes happen to be + # ``{"success":true,"url":"..."}`` is streamed under the backend + # strategy with its own (``application/json``) content-type and + # can never carry this vendor type, so it can't be mistaken for a + # redirect (open-redirect gadget). Content-Type is CORS-safelisted, + # so this stays readable cross-origin without expose-headers. + resp.headers["Content-Type"] = _ARTIFACT_URL_ENVELOPE_MIME + return resp return redirect(url, code=302) # Stream the bytes in chunks instead of buffering the whole object in @@ -411,18 +343,12 @@ class DownloadArtifact(Resource): return Response( stream_with_context(_stream()), mimetype=mime_type, - headers={ - "Content-Disposition": f'attachment; filename="{filename}"' - }, + headers={"Content-Disposition": f'attachment; filename="{filename}"'}, ) except FileNotFoundError: - return make_response( - jsonify({"success": False, "message": "File not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "File not found"}), 404) except Exception as err: - current_app.logger.error( - f"Error downloading artifact: {err}", exc_info=True - ) + current_app.logger.error(f"Error downloading artifact: {err}", exc_info=True) return make_response(jsonify({"success": False}), 400) @@ -431,43 +357,31 @@ class RestoreArtifact(Resource): @api.doc(description="Restore a prior version by appending it as the new current version") def post(self, artifact_id: str): if not looks_like_uuid(artifact_id): - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) principal = resolve_principal() data = request.get_json(silent=True) or {} target_version = data.get("version") if target_version is None: - return make_response( - jsonify({"success": False, "message": "Missing version"}), 400 - ) + return make_response(jsonify({"success": False, "message": "Missing version"}), 400) try: target_version = int(target_version) except (ValueError, TypeError): - return make_response( - jsonify({"success": False, "message": "Invalid version"}), 400 - ) + return make_response(jsonify({"success": False, "message": "Invalid version"}), 400) try: with db_session() as conn: repo = ArtifactsRepository(conn) artifact = repo.get_artifact(artifact_id) if artifact is None: - return make_response( - jsonify({"success": False, "message": "Artifact not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Artifact not found"}), 404) # Restore is a WRITE (it appends a new current version); share # links / shared_with collaborators inherit read access only, so # gate on the stricter owner-required write check. if not authorize_artifact_write(conn, artifact, principal): - return make_response( - jsonify({"success": False, "message": "Forbidden"}), 403 - ) + return make_response(jsonify({"success": False, "message": "Forbidden"}), 403) source = repo.get_version(artifact_id, target_version) if source is None: - return make_response( - jsonify({"success": False, "message": "Version not found"}), 404 - ) + return make_response(jsonify({"success": False, "message": "Version not found"}), 404) new_version = repo.append_version( artifact_id, mime_type=source.get("mime_type"), @@ -480,9 +394,7 @@ class RestoreArtifact(Resource): produced_by=source.get("produced_by"), ) return make_response( - jsonify( - {"success": True, "version": _version_summary(new_version, include_spec=True)} - ), + jsonify({"success": True, "version": _version_summary(new_version, include_spec=True)}), 200, ) except Exception as err: diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index bd3d644a..00efb718 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -395,6 +395,33 @@ def reap_sandbox_sessions(self): return {"reaped": len(reaped)} +@celery.task(bind=True, acks_late=False) +def reap_stale_workflow_runs(self): + """Fail workflow runs stranded in ``running`` past the stale deadline. + + A run row is pre-created as ``running`` and finalized when its generator + finishes; a client disconnect or worker crash can leave it ``running`` + forever. This closes those rows out so the UI/API stop showing a run that + will never complete. + """ + from datetime import datetime, timezone + + from application.core.settings import settings + from application.storage.db.engine import get_engine + from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository + + try: + stale_seconds = max(60, int(settings.WORKFLOW_RUN_STALE_SECONDS)) + cutoff = datetime.now(timezone.utc) - timedelta(seconds=stale_seconds) + engine = get_engine() + with engine.begin() as conn: + reaped = WorkflowRunsRepository(conn).mark_stale_running_failed(cutoff) + except Exception: # noqa: BLE001 - housekeeping must never crash the beat loop + logging.getLogger(__name__).exception("reap_stale_workflow_runs failed") + return {"reaped": 0, "error": True} + return {"reaped": reaped} + + @celery.on_after_configure.connect def setup_periodic_tasks(sender, **kwargs): from application.core.settings import settings @@ -468,6 +495,13 @@ def setup_periodic_tasks(sender, **kwargs): reap_sandbox_sessions.s(), name="reap-sandbox-sessions", ) + # Fail workflow runs stranded in ``running`` (client disconnect / crash) so + # they don't linger forever. Every few minutes is plenty; the cutoff is hours. + sender.add_periodic_task( + timedelta(seconds=300), + reap_stale_workflow_runs.s(), + name="reap-stale-workflow-runs", + ) # Bound time limits so a hung OAuth discovery (user never finishes the diff --git a/application/core/settings.py b/application/core/settings.py index 88fcd108..e02c83ed 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -251,18 +251,10 @@ 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. ``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", - "artifact_generator", ] # Conversation Compression Settings @@ -358,6 +350,8 @@ class Settings(BaseSettings): DOCUMENT_PARSE_QUEUE: str = "parsing" # queue the parse_document task is routed to DOCUMENT_PARSE_TIMEOUT: int = 120 # seconds the tool awaits the enqueued parse before degrading DOCUMENT_PARSE_MAX_BYTES: int = 0 # cap on a parsed document's bytes (0 = reuse SANDBOX_MAX_INPUT_BYTES) + DOCUMENT_MAX_DECOMPRESSED_BYTES: int = 300 * 1024 * 1024 + DOCUMENT_MAX_ARCHIVE_ENTRIES: int = 10000 # Per-agent-node cap on files passed natively to the node's LLM (vision/doc # inputs). Files past the cap are extracted to text or dropped, not attached # natively, to bound context/cost. Re-uses SANDBOX_MAX_INPUT_BYTES per file. @@ -368,6 +362,11 @@ class Settings(BaseSettings): # avoid serializing dozens of parses; documents past the cap are skipped with # a truncation note instead of extracted. WORKFLOW_NODE_EXTRACT_MAX_FILES: int = 5 + # A workflow run row is pre-created as ``running`` and finalized when its + # generator completes; a client disconnect or worker crash can strand it in + # ``running`` forever. The beat reaper fails runs still ``running`` past this + # many seconds. Generous so a legitimately long run is never cut off. + WORKFLOW_RUN_STALE_SECONDS: int = 3600 # Runner container resource caps — consumed by the docsgpt-sandbox compose # service (deployment/sandbox), not by the app client. cgroup CPU/mem caps # are part of the untrusted-code security boundary. diff --git a/application/parser/document_reader.py b/application/parser/document_reader.py index 4524394d..c9483759 100644 --- a/application/parser/document_reader.py +++ b/application/parser/document_reader.py @@ -9,9 +9,11 @@ temp cleanup — so a hostile filename or document is treated as inert data. from __future__ import annotations +import io import logging import os import tempfile +import zipfile from pathlib import Path from typing import Any, Dict, List, Optional @@ -98,6 +100,31 @@ def _max_input_bytes() -> int: return int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024)) +_ZIP_CONTAINER_EXTENSIONS = frozenset({".docx", ".xlsx", ".pptx", ".epub"}) + + +def _reject_zip_bomb(data: bytes, suffix: str) -> Optional[str]: + """Return an error string if a zip-based document declares an implausible expansion, else None.""" + if suffix not in _ZIP_CONTAINER_EXTENSIONS: + return None + max_entries = int(getattr(settings, "DOCUMENT_MAX_ARCHIVE_ENTRIES", 10000)) + cap = int(getattr(settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 300 * 1024 * 1024)) + try: + with zipfile.ZipFile(io.BytesIO(data)) as zf: + infos = zf.infolist() + if len(infos) > max_entries: + return f"document archive has too many entries ({len(infos)} > {max_entries})." + total = 0 + for info in infos: + total += info.file_size + if total > cap: + return f"document decompresses to too much data: exceeds the {cap}-byte cap." + except zipfile.BadZipFile: + # Not a readable zip; the format-specific parser will surface a clean error. + return None + return None + + def _resolve_ocr_enabled(ocr: str) -> bool: """Resolve the OCR flag from the ``ocr`` arg and the deployment setting.""" if ocr == "on": @@ -178,9 +205,7 @@ def _compact_table(table: Dict[str, Any]) -> Dict[str, Any]: return compact -def _docling_structured( - path: Path, *, ocr_enabled: bool, include_tables: bool, parser: Any = None -) -> Dict[str, Any]: +def _docling_structured(path: Path, *, ocr_enabled: bool, include_tables: bool, parser: Any = None) -> Dict[str, Any]: """Convert a document with Docling and return markdown + structured dict + bounded tables. When ``parser`` is the configured ``DoclingParser`` (the collapse-the-double-convert @@ -307,6 +332,10 @@ def parse_document_bytes( if len(data) > cap: return {"error": f"input document is too large: {len(data)} bytes exceeds the {cap}-byte cap."} + bomb = _reject_zip_bomb(data, suffix) + if bomb is not None: + return {"error": bomb} + ocr_enabled = _resolve_ocr_enabled(ocr) tmp_dir = tempfile.mkdtemp(prefix="docparse-") tmp_path = Path(tmp_dir) / safe_name @@ -364,9 +393,7 @@ def _shape( # (Docling/torch conversion dominates the cost, so a re-convert ~doubles latency). if wants_tables and _is_docling_parser(parser): try: - extracted = _docling_structured( - path, ocr_enabled=ocr_enabled, include_tables=True, parser=parser - ) + extracted = _docling_structured(path, ocr_enabled=ocr_enabled, include_tables=True, parser=parser) text = extracted["markdown"] tables: List[Dict[str, Any]] = [_compact_table(t) for t in extracted["tables"]] except Exception: @@ -392,8 +419,10 @@ def _shape( tables: List[Dict[str, Any]] = [] if wants_tables: try: - tables = [_compact_table(t) for t in _docling_structured( - path, ocr_enabled=ocr_enabled, include_tables=True)["tables"]] + tables = [ + _compact_table(t) + for t in _docling_structured(path, ocr_enabled=ocr_enabled, include_tables=True)["tables"] + ] except Exception: tables = [] bounded, truncated = _bounded(text) diff --git a/application/prompts/agentic/creative.txt b/application/prompts/agentic/creative.txt index e88cec0f..5a610875 100644 --- a/application/prompts/agentic/creative.txt +++ b/application/prompts/agentic/creative.txt @@ -38,5 +38,7 @@ The user attached these files to this message: {% for f in attachments.files %} - {{ f.filename }} ({{ f.mime_type }}{% if f.size %}, {{ f.size }} bytes{% endif %}) {% endfor %} -You can pass any attached file to the code tool by its name. If a file is not one you can read directly, use the code tool to read or process it. +{% if tools.enabled is not defined or 'code_executor' in tools.enabled or 'read_document' in tools.enabled %} +You can pass any attached file to a code or document tool by its name to read or process it. +{% endif %} {% endif %} diff --git a/application/prompts/agentic/default.txt b/application/prompts/agentic/default.txt index a25aafed..eefa822b 100644 --- a/application/prompts/agentic/default.txt +++ b/application/prompts/agentic/default.txt @@ -37,5 +37,7 @@ The user attached these files to this message: {% for f in attachments.files %} - {{ f.filename }} ({{ f.mime_type }}{% if f.size %}, {{ f.size }} bytes{% endif %}) {% endfor %} -You can pass any attached file to the code tool by its name. If a file is not one you can read directly, use the code tool to read or process it. +{% if tools.enabled is not defined or 'code_executor' in tools.enabled or 'read_document' in tools.enabled %} +You can pass any attached file to a code or document tool by its name to read or process it. +{% endif %} {% endif %} diff --git a/application/prompts/agentic/strict.txt b/application/prompts/agentic/strict.txt index c50a2718..144ed88f 100644 --- a/application/prompts/agentic/strict.txt +++ b/application/prompts/agentic/strict.txt @@ -37,5 +37,7 @@ The user attached these files to this message: {% for f in attachments.files %} - {{ f.filename }} ({{ f.mime_type }}{% if f.size %}, {{ f.size }} bytes{% endif %}) {% endfor %} -You can pass any attached file to the code tool by its name. If a file is not one you can read directly, use the code tool to read or process it. +{% if tools.enabled is not defined or 'code_executor' in tools.enabled or 'read_document' in tools.enabled %} +You can pass any attached file to a code or document tool by its name to read or process it. +{% endif %} {% endif %} diff --git a/application/prompts/chat_combine_creative.txt b/application/prompts/chat_combine_creative.txt index 0f46fd27..8420f852 100644 --- a/application/prompts/chat_combine_creative.txt +++ b/application/prompts/chat_combine_creative.txt @@ -48,5 +48,7 @@ The user attached these files to this message: {% for f in attachments.files %} - {{ f.filename }} ({{ f.mime_type }}{% if f.size %}, {{ f.size }} bytes{% endif %}) {% endfor %} -You can pass any attached file to the code tool by its name. If a file is not one you can read directly, use the code tool to read or process it. +{% if tools.enabled is not defined or 'code_executor' in tools.enabled or 'read_document' in tools.enabled %} +You can pass any attached file to a code or document tool by its name to read or process it. +{% endif %} {% endif %} diff --git a/application/prompts/chat_combine_default.txt b/application/prompts/chat_combine_default.txt index b7e5c7ba..d6b448e7 100644 --- a/application/prompts/chat_combine_default.txt +++ b/application/prompts/chat_combine_default.txt @@ -47,5 +47,7 @@ The user attached these files to this message: {% for f in attachments.files %} - {{ f.filename }} ({{ f.mime_type }}{% if f.size %}, {{ f.size }} bytes{% endif %}) {% endfor %} -You can pass any attached file to the code tool by its name. If a file is not one you can read directly, use the code tool to read or process it. +{% if tools.enabled is not defined or 'code_executor' in tools.enabled or 'read_document' in tools.enabled %} +You can pass any attached file to a code or document tool by its name to read or process it. +{% endif %} {% endif %} diff --git a/application/prompts/chat_combine_strict.txt b/application/prompts/chat_combine_strict.txt index 74d730f5..aa12d5f5 100644 --- a/application/prompts/chat_combine_strict.txt +++ b/application/prompts/chat_combine_strict.txt @@ -47,5 +47,7 @@ The user attached these files to this message: {% for f in attachments.files %} - {{ f.filename }} ({{ f.mime_type }}{% if f.size %}, {{ f.size }} bytes{% endif %}) {% endfor %} -You can pass any attached file to the code tool by its name. If a file is not one you can read directly, use the code tool to read or process it. +{% if tools.enabled is not defined or 'code_executor' in tools.enabled or 'read_document' in tools.enabled %} +You can pass any attached file to a code or document tool by its name to read or process it. +{% endif %} {% endif %} diff --git a/application/sandbox/artifacts_capture.py b/application/sandbox/artifacts_capture.py index 07270a84..e484e8d4 100644 --- a/application/sandbox/artifacts_capture.py +++ b/application/sandbox/artifacts_capture.py @@ -149,6 +149,7 @@ def capture_artifacts( user_id: str, conversation_id: Optional[str] = None, workflow_run_id: Optional[str] = None, + message_id: Optional[str] = None, produced_by: Optional[Dict[str, Any]] = None, outputs: Optional[List[str]] = None, ) -> List[Dict[str, Any]]: @@ -198,6 +199,7 @@ def capture_artifacts( user_id=user_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id, + message_id=message_id, produced_by=produced_by, ) except QuotaExceeded: @@ -217,6 +219,7 @@ def persist_artifact( user_id: str, conversation_id: Optional[str] = None, workflow_run_id: Optional[str] = None, + message_id: Optional[str] = None, produced_by: Optional[Dict[str, Any]] = None, ) -> Optional[Dict[str, Any]]: """Store a captured workspace file as a new artifact (kind/mime inferred from its name).""" @@ -234,6 +237,7 @@ def persist_artifact( title=display_name, conversation_id=conversation_id, workflow_run_id=workflow_run_id, + message_id=message_id, produced_by=produced_by, ) @@ -341,6 +345,7 @@ def persist_new_artifact( title: Optional[str] = None, conversation_id: Optional[str] = None, workflow_run_id: Optional[str] = None, + message_id: Optional[str] = None, spec: Any = None, preview_text: Optional[str] = None, produced_by: Any = None, @@ -366,6 +371,7 @@ def persist_new_artifact( kind, conversation_id=conversation_id, workflow_run_id=workflow_run_id, + message_id=message_id, title=title or safe_name, mime_type=mime_type, filename=safe_name, diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py index 88fd7643..2c4b44f6 100644 --- a/application/storage/db/repositories/artifacts.py +++ b/application/storage/db/repositories/artifacts.py @@ -392,13 +392,29 @@ class ArtifactsRepository: return int(row[0]) if row is not None else 0 def total_bytes_for_user(self, user_id: str) -> int: - """Return the summed byte size of every version a user owns (quota accounting).""" + """Return the stored byte size a user owns (quota accounting). + + Deduplicated by ``storage_path``: ``restore`` appends a version that + re-points at an existing version's stored object (same key + size), so + summing every version row would charge the same bytes N times and inflate + usage without any new bytes being written. Each distinct stored object is + counted once; versions with no object yet (``storage_path IS NULL`` -- + spec-only, whose size is normally NULL too) are counted individually since + they share no key to dedupe on. + """ row = self._conn.execute( text( - "SELECT COALESCE(SUM(v.size), 0) " - "FROM artifact_versions v " - "JOIN artifacts a ON a.id = v.artifact_id " - "WHERE a.user_id = :user_id" + "SELECT COALESCE(SUM(t.size), 0) FROM (" + " SELECT DISTINCT v.storage_path, v.size " + " FROM artifact_versions v " + " JOIN artifacts a ON a.id = v.artifact_id " + " WHERE a.user_id = :user_id AND v.storage_path IS NOT NULL " + " UNION ALL " + " SELECT v.storage_path, v.size " + " FROM artifact_versions v " + " JOIN artifacts a ON a.id = v.artifact_id " + " WHERE a.user_id = :user_id AND v.storage_path IS NULL" + ") t" ), {"user_id": user_id}, ).fetchone() diff --git a/application/storage/db/repositories/workflow_runs.py b/application/storage/db/repositories/workflow_runs.py index 8d4ade7f..25e105e9 100644 --- a/application/storage/db/repositories/workflow_runs.py +++ b/application/storage/db/repositories/workflow_runs.py @@ -9,7 +9,7 @@ from __future__ import annotations from datetime import datetime from typing import Optional -from sqlalchemy import Connection, text +from sqlalchemy import Connection, func, text from sqlalchemy.dialects.postgresql import insert as pg_insert from application.storage.db.base_repository import row_to_dict @@ -87,6 +87,28 @@ class WorkflowRunsRepository: res = self._conn.execute(stmt) return res.rowcount > 0 + def mark_stale_running_failed(self, older_than: datetime) -> int: + """Fail runs left in ``running`` (no ``ended_at``) since before ``older_than``. + + The run row is pre-created as ``running`` and finalized when its generator + finishes; a client disconnect or a worker crash can strand it in ``running`` + forever, since nothing else finalizes it. This closes those rows out so they + don't linger. Returns the number of rows updated. + """ + stmt = ( + workflow_runs_table.update() + .where(workflow_runs_table.c.status == "running") + .where(workflow_runs_table.c.ended_at.is_(None)) + .where(workflow_runs_table.c.started_at < older_than) + .values( + status="failed", + ended_at=func.now(), + result={"error": "Run did not complete (timed out or the client disconnected)."}, + ) + ) + res = self._conn.execute(stmt) + return res.rowcount + def get(self, run_id: str) -> Optional[dict]: res = self._conn.execute( text("SELECT * FROM workflow_runs WHERE id = CAST(:id AS uuid)"), diff --git a/application/worker.py b/application/worker.py index d7a05a67..5bd8abb7 100755 --- a/application/worker.py +++ b/application/worker.py @@ -1718,6 +1718,7 @@ def _persist_parse_result(result, title, user_id, parent, options): title=f"{title} (parsed)", conversation_id=parent.get("conversation_id"), workflow_run_id=parent.get("workflow_run_id"), + message_id=parent.get("message_id"), produced_by={"tool": "read_document", "action": "read_document", "tool_id": options.get("tool_id")}, ) except QuotaExceeded as exc: diff --git a/deployment/docker-compose-azure.yaml b/deployment/docker-compose-azure.yaml index 07bf9cdd..3295a12c 100644 --- a/deployment/docker-compose-azure.yaml +++ b/deployment/docker-compose-azure.yaml @@ -18,21 +18,8 @@ services: - CELERY_BROKER_URL=redis://redis:6379/0 - CELERY_RESULT_BACKEND=redis://redis:6379/1 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt - # Code-execution runner reached over HTTP + WebSocket (no docker socket). - - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 - # Env-scrubbing kernelspec the docsgpt-sandbox image installs. - - SANDBOX_KERNEL_NAME=docsgpt-python ports: - "7091:7091" - 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 @@ -55,75 +42,32 @@ services: - CELERY_RESULT_BACKEND=redis://redis:6379/1 - API_URL=http://backend:7091 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt - - 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 - # 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 - read_only: true - environment: - # Keep Jupyter's runtime/connection files on the writable tmpfs. - - JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime - - JUPYTER_DATA_DIR=/tmp/jupyter-data - tmpfs: - - /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 - - sandbox-egress + # NOTE: sandboxed code execution (the `code_executor` / `artifact_generator` + # tools) is an OPT-IN feature and is intentionally NOT part of this default + # stack. To run it, layer the sandbox overlay on top of this file: + # docker compose -f docker-compose-azure.yaml \ + # -f optional/docker-compose.optional.sandbox.yaml up + # See deployment/sandbox/README.md. redis: image: redis:6-alpine - # Data plane only: kept off the sandbox's networks (a reachable, unauthed - # broker = Celery task injection -> worker RCE). - networks: - - data-net ports: - # 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 + - 6379:6379 postgres: image: postgres:16-alpine - # Data plane only (see redis): kept off the sandbox's networks. - networks: - - data-net environment: - POSTGRES_USER=docsgpt - POSTGRES_PASSWORD=docsgpt - POSTGRES_DB=docsgpt ports: - # Loopback only (see redis): not reachable from the sandbox via the gateway. - - "127.0.0.1:5432:5432" + - "5432:5432" volumes: - postgres_data:/var/lib/postgresql/data healthcheck: @@ -132,21 +76,5 @@ services: timeout: 5s retries: 10 -networks: - # 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: - internal: true - volumes: postgres_data: diff --git a/deployment/docker-compose-hub.yaml b/deployment/docker-compose-hub.yaml index 7d950473..df41d038 100644 --- a/deployment/docker-compose-hub.yaml +++ b/deployment/docker-compose-hub.yaml @@ -23,21 +23,8 @@ services: - CELERY_RESULT_BACKEND=redis://redis:6379/1 - CACHE_REDIS_URL=redis://redis:6379/2 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt - # Code-execution runner reached over HTTP + WebSocket (no docker socket). - - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 - # Env-scrubbing kernelspec the docsgpt-sandbox image installs. - - SANDBOX_KERNEL_NAME=docsgpt-python ports: - "7091:7091" - 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 @@ -62,14 +49,6 @@ services: - API_URL=http://backend:7091 - CACHE_REDIS_URL=redis://redis:6379/2 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt - - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 - - 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 @@ -80,64 +59,26 @@ services: postgres: condition: service_healthy - # 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 - read_only: true - environment: - # Keep Jupyter's runtime/connection files on the writable tmpfs. - - JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime - - JUPYTER_DATA_DIR=/tmp/jupyter-data - tmpfs: - - /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 - - sandbox-egress + # NOTE: sandboxed code execution (the `code_executor` / `artifact_generator` + # tools) is an OPT-IN feature and is intentionally NOT part of this default + # stack. To run it, layer the sandbox overlay on top of this file: + # docker compose -f docker-compose-hub.yaml \ + # -f optional/docker-compose.optional.sandbox.yaml up + # See deployment/sandbox/README.md. redis: image: redis:6-alpine - # Data plane only: kept off the sandbox's networks (a reachable, unauthed - # broker = Celery task injection -> worker RCE). - networks: - - data-net ports: - # 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 + - 6379:6379 postgres: image: postgres:16-alpine - # Data plane only (see redis): kept off the sandbox's networks. - networks: - - data-net environment: - POSTGRES_USER=docsgpt - POSTGRES_PASSWORD=docsgpt - POSTGRES_DB=docsgpt ports: - # Loopback only (see redis): not reachable from the sandbox via the gateway. - - "127.0.0.1:5432:5432" + - "5432:5432" volumes: - postgres_data:/var/lib/postgresql/data healthcheck: @@ -146,21 +87,5 @@ services: timeout: 5s retries: 10 -networks: - # 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: - internal: true - volumes: postgres_data: diff --git a/deployment/docker-compose.yaml b/deployment/docker-compose.yaml index 193f156c..b027ac1a 100644 --- a/deployment/docker-compose.yaml +++ b/deployment/docker-compose.yaml @@ -24,23 +24,8 @@ services: - CELERY_RESULT_BACKEND=redis://redis:6379/1 - CACHE_REDIS_URL=redis://redis:6379/2 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt - # Code-execution runner reached over HTTP + WebSocket (no docker socket). - - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 - # Select the runner's env-scrubbing kernelspec (distinct name; never - # shadowed by the stock "python3" spec). Must match the kernel the - # docsgpt-sandbox image installs. - - SANDBOX_KERNEL_NAME=docsgpt-python ports: - "7091:7091" - 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,15 +41,8 @@ services: build: ../application # Consumes the default queue AND the dedicated `parsing` queue (read_document / # parse_document). Without `parsing` here the read_document await never resolves. - # For heavy/OCR parsing run a separate worker with `-Q parsing` (see - # deployment/sandbox/README.md). + # For heavy/OCR parsing run a separate worker with `-Q parsing`. 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: @@ -74,9 +52,6 @@ services: - API_URL=http://backend:7091 - CACHE_REDIS_URL=redis://redis:6379/2 - POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt - - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 - # Env-scrubbing kernelspec selected by name (see backend service). - - SANDBOX_KERNEL_NAME=docsgpt-python volumes: - ../application/indexes:/app/indexes - ../application/inputs:/app/inputs @@ -87,69 +62,19 @@ services: postgres: condition: service_healthy - # 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 - # boundary. The custom kernelspec scrubs secrets from the kernel env, but - # sibling workspaces are readable under the shared uid and kernels share one - # address space. Do NOT add `env_file: ../.env` here (the runner needs no app - # secrets). For cross-tenant / 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 - read_only: true - environment: - # Keep Jupyter's runtime/connection files on the writable tmpfs. - - JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime - - JUPYTER_DATA_DIR=/tmp/jupyter-data - tmpfs: - # Per-session workspaces (/tmp/docsgpt-sandbox/) and Jupyter - # runtime files live on tmpfs; the root FS is read-only everywhere else. - - /tmp - networks: - # 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 - - sandbox-egress - redis: image: redis:6-alpine - # Data plane only: NOT reachable from the code-exec sandbox (which has no - # auth on this broker, so a reachable Redis = Celery task injection -> RCE). - networks: - - data-net ports: - # 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 + - 6379:6379 postgres: image: postgres:16-alpine - # Data plane only (see redis): kept off the sandbox's networks. - networks: - - data-net environment: - POSTGRES_USER=docsgpt - POSTGRES_PASSWORD=docsgpt - POSTGRES_DB=docsgpt ports: - # Loopback only (see redis): not reachable from the sandbox via the gateway. - - "127.0.0.1:5432:5432" + - "5432:5432" volumes: - postgres_data:/var/lib/postgresql/data healthcheck: @@ -158,24 +83,5 @@ services: timeout: 5s retries: 10 -networks: - # 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. - data-net: - internal: true - volumes: postgres_data: - diff --git a/deployment/k8s/deployments/docsgpt-deploy.yaml b/deployment/k8s/deployments/docsgpt-deploy.yaml index 16159f53..345b1449 100644 --- a/deployment/k8s/deployments/docsgpt-deploy.yaml +++ b/deployment/k8s/deployments/docsgpt-deploy.yaml @@ -57,14 +57,6 @@ spec: value: "false" - name: AUTO_CREATE_DB value: "false" - # Reach the always-on code-execution runner over HTTP + WebSocket. - - name: SANDBOX_GATEWAY_URL - value: "http://docsgpt-sandbox:8888" - # Select the runner's env-scrubbing kernelspec by its distinct name so - # it is never shadowed by the stock "python3" spec. MUST be set here: - # the runner only ships the kernelspec; the app (this pod) chooses it. - - name: SANDBOX_KERNEL_NAME - value: "docsgpt-python" --- apiVersion: apps/v1 kind: Deployment @@ -121,14 +113,6 @@ spec: value: "false" - name: AUTO_CREATE_DB value: "false" - # Reach the always-on code-execution runner over HTTP + WebSocket. - - name: SANDBOX_GATEWAY_URL - value: "http://docsgpt-sandbox:8888" - # Select the runner's env-scrubbing kernelspec by its distinct name so - # it is never shadowed by the stock "python3" spec. MUST be set here: - # the runner only ships the kernelspec; the app (this pod) chooses it. - - name: SANDBOX_KERNEL_NAME - value: "docsgpt-python" --- apiVersion: apps/v1 kind: Deployment diff --git a/deployment/k8s/deployments/sandbox-deploy.yaml b/deployment/k8s/deployments/sandbox-deploy.yaml index 24f92e81..ca82601b 100644 --- a/deployment/k8s/deployments/sandbox-deploy.yaml +++ b/deployment/k8s/deployments/sandbox-deploy.yaml @@ -1,7 +1,25 @@ -# docsgpt-sandbox runner (Jupyter Kernel Gateway). Sessions are in-process -# kernels, never child containers; the Docker socket is NOT mounted. Network -# egress is constrained by the NetworkPolicy under -# deployment/k8s/network-policies/sandbox-egress-policy.yaml -- apply both. +# docsgpt-sandbox runner (Jupyter Kernel Gateway). OPT-IN: this manifest and its +# NetworkPolicy are NOT in kustomization.yaml; apply them only when enabling +# code execution. Sessions are in-process kernels, never child containers; the +# Docker socket is NOT mounted. Network egress is constrained by the NetworkPolicy +# under deployment/k8s/network-policies/sandbox-egress-policy.yaml -- apply both. +# +# REQUIRED before enabling: create the shared gateway-token Secret and add the +# runner env to the app pods (docsgpt-api / docsgpt-worker), which the default +# docsgpt-deploy.yaml intentionally omits (sandbox is off by default): +# kubectl create secret generic docsgpt-sandbox-gateway \ +# --from-literal=token="$(openssl rand -hex 32)" +# # on docsgpt-api and docsgpt-worker containers, add: +# # - name: SANDBOX_GATEWAY_URL +# # value: "http://docsgpt-sandbox:8888" +# # - name: SANDBOX_KERNEL_NAME +# # value: "docsgpt-python" +# # - name: SANDBOX_GATEWAY_AUTH_TOKEN +# # valueFrom: { secretKeyRef: { name: docsgpt-sandbox-gateway, key: token } } +# The gateway authenticates every HTTP + WebSocket request with the token, so +# kernel code cannot reach the gateway control API over loopback (enumerate/kill +# sibling kernels or bypass the session cap). The gateway fails closed if the +# token is unset. # # On Linux prod, schedule this onto a gVisor `runsc` RuntimeClass for kernel # isolation (uncomment `runtimeClassName` once the node has it installed). @@ -52,6 +70,13 @@ spec: drop: - ALL env: + # Shared gateway token; the launcher fails closed if it is unset. + # Must match SANDBOX_GATEWAY_AUTH_TOKEN on the app pods. + - name: SANDBOX_GATEWAY_AUTH_TOKEN + valueFrom: + secretKeyRef: + name: docsgpt-sandbox-gateway + key: token - name: JUPYTER_RUNTIME_DIR value: /tmp/jupyter-runtime - name: JUPYTER_DATA_DIR diff --git a/deployment/optional/docker-compose.optional.sandbox-egress.yaml b/deployment/optional/docker-compose.optional.sandbox-egress.yaml index eca41f23..b9e37c81 100644 --- a/deployment/optional/docker-compose.optional.sandbox-egress.yaml +++ b/deployment/optional/docker-compose.optional.sandbox-egress.yaml @@ -7,9 +7,10 @@ # 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 +# The sandbox overlay (docker-compose.optional.sandbox.yaml) puts docsgpt-sandbox +# on two networks: the internal `sandbox-net` (control plane to backend/worker) +# and `sandbox-egress` (its internet route). Layer THIS overlay on top of it to +# flip `sandbox-egress` to `internal: true`, which removes the # runner's direct route to the internet / host / RFC1918 / metadata entirely -- # so raw sockets in arbitrary sandbox code have no route OFF the host except the # deny-private proxy. The proxy is the ONLY container with an internet route @@ -50,7 +51,7 @@ # `{{.NetworkSettings.Networks.docsgpt-oss_sandbox-egress.IPAddress}}` form # never parsed. Both IPs are read on sandbox-net so source+dest match the # shared path: -# CF="-f deployment/docker-compose.yaml -f deployment/optional/docker-compose.optional.sandbox-egress.yaml --profile sandbox" +# CF="-f deployment/docker-compose.yaml -f deployment/optional/docker-compose.optional.sandbox.yaml -f deployment/optional/docker-compose.optional.sandbox-egress.yaml" # NET=docsgpt-oss_sandbox-net # SBX=$(docker inspect -f "{{(index .NetworkSettings.Networks \"$NET\").IPAddress}}" "$(docker compose $CF ps -q docsgpt-sandbox)") # BE=$( docker inspect -f "{{(index .NetworkSettings.Networks \"$NET\").IPAddress}}" "$(docker compose $CF ps -q backend)") @@ -74,11 +75,10 @@ # This does NOT stop runner->backend/worker on sandbox-net -- pair it with # backend auth and/or approach (1). # -# Apply alongside the base stack: +# Apply on top of the base stack AND the sandbox overlay (order matters): # docker compose -f deployment/docker-compose.yaml \ +# -f deployment/optional/docker-compose.optional.sandbox.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: because `sandbox-egress` is internal here, non-HTTP raw egress TO THE # INTERNET / HOST / RFC1918 / METADATA is blocked at L3 regardless of what the @@ -102,14 +102,10 @@ services: - sandbox-egress-out # the ONLY container with an internet route docsgpt-sandbox: - # 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. + # No `networks:` override here on purpose: the sandbox overlay 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 diff --git a/deployment/optional/docker-compose.optional.sandbox.yaml b/deployment/optional/docker-compose.optional.sandbox.yaml new file mode 100644 index 00000000..8b6f32bb --- /dev/null +++ b/deployment/optional/docker-compose.optional.sandbox.yaml @@ -0,0 +1,91 @@ +# Optional code-execution sandbox overlay. +# +# Sandboxed code execution (the `code_executor` / `artifact_generator` tools) is +# an OPT-IN feature and is NOT part of the default stack. This overlay adds the +# `docsgpt-sandbox` runner (a Jupyter Kernel Gateway; sessions are in-process +# kernels, never child containers, and the Docker socket is NOT mounted) and +# wires the app to it. +# +# Apply on top of the base stack (order matters -- base first): +# docker compose -f deployment/docker-compose.yaml \ +# -f deployment/optional/docker-compose.optional.sandbox.yaml up +# +# REQUIRED: set a shared gateway token before starting. The runner's gateway +# authenticates every HTTP + WebSocket request with it (so kernel code cannot +# reach the gateway control API over loopback -- enumerate/kill sibling kernels +# or bypass the session cap), and the app sends it as `Authorization: token`. +# Compose fails to start if it is unset. Put it in `deployment/.env` or export it: +# export SANDBOX_GATEWAY_AUTH_TOKEN=$(openssl rand -hex 32) +# +# SECURITY POSTURE (single trust domain): all sessions share this one +# container/uid and are isolated by working directory only -- sibling workspaces +# are readable and kernels share one address space. The kernelspec scrubs secrets +# from the kernel env, and the runner sits on the internal `sandbox-net` + +# `sandbox-egress` only (NOT the `default` net), so it has no Docker-DNS route to +# redis/postgres. Two residual gaps remain, exactly as documented in the egress +# overlay: (a) internet/host/RFC1918/metadata egress is open until you ALSO layer +# docker-compose.optional.sandbox-egress.yaml; (b) the runner shares sandbox-net +# with backend/worker (its control path), so on `AUTH_TYPE=none` it can reach the +# API unauthenticated -- run real auth and/or the egress overlay's host-firewall +# DROP. For untrusted multi-tenant workloads use the Daytona backend +# (SANDBOX_BACKEND=daytona, a per-session VM) or the Kubernetes manifests, whose +# NetworkPolicy blocks the internal path too. + +services: + backend: + environment: + # Code-execution runner reached over HTTP + WebSocket (no docker socket). + - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 + # Select the runner's env-scrubbing kernelspec (distinct name; never + # shadowed by the stock "python3" spec). + - SANDBOX_KERNEL_NAME=docsgpt-python + # Shared secret the app sends as `Authorization: token <...>`. Must match + # the runner's token below; compose errors if it is unset. + - SANDBOX_GATEWAY_AUTH_TOKEN=${SANDBOX_GATEWAY_AUTH_TOKEN:?set SANDBOX_GATEWAY_AUTH_TOKEN to a shared gateway token} + networks: + - default + - sandbox-net + + worker: + environment: + - SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888 + - SANDBOX_KERNEL_NAME=docsgpt-python + - SANDBOX_GATEWAY_AUTH_TOKEN=${SANDBOX_GATEWAY_AUTH_TOKEN:?set SANDBOX_GATEWAY_AUTH_TOKEN to a shared gateway token} + networks: + - default + - sandbox-net + + docsgpt-sandbox: + build: ./sandbox + mem_limit: ${SANDBOX_MEMORY:-1g} + cpus: ${SANDBOX_CPUS:-1.0} + pids_limit: 256 + read_only: true + environment: + # The gateway REQUIRES this and fails closed if unset (see gateway-launch.sh). + # Do NOT add `env_file: ../.env` here -- the runner needs no other app secret. + - SANDBOX_GATEWAY_AUTH_TOKEN=${SANDBOX_GATEWAY_AUTH_TOKEN:?set SANDBOX_GATEWAY_AUTH_TOKEN to a shared gateway token} + # Keep Jupyter's runtime/connection files on the writable tmpfs. + - JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime + - JUPYTER_DATA_DIR=/tmp/jupyter-data + tmpfs: + # Per-session workspaces (/tmp/docsgpt-sandbox/) and Jupyter + # runtime files live on tmpfs; the root FS is read-only everywhere else. + - /tmp + networks: + # Reachable by backend/worker over the internal sandbox-net; internet egress + # (runtime pip install, etc.) via the dedicated sandbox-egress net. NOT on + # `default`: the egress overlay cuts internet by flipping sandbox-egress to + # internal, without severing the control path. + - sandbox-net + - sandbox-egress + +networks: + # Control plane between backend/worker and the 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. Internet-facing by default so runtime + # pip install works; docker-compose.optional.sandbox-egress.yaml flips it to + # internal and forces egress through a deny-private proxy for SSRF containment. + sandbox-egress: {} diff --git a/deployment/sandbox/Dockerfile b/deployment/sandbox/Dockerfile index b1ac1816..c70bd547 100644 --- a/deployment/sandbox/Dockerfile +++ b/deployment/sandbox/Dockerfile @@ -56,6 +56,8 @@ RUN if [ "$INSTALL_DOCLING" = "true" ]; then \ # runner image itself should stay free of app secrets it has no use for. COPY kernel-launch.sh /opt/docsgpt/kernel-launch.sh RUN chmod 0555 /opt/docsgpt/kernel-launch.sh +COPY gateway-launch.sh /opt/docsgpt/gateway-launch.sh +RUN chmod 0555 /opt/docsgpt/gateway-launch.sh COPY kernels/docsgpt-python/kernel.json /usr/local/share/jupyter/kernels/docsgpt-python/kernel.json # Numeric UID (not the name) so a kubelet with `runAsNonRoot: true` can verify @@ -66,12 +68,8 @@ WORKDIR /home/sandbox EXPOSE 8888 -# ip=0.0.0.0 so the backend can reach it over the internal compose network. -# allow_origin is intentionally NOT set to "*" — publishing port 8888 requires -# setting SANDBOX_GATEWAY_AUTH_TOKEN (see README). limit_rate=False raises the -# iopub data-rate cap so large get_file base64 payloads are not truncated; the -# get_file integrity check still guards against truncation if this is ever off. -CMD ["jupyter", "kernelgateway", \ - "--KernelGatewayApp.ip=0.0.0.0", \ - "--KernelGatewayApp.port=8888", \ - "--ZMQChannelsWebsocketConnection.limit_rate=False"] +# The launcher REQUIRES SANDBOX_GATEWAY_AUTH_TOKEN and fails closed if it is +# unset: the gateway control API is reachable from kernel code over loopback, so +# it must never run unauthenticated (see gateway-launch.sh). The token is shared +# with the app and is scrubbed from the kernel env by kernel-launch.sh. +CMD ["/opt/docsgpt/gateway-launch.sh"] diff --git a/deployment/sandbox/README.md b/deployment/sandbox/README.md index e68a80e8..8de18717 100644 --- a/deployment/sandbox/README.md +++ b/deployment/sandbox/README.md @@ -7,28 +7,34 @@ 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: +The runner is **opt-in**. Neither `code_executor` nor `artifact_generator` is a +default chat tool (both were removed from `DEFAULT_CHAT_TOOLS`), and the runner +is **not** part of the base compose stack — a plain `docker compose up` does +**not** start `docsgpt-sandbox`. Enable it by layering the sandbox overlay and +setting a shared gateway token: ```bash -docker compose -f deployment/docker-compose.yaml --profile sandbox up +export SANDBOX_GATEWAY_AUTH_TOKEN=$(openssl rand -hex 32) +docker compose \ + -f deployment/docker-compose.yaml \ + -f deployment/optional/docker-compose.optional.sandbox.yaml up ``` -With the egress-firewall overlay (see *Network egress / SSRF* below): +The token is **required** — the gateway fails closed if it is unset (see +*Gateway authentication* below). Add the egress-firewall overlay for SSRF +containment (see *Network egress / SSRF*): ```bash docker compose \ -f deployment/docker-compose.yaml \ - -f deployment/optional/docker-compose.optional.sandbox-egress.yaml \ - --profile sandbox up + -f deployment/optional/docker-compose.optional.sandbox.yaml \ + -f deployment/optional/docker-compose.optional.sandbox-egress.yaml 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. +Then enable `code_executor` / `artifact_generator` **per-agent** in the agent +tool picker. Agents without them never call the runner, and the backend/worker +degrade gracefully when the runner is absent. The `-hub` and `-azure` compose +variants take the same sandbox overlay. ## Isolation model @@ -70,6 +76,10 @@ Residual gaps (treat all sessions in one runner as mutually trusting): - **In-memory / cross-kernel.** Kernels are child processes of one gateway under one uid; OS-level process isolation is the only boundary, and it is not a sandbox boundary against a determined escape. No gVisor in the base posture. + (The gateway's HTTP/WebSocket control API is reachable from kernel code over + loopback, but it is **authenticated** — see *Gateway authentication* — and the + token is scrubbed from the kernel env, so kernel code cannot drive it to + enumerate/kill sibling kernels or spawn kernels past the session cap.) - **Egress.** Outbound is broad by design (so code can `pip install` / call public APIs). Private/link-local/metadata ranges are blocked **only** by the network layer — the k8s NetworkPolicy or a host/cloud firewall (see *Network @@ -88,15 +98,19 @@ Build and run the runner on its own, then point the app at it: ```bash docker build -t docsgpt-sandbox deployment/sandbox -docker run --rm -p 8888:8888 docsgpt-sandbox +docker run --rm -p 8888:8888 -e SANDBOX_GATEWAY_AUTH_TOKEN=devtoken docsgpt-sandbox # in the app's .env: SANDBOX_GATEWAY_URL=http://localhost:8888 +# SANDBOX_GATEWAY_AUTH_TOKEN=devtoken ``` -Without Docker (matches the test harness) you can run the gateway directly from -a venv that has `jupyter-kernel-gateway` installed: +The token is required — the image's entrypoint refuses to start without it (see +*Gateway authentication*). Without Docker (matches the test harness) you can run +the gateway directly from a venv that has `jupyter-kernel-gateway` installed; set +a matching `--KernelGatewayApp.auth_token`: ```bash jupyter kernelgateway --KernelGatewayApp.ip=0.0.0.0 --KernelGatewayApp.port=8888 \ + --KernelGatewayApp.auth_token=devtoken \ --ZMQChannelsWebsocketConnection.limit_rate=False ``` @@ -115,26 +129,35 @@ copy `kernels/docsgpt-python/kernel.json` (pointing `argv` at a local copy of `kernel-launch.sh`) into a Jupyter data dir on the kernelspec search path and set `SANDBOX_KERNEL_NAME=docsgpt-python` before launching. -## Exposing the port requires auth +## Gateway authentication -The image does **not** set `--KernelGatewayApp.allow_origin=*`. If you publish -port 8888 (e.g. `docker run -p 8888:8888`), set `SANDBOX_GATEWAY_AUTH_TOKEN` -and launch the gateway with a matching `--KernelGatewayApp.auth_token` so the -runner is not an open arbitrary-code-execution endpoint. In compose the runner -stays on the internal-only network with no published port, so no token is -required there. +The gateway **requires** an auth token and **fails closed** if it is unset — the +image's entrypoint (`gateway-launch.sh`) refuses to start an unauthenticated +gateway. This matters even on an internal-only network: the gateway and every +session kernel share one container, so kernel code can reach the gateway's +control API over **loopback** (`http://localhost:8888`). Without auth, that +control API would let kernel code enumerate/attach/kill sibling sessions' +kernels and spawn kernels without bound (bypassing the app-side session cap). + +Set the same token on the runner and the app via `SANDBOX_GATEWAY_AUTH_TOKEN` +(the app sends it as `Authorization: token <...>`; the runner's gateway +validates it on every HTTP + WebSocket request). Kernel code cannot read it: the +kernelspec launcher scrubs it from the kernel env (see *Isolation model*), so it +is present for the gateway process only. The image also does **not** set +`--KernelGatewayApp.allow_origin=*`. ## In docker-compose -The `docsgpt-sandbox` service is defined in `deployment/docker-compose.yaml` on -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` -is set on the `docsgpt-api` and `docsgpt-worker` deployments in -`deployment/k8s/deployments/docsgpt-deploy.yaml`. +The `docsgpt-sandbox` service lives in the opt-in overlay +`deployment/optional/docker-compose.optional.sandbox.yaml` (layered on the base +stack; see *Enabling code execution (opt-in)*) on an internal-only network with +no published host port. The overlay puts the backend and worker on `sandbox-net` +to reach the runner at `http://docsgpt-sandbox:8888`, and sets +`SANDBOX_KERNEL_NAME=docsgpt-python` on them (the runner only ships the +kernelspec; the app chooses it) plus the shared `SANDBOX_GATEWAY_AUTH_TOKEN`. In +k8s these are added to the `docsgpt-api` and `docsgpt-worker` deployments when +enabling the opt-in `sandbox-deploy.yaml` (the default `docsgpt-deploy.yaml` +omits them); see that manifest's header for the exact env and the token Secret. ## Artifact rendering on Daytona (snapshot) @@ -221,7 +244,7 @@ The hardened container runs **without `NET_ADMIN`**, so it cannot self-apply ``` - **docker-compose** — compose cannot express L3 egress filtering natively. The - base stack reaches the runner over an `internal: true` control network + sandbox overlay 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 @@ -241,16 +264,20 @@ The hardened container runs **without `NET_ADMIN`**, so it cannot self-apply real authentication (`AUTH_TYPE` != none / a real auth provider) so a reachable API rejects unauthenticated requests — **required** — and/or add a host-firewall `DROP` for runner→backend/worker on `sandbox-net` (see approach (1) in the - overlay file's header comment). Note the broker/DB published ports are bound to - `127.0.0.1` so the runner cannot reach them via the host gateway either. + overlay file's header comment). The runner is not on the `default` network, so + it has no Docker-DNS route to redis/postgres; if the broker/DB publish host + ports on a cloud VM, also apply the egress overlay (its `internal` flip removes + the runner's route to the host gateway / RFC1918) or bind those ports to + `127.0.0.1`. ## Other hardening (deployment-level) The gVisor `runsc` runtime (kernel isolation for untrusted code), seccomp profile, read-only root FS, non-root, and cgroup CPU/mem/PID caps (wired from `SANDBOX_MEMORY` / `SANDBOX_CPUS`) are deployment-level concerns. The compose -service in `deployment/docker-compose.yaml` already sets `read_only`, -`mem_limit`, `cpus`, and `pids_limit`; the k8s `sandbox-deploy.yaml` sets the +service in `deployment/optional/docker-compose.optional.sandbox.yaml` already +sets `read_only`, `mem_limit`, `cpus`, and `pids_limit`; the k8s +`sandbox-deploy.yaml` sets the equivalent `securityContext` + resource limits and has a commented `runtimeClassName: gvisor` to enable on nodes with the `runsc` RuntimeClass installed. These complement — they do not replace — the network egress policy diff --git a/deployment/sandbox/gateway-launch.sh b/deployment/sandbox/gateway-launch.sh new file mode 100755 index 00000000..0e1ab601 --- /dev/null +++ b/deployment/sandbox/gateway-launch.sh @@ -0,0 +1,33 @@ +#!/bin/sh +# Entrypoint for the docsgpt-sandbox Jupyter Kernel Gateway. +# +# The gateway and every session kernel share this one container, so LLM-authored +# kernel code can reach the gateway's control API over loopback +# (http://localhost:8888). WITHOUT authentication that means kernel code could +# enumerate/attach/kill sibling sessions' kernels and spawn kernels without bound +# (bypassing the app-side session cap). We therefore REQUIRE an auth token and +# fail closed if it is missing -- an unauthenticated gateway must never start. +# +# The token is shared with the app via SANDBOX_GATEWAY_AUTH_TOKEN (the app sends +# it as `Authorization: token <...>`). Kernel code cannot read it: kernel-launch.sh +# re-execs ipykernel under a scrubbed `env -i` allowlist that excludes the token, +# so it is present for the gateway process only. +set -eu + +TOKEN="${SANDBOX_GATEWAY_AUTH_TOKEN:-}" +if [ -z "$TOKEN" ]; then + echo "docsgpt-sandbox: refusing to start an unauthenticated gateway." >&2 + echo "Set SANDBOX_GATEWAY_AUTH_TOKEN (same value on the app and this container)." >&2 + exit 1 +fi + +# ip=0.0.0.0 so the backend/worker can reach it over the internal sandbox network. +# auth_token gates every HTTP + WebSocket request, including loopback ones from +# kernel code. limit_rate=False raises the iopub data-rate cap so large get_file +# base64 payloads are not truncated (the get_file integrity check still guards +# truncation if this is ever off). +exec jupyter kernelgateway \ + --KernelGatewayApp.ip=0.0.0.0 \ + --KernelGatewayApp.port=8888 \ + --KernelGatewayApp.auth_token="$TOKEN" \ + --ZMQChannelsWebsocketConnection.limit_rate=False diff --git a/deployment/sandbox/kernel-launch.sh b/deployment/sandbox/kernel-launch.sh old mode 100644 new mode 100755 diff --git a/docs/content/Tools/artifacts-and-code-execution.mdx b/docs/content/Tools/artifacts-and-code-execution.mdx index 16adbd13..f8e5617b 100644 --- a/docs/content/Tools/artifacts-and-code-execution.mdx +++ b/docs/content/Tools/artifacts-and-code-execution.mdx @@ -13,7 +13,7 @@ DocsGPT can generate documents, run code, and read files for you during a chat o - **Code Executor** runs code in a sandboxed session and turns any files the code writes into downloadable artifacts. - **Read Document** parses an uploaded or produced file (PDF, Word, PowerPoint, and more) into text, markdown, or structured data. -The first two are enabled by default for new agents. Read Document is available inside workflows. +All three are opt-in: enable Artifact and Code Executor per agent in the tool picker (they need a running [sandbox runner](https://github.com/arc53/DocsGPT/tree/main/deployment/sandbox)), and Read Document is available inside workflows. A new agent starts with no tools enabled. ## Artifacts diff --git a/frontend/src/components/artifactViewUtils.test.ts b/frontend/src/components/artifactViewUtils.test.ts index 5f08d253..25ddcc07 100644 --- a/frontend/src/components/artifactViewUtils.test.ts +++ b/frontend/src/components/artifactViewUtils.test.ts @@ -1,6 +1,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { + ARTIFACT_URL_ENVELOPE_MIME, buildPreviewDocument, bytesPreviewModeForMime, displayFilename, @@ -237,14 +238,14 @@ describe('buildPreviewDocument', () => { }); describe('readPresignedUrlEnvelope', () => { - const json = (body: unknown, contentType = 'application/json') => + const envelope = (body: unknown, contentType = ARTIFACT_URL_ENVELOPE_MIME) => new Response(JSON.stringify(body), { headers: { 'Content-Type': contentType }, }); it('extracts the presigned url from the s3 JSON envelope', async () => { const url = await readPresignedUrlEnvelope( - json({ success: true, url: 'https://signed.example/x?sig=1' }), + envelope({ success: true, url: 'https://signed.example/x?sig=1' }), ); expect(url).toBe('https://signed.example/x?sig=1'); }); @@ -258,20 +259,36 @@ describe('readPresignedUrlEnvelope', () => { }); it('returns null for a JSON artifact whose bytes are not the envelope', async () => { - const res = json({ foo: 1 }); + const res = envelope({ foo: 1 }, 'application/json'); expect(await readPresignedUrlEnvelope(res)).toBeNull(); // clone() was used to peek, so the original body is still consumable. expect(await res.json()).toEqual({ foo: 1 }); }); + it('ignores an attacker JSON artifact that mimics the envelope body', async () => { + // Under URL_STRATEGY=backend a `data` artifact is streamed with its own + // application/json content-type; even with the exact envelope shape it must + // NOT be treated as a redirect (open-redirect gadget) — only the vendor + // media type marks a real envelope. + const res = envelope( + { success: true, url: 'https://attacker.example/phish' }, + 'application/json', + ); + expect(await readPresignedUrlEnvelope(res)).toBeNull(); + expect(await res.json()).toEqual({ + success: true, + url: 'https://attacker.example/phish', + }); + }); + it('rejects a non-http(s) url in the envelope', async () => { expect( await readPresignedUrlEnvelope( - json({ success: true, url: 'javascript:alert(1)' }), + envelope({ success: true, url: 'javascript:alert(1)' }), ), ).toBeNull(); expect( - await readPresignedUrlEnvelope(json({ success: true, url: '' })), + await readPresignedUrlEnvelope(envelope({ success: true, url: '' })), ).toBeNull(); }); }); @@ -308,7 +325,7 @@ describe('triggerResponseDownload', () => { it('navigates top-level to the presigned url for the s3 envelope', async () => { const res = new Response( JSON.stringify({ success: true, url: 'https://signed.example/x' }), - { headers: { 'Content-Type': 'application/json' } }, + { headers: { 'Content-Type': ARTIFACT_URL_ENVELOPE_MIME } }, ); expect(await triggerResponseDownload(res, 'f.bin')).toBe(true); // No blob path for the presigned envelope. diff --git a/frontend/src/components/artifactViewUtils.ts b/frontend/src/components/artifactViewUtils.ts index 75b7420b..cc986a74 100644 --- a/frontend/src/components/artifactViewUtils.ts +++ b/frontend/src/components/artifactViewUtils.ts @@ -259,19 +259,29 @@ export function filenameFromContentDisposition( /** * Detect the `URL_STRATEGY=s3` presigned-URL envelope and return its `url`, or - * null when the response is a byte stream (`URL_STRATEGY=backend`). The download - * endpoint returns `{ success: true, url }` as JSON only when asked via - * `?disposition=url`; a plain byte stream (even a JSON artifact's own bytes) - * carries the file's own content-type and never this exact shape. The response - * is `clone()`d before reading so the caller can still consume the original body - * as bytes when this is not the envelope. The `url` must be an absolute http(s) - * URL to further guard against a JSON artifact coincidentally matching. + * null when the response is a byte stream (`URL_STRATEGY=backend`). + * + * Detection keys off the server-set envelope media type, NOT the body shape: + * under the backend strategy the endpoint streams the artifact's own bytes, so a + * `data` artifact whose bytes are literally + * `{"success":true,"url":"https://attacker.example"}` would otherwise be read as + * a redirect target and navigated to (an open-redirect / phishing gadget). Only + * the real envelope carries `ARTIFACT_URL_ENVELOPE_MIME` (a stored artifact's own + * mime can never be this vendor value, and Content-Type is CORS-safelisted so it + * stays readable cross-origin), so attacker-controlled bytes can never be + * mistaken for one. The response is `clone()`d before reading so the caller can + * still consume the original body as bytes when this is not the envelope; the + * shape/absolute-http(s) checks remain as defense in depth. */ +export const ARTIFACT_URL_ENVELOPE_MIME = + 'application/vnd.docsgpt.artifact-url+json'; + export async function readPresignedUrlEnvelope( response: Response, ): Promise { const contentType = response.headers.get('Content-Type') ?? ''; - if (!contentType.toLowerCase().includes('application/json')) return null; + if (!contentType.toLowerCase().includes(ARTIFACT_URL_ENVELOPE_MIME)) + return null; const data = await response .clone() .json() diff --git a/tests/agents/test_default_tools.py b/tests/agents/test_default_tools.py index c61b937a..5c057f6e 100644 --- a/tests/agents/test_default_tools.py +++ b/tests/agents/test_default_tools.py @@ -175,21 +175,22 @@ 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. + # The real shipped DEFAULT_CHAT_TOOLS must pass startup validation. Neither + # sandbox-backed tool is shipped default-on (both need a running runner). usable = default_tools.validate_default_chat_tools() assert "code_executor" not in default_tools.settings.DEFAULT_CHAT_TOOLS - assert "artifact_generator" in usable + assert "artifact_generator" not in default_tools.settings.DEFAULT_CHAT_TOOLS + assert usable # the remaining defaults (memory/read_webpage/scheduler) validate - 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. + def test_sandbox_tools_not_shipped_defaults(self): + # code_executor and artifact_generator both render/execute through the + # sandbox runner, an opt-in service; a fresh deploy without a runner must + # not advertise tools that hard-fail on every call. Enable per-agent. assert "code_executor" not in default_tools.settings.DEFAULT_CHAT_TOOLS - assert "artifact_generator" in default_tools.settings.DEFAULT_CHAT_TOOLS + assert "artifact_generator" not 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 + assert "artifact_generator" not in names def test_tool_with_required_config_is_rejected(self, monkeypatch): # ``brave`` needs an API key. @@ -339,18 +340,20 @@ class TestResolveToolById: assert row["builtin"] is True assert row["default"] is True - @pytest.mark.parametrize("name", ["artifact_generator"]) - def test_sandbox_default_id_resolves_in_memory(self, name): - # 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). + @pytest.mark.parametrize("name", ["code_executor", "artifact_generator"]) + def test_sandbox_tool_id_resolves_in_memory(self, name): + # Both sandbox tools are opt-in per agent (removed from DEFAULT_CHAT_TOOLS) + # but stay registered as agent-selectable builtins, so their synthetic id + # still resolves to an in-memory row (loaded user-scoped at execute time) + # — an agent that enabled one never silently loses it. + assert name not in default_tools.settings.DEFAULT_CHAT_TOOLS + assert name in default_tools.BUILTIN_AGENT_TOOLS 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") assert row is not None assert row["name"] == name assert row["id"] == tool_id + assert row["builtin"] is True def test_read_document_builtin_id_resolves_workflow_only(self): # read_document is a workflow-only builtin: its synthetic id resolves diff --git a/tests/api/user/test_artifacts_routes.py b/tests/api/user/test_artifacts_routes.py index a796b0d2..0a15a6a9 100644 --- a/tests/api/user/test_artifacts_routes.py +++ b/tests/api/user/test_artifacts_routes.py @@ -609,6 +609,9 @@ class TestDownloadArtifact: assert resp.status_code == 200 assert resp.json["success"] is True assert resp.json["url"].startswith("https://signed.example/") + # The envelope carries a distinctive media type so the client keys off + # a server signal, not the JSON body shape (open-redirect gadget guard). + assert resp.mimetype == "application/vnd.docsgpt.artifact-url+json" def test_stranger_denied(self, _patch_db, flask_app, monkeypatch): from application.api.user.artifacts.routes import DownloadArtifact diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index 66f4fd02..809d5a04 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -276,7 +276,7 @@ class TestSetupPeriodicTasks: setup_periodic_tasks(sender) - assert sender.add_periodic_task.call_count == 12 + assert sender.add_periodic_task.call_count == 13 calls = sender.add_periodic_task.call_args_list @@ -311,6 +311,9 @@ class TestSetupPeriodicTasks: # sandbox session reaper (60s) assert calls[11][0][0] == timedelta(seconds=60) assert calls[11][1].get("name") == "reap-sandbox-sessions" + # stale workflow-run reaper (5m) + assert calls[12][0][0] == timedelta(seconds=300) + assert calls[12][1].get("name") == "reap-stale-workflow-runs" class TestMcpOauthTask: diff --git a/tests/parser/test_document_reader.py b/tests/parser/test_document_reader.py index 248a6bd8..dbf464c0 100644 --- a/tests/parser/test_document_reader.py +++ b/tests/parser/test_document_reader.py @@ -28,6 +28,44 @@ def test_unknown_extension_is_rejected(): assert "error" in out and "unsupported file type" in out["error"] +def _make_zip(entries: Dict[str, bytes]) -> bytes: + import io + import zipfile + + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: + for name, data in entries.items(): + zf.writestr(name, data) + return buf.getvalue() + + +@pytest.mark.unit +def test_zip_bomb_declared_size_over_cap_is_rejected(monkeypatch): + # A tiny highly-compressible archive whose decompressed size exceeds the cap + # must be rejected before any parser reads it (guards a zip-bomb OOM). + monkeypatch.setattr(dr.settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 1000, raising=False) + data = _make_zip({"word/document.xml": b"A" * 50_000}) + assert len(data) < 1000 # the on-disk archive is well under the byte cap + out = parse_document_bytes(data, "bomb.docx") + assert "error" in out and "too much data" in out["error"] + + +@pytest.mark.unit +def test_zip_too_many_entries_is_rejected(monkeypatch): + monkeypatch.setattr(dr.settings, "DOCUMENT_MAX_ARCHIVE_ENTRIES", 3, raising=False) + data = _make_zip({f"f{i}.xml": b"x" for i in range(10)}) + out = parse_document_bytes(data, "many.xlsx") + assert "error" in out and "too many entries" in out["error"] + + +@pytest.mark.unit +def test_reject_zip_bomb_ignores_non_zip_formats(): + # A non-zip extension (or a non-zip payload named .docx) is not gated here; + # the format parser surfaces its own error downstream. + assert dr._reject_zip_bomb(b"plain text", ".txt") is None + assert dr._reject_zip_bomb(b"not a zip", ".docx") is None + + @pytest.mark.unit def test_size_cap_rejects_oversize(monkeypatch): monkeypatch.setattr(dr.settings, "DOCUMENT_PARSE_MAX_BYTES", 8, raising=False) diff --git a/tests/storage/db/repositories/test_artifacts.py b/tests/storage/db/repositories/test_artifacts.py index 3673ca01..b37b980f 100644 --- a/tests/storage/db/repositories/test_artifacts.py +++ b/tests/storage/db/repositories/test_artifacts.py @@ -56,6 +56,23 @@ class TestCreateArtifact: assert v1["sha256"] == "abc" assert v1["spec"] == {"slides": []} + def test_persists_message_id_for_share_link_scope(self, pg_conn): + # Share-token artifact access gates on the artifact's message being within + # the conversation's shared first_n_queries snapshot, so creation paths + # must tag the artifact with the producing message. Without a stored + # message_id the share-link inheritance is dead (always denied). + repo = _repo(pg_conn) + msg_id = str(uuid.uuid4()) + artifact = repo.create_artifact( + "user-1", + "document", + conversation_id=_conversation_id(), + message_id=msg_id, + filename="d.md", + ) + assert str(artifact["message_id"]) == msg_id + assert str(repo.get_artifact(artifact["id"])["message_id"]) == msg_id + def test_spec_only_version_allows_null_storage_path(self, pg_conn): repo = _repo(pg_conn) artifact = repo.create_artifact( @@ -308,6 +325,37 @@ class TestQuotaAccounting: repo.create_artifact("alice", "document", conversation_id=_conversation_id()) assert repo.total_bytes_for_user("alice") == 0 + def test_total_bytes_dedupes_restored_version_by_storage_path(self, pg_conn): + # Restore appends a version that re-points at an existing version's stored + # object (same storage_path + size). It stores no new bytes, so it must not + # inflate the user's usage. + repo = _repo(pg_conn) + art = repo.create_artifact( + "alice", + "document", + conversation_id=_conversation_id(), + storage_path="inputs/alice/artifacts/x/v1/deck.pptx", + size=500, + filename="deck.pptx", + ) + # A genuine new version stores new bytes at a new key (+300 -> 800). + repo.append_version( + art["id"], + storage_path="inputs/alice/artifacts/x/v2/deck.pptx", + size=300, + filename="deck.pptx", + ) + assert repo.total_bytes_for_user("alice") == 800 + # Restoring v1 re-points at v1's object; usage stays 800, not 1300. + source = repo.get_version(art["id"], 1) + repo.append_version( + art["id"], + storage_path=source["storage_path"], + size=source["size"], + filename=source["filename"], + ) + assert repo.total_bytes_for_user("alice") == 800 + class TestQuotaEnforcement: """Exercises the shared enforcement helper against a real repo + connection.""" diff --git a/tests/storage/db/repositories/test_workflow_runs.py b/tests/storage/db/repositories/test_workflow_runs.py index 2bca3900..84f5c7e1 100644 --- a/tests/storage/db/repositories/test_workflow_runs.py +++ b/tests/storage/db/repositories/test_workflow_runs.py @@ -2,7 +2,7 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from application.storage.db.repositories.workflows import WorkflowsRepository @@ -98,3 +98,36 @@ class TestListForWorkflow: wf = _wf(pg_conn) repo = _repo(pg_conn) assert repo.list_for_workflow(wf["id"]) == [] + + +class TestMarkStaleRunningFailed: + def test_fails_old_running_run(self, pg_conn): + wf = _wf(pg_conn) + repo = _repo(pg_conn) + old = datetime.now(timezone.utc) - timedelta(hours=2) + run = repo.create(wf["id"], "user-1", "running", started_at=old) + + reaped = repo.mark_stale_running_failed( + datetime.now(timezone.utc) - timedelta(hours=1) + ) + assert reaped == 1 + fetched = repo.get(run["id"]) + assert fetched["status"] == "failed" + assert fetched["ended_at"] is not None + assert "did not complete" in (fetched["result"] or {}).get("error", "") + + def test_leaves_recent_running_and_terminal_runs(self, pg_conn): + wf = _wf(pg_conn) + repo = _repo(pg_conn) + now = datetime.now(timezone.utc) + recent = repo.create(wf["id"], "user-1", "running", started_at=now) + # A terminal run that started long ago must not be touched (has ended_at). + done = repo.create( + wf["id"], "user-1", "completed", + started_at=now - timedelta(hours=2), ended_at=now - timedelta(hours=2), + ) + + reaped = repo.mark_stale_running_failed(now - timedelta(hours=1)) + assert reaped == 0 + assert repo.get(recent["id"])["status"] == "running" + assert repo.get(done["id"])["status"] == "completed"