tiennm99/DocsGPT
1
0 Fork 0
mirror of https://github.com/tiennm99/DocsGPT.git synced 2026-10-03 07:11:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: remove code exec as default and sec improvements

Browse Source
This commit is contained in:
Alex committed 2026-07-08 18:50:10 +01:00
1 parent 2d79ff2227
commit 66786c2760
41 files changed
+800 -772
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.github/THREAT_MODEL.md
+3 -3
View File
@@ -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
application/agents/default_tools.py
+10 -5
View File
@@ -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
application/agents/tool_executor.py
+81 -152
View File
@@ -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",
}
application/agents/tools/artifact_generator.py
+5 -3
View File
@@ -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"
" '<!doctype html><html lang=\"en\"><head><meta charset=\"utf-8\">'\n"
" '<meta name=\"viewport\" content=\"width=device-width,initial-scale=1\">'\n"
' \'<!doctype html><html lang="en"><head><meta charset="utf-8">\'\n'
' \'<meta name="viewport" content="width=device-width,initial-scale=1">\'\n'
" '<title>' + esc(title or 'Report') + '</title>'\n"
" '<style>' + css + '</style></head><body>'\n"
" + ''.join(parts) + '</body></html>'\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),
)
application/agents/tools/code_executor.py
+10 -8
View File
@@ -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",
application/agents/tools/read_document.py
+7 -7
View File
@@ -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}
# ------------------------------------------------------------------
application/agents/workflow_agent.py
+34 -43
View File
@@ -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):
application/api/user/artifacts/routes.py
+62 -150
View File
@@ -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/<artifact_id>/versions/<int:version>")
@@ -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:
application/api/user/tasks.py
+34
View File
@@ -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
application/core/settings.py
+7 -8
View File
@@ -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.
application/parser/document_reader.py
+37 -8
View File
@@ -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)
application/prompts/agentic/creative.txt
+3 -1
View File
@@ -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 %}
application/prompts/agentic/default.txt
+3 -1
View File
@@ -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 %}
application/prompts/agentic/strict.txt
+3 -1
View File
@@ -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 %}
application/prompts/chat_combine_creative.txt
+3 -1
View File
@@ -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 %}
application/prompts/chat_combine_default.txt
+3 -1
View File
@@ -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 %}
application/prompts/chat_combine_strict.txt
+3 -1
View File
@@ -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 %}
application/sandbox/artifacts_capture.py
+6
View File
@@ -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,
application/storage/db/repositories/artifacts.py
+21 -5
View File
@@ -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()
application/storage/db/repositories/workflow_runs.py
+23 -1
View File
@@ -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)"),
application/worker.py
+1
View File
@@ -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:
deployment/docker-compose-azure.yaml
+8 -80
View File
@@ -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:
deployment/docker-compose-hub.yaml
+8 -83
View File
@@ -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:
deployment/docker-compose.yaml
+3 -97
View File
@@ -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/<session_id>) 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:
deployment/k8s/deployments/docsgpt-deploy.yaml
-16
View File
@@ -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
deployment/k8s/deployments/sandbox-deploy.yaml
+29 -4
View File
@@ -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
deployment/optional/docker-compose.optional.sandbox-egress.yaml
+11 -15
View File
@@ -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
deployment/optional/docker-compose.optional.sandbox.yaml
+91
View File
@@ -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/<session_id>) 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: {}
deployment/sandbox/Dockerfile
+7 -9
View File
@@ -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"]
deployment/sandbox/README.md
+63 -36
View File
@@ -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
deployment/sandbox/gateway-launch.sh
Executable
+33
View File
@@ -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
deployment/sandbox/kernel-launch.sh
Regular → Executable
View File
File mode changed.
docs/content/Tools/artifacts-and-code-execution.mdx
+1 -1
View File
@@ -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
frontend/src/components/artifactViewUtils.test.ts
+23 -6
View File
@@ -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.
frontend/src/components/artifactViewUtils.ts
+18 -8
View File
@@ -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<string | null> {
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()
tests/agents/test_default_tools.py
+19 -16
View File
@@ -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
tests/api/user/test_artifacts_routes.py
+3
View File
@@ -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
tests/api/user/test_tasks.py
+4 -1
View File
@@ -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:
tests/parser/test_document_reader.py
+38
View File
@@ -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)
tests/storage/db/repositories/test_artifacts.py
+48
View File
@@ -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."""
tests/storage/db/repositories/test_workflow_runs.py
+34 -1
View File
@@ -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"