Merge pull request #2568 from arc53/artifacts-and-code-executors

Artifacts and sandboxed code execution: tools, workflows, and document I/O
This commit is contained in:
Alex authored and GitHub committed 2026-07-10 15:33:28 +02:00
commit 3e978759de
159 files changed
+23593 -732

No files matched your search

+5 -1
View File
@@ -29,9 +29,13 @@ serves only the WSGI Flask app — it omits `/mcp` and the reconnect reader
### Celery (Task Queue)
```bash
celery -A application.app.celery worker -l INFO
celery -A application.app.celery worker -l INFO -Q docsgpt,parsing
```
The `parsing` queue serves document parsing (the `read_document` tool / workflow
native-file parse); without it those calls hang `DOCUMENT_PARSE_TIMEOUT` then
error. A dedicated `-Q parsing` worker can be GPU-enabled for heavier parsers.
## Github Codespaces Instructions
### 1. Make Ports Public:
+11 -1
View File
@@ -1,7 +1,7 @@
# DocsGPT Public Threat Model
**Classification:** Public
**Last updated:** 2026-04-15
**Last updated:** 2026-06-25
**Applies to:** Open-source and self-hosted DocsGPT deployments
## 1) Overview
@@ -41,6 +41,7 @@ Out of scope:
- Documents, attachments, chunks/embeddings, summaries.
- Conversations, agents, workflows, prompt templates.
- Generated artifacts and their versions; sandbox code-execution sessions.
- Secrets (JWT secret, `INTERNAL_KEY`, provider/API/OAuth credentials).
- Operational capacity (worker throughput, queue depth, model quota/cost).
@@ -65,6 +66,7 @@ Untrusted input includes API payloads, file uploads, remote URLs, OAuth/webhook
6. Frontend rendering + token storage.
7. Internal service endpoints (`INTERNAL_KEY`).
8. High-impact integrations (SQL tool, generic API tool, remote MCP tools).
9. Sandboxed code execution (LLM-authored code, document/artifact generation, workflow code nodes).
## 7) Key threats and expected mitigations
@@ -103,6 +105,12 @@ Untrusted input includes API payloads, file uploads, remote URLs, OAuth/webhook
- Threat: request floods, large ingestion jobs, expensive prompts/crawls.
- Mitigations: rate limits, quotas, timeouts, queue backpressure, usage budgets.
### 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. 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
- Internet-exposed deployment runs with weak/no auth and receives unauthorized data access/abuse.
@@ -111,6 +119,7 @@ Untrusted input includes API payloads, file uploads, remote URLs, OAuth/webhook
- Malicious URL/redirect chain targets internal services.
- Poisoned document causes data exfiltration through tool calls.
- Over-privileged SQL/API/MCP tool performs destructive side effects.
- A poisoned document drives a workflow code node or the code-execution tool to run attacker-chosen Python inside a shared runner and read another session's workspace.
## 9) Severity calibration
@@ -131,6 +140,7 @@ Untrusted input includes API payloads, file uploads, remote URLs, OAuth/webhook
8. Monitor auth failures, tool anomalies, ingestion spikes, and cost anomalies.
9. Keep dependencies/images patched and scanned.
10. Validate multi-tenant isolation with explicit tests.
11. Run untrusted code execution with per-tenant isolation (Daytona per-session VM or gVisor), scrubbed kernel secrets, and network-layer egress controls; treat a shared self-hosted runner as a single trust domain.
## 11) Maintenance
+6
View File
@@ -66,6 +66,12 @@ On macOS, prefer the solo pool for Celery:
python -m celery -A application.app.celery worker -l INFO --pool=solo
```
A bare worker (no `-Q`) consumes every configured queue, so one worker does the
whole job — app tasks and document parsing (the `read_document` tool / workflow
native-file parse) alike. Use `-Q` only to split load: run the main worker with
`-Q docsgpt` and a dedicated (e.g. GPU-enabled) parser worker with `-Q parsing`
for heavy OCR.
### Frontend
Install dependencies only when needed, then run the dev server:
+11 -7
View File
@@ -4,7 +4,11 @@ import uuid
from abc import ABC, abstractmethod
from typing import Any, Dict, Generator, List, Optional
from application.agents.tool_executor import ToolExecutor
from application.agents.tool_executor import (
ToolExecutor,
result_status,
truncate_tool_result,
)
from application.core.json_schema_utils import (
JsonSchemaValidationError,
normalize_json_schema_payload,
@@ -298,12 +302,8 @@ class BaseAgent(ABC):
"call_id": call_id,
"action_name": pending.get("llm_name", pending["name"]),
"arguments": args,
"result": (
result_str[:50] + "..."
if len(result_str) > 50
else result_str
),
"status": "completed",
"result": truncate_tool_result(result_str),
"status": result_status(result),
},
}
@@ -340,6 +340,10 @@ class BaseAgent(ABC):
self.tools = self.tool_executor.prepare_tools_for_llm(tools_dict)
def _execute_tool_action(self, tools_dict, call):
# Mirror the request's attachments onto the executor so sandbox tools
# can lazily bridge a referenced chat attachment to a conversation
# artifact; only the caller's own (user-scoped) attachments are passed.
self.tool_executor.attachments = self.attachments
return self.tool_executor.execute(
tools_dict, call, self.llm.__class__.__name__
)
+16 -1
View File
@@ -31,7 +31,21 @@ _HEADLESS_EXCLUDED_TOOLS = frozenset({"scheduler"})
# default tools. Names may overlap with DEFAULT_CHAT_TOOLS (e.g. ``scheduler``)
# — both registries share ``_DEFAULT_TOOL_NAMESPACE`` so the same uuid5
# resolves either way (the dual-flag row carries ``default`` AND ``builtin``).
BUILTIN_AGENT_TOOLS: tuple = ("scheduler",)
# ``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
# can filter; execution still reuses the builtin synthetic-id path.
WORKFLOW_ONLY_BUILTINS = frozenset({"read_document"})
_tool_cache: Dict[str, Optional[Any]] = {}
_ids_cache: Dict[tuple, Dict[str, str]] = {}
@@ -263,6 +277,7 @@ def synthesize_builtin_agent_tool(tool_name: str) -> Optional[Dict[str, Any]]:
"status": True,
"default": False,
"builtin": True,
"workflow_only": tool_name in WORKFLOW_ONLY_BUILTINS,
}
+1
View File
@@ -117,6 +117,7 @@ def run_agent_headless(
user_id=owner,
docs=retrieved_docs or None,
docs_together=format_docs_for_prompt(retrieved_docs),
artifact_parent={"conversation_id": conversation_id},
)
except Exception as exc:
logger.warning("Headless prompt rendering failed; using raw prompt: %s", exc)
+283 -154
View File
@@ -33,6 +33,70 @@ def _sanitize_tool_prefix(tool_name: Optional[str]) -> str:
return re.sub(r"[^a-zA-Z0-9_-]+", "_", str(tool_name or "")).strip("_")
# Longest string value rendered into a debug log line; longer values (e.g. an
# LLM-authored ``code`` body or an api_tool ``body``) are truncated so the full
# program/secret is never written to logs even at DEBUG level.
_LOG_VALUE_PREVIEW_LEN = 80
# Longest tool result persisted on the message / streamed to the UI. The LLM
# and the ``tool_call_attempts`` journal always receive the full result; this
# only bounds the message JSONB copy. 50 chars hid every real error behind
# "...", making retry storms undiagnosable from the stored conversation.
PERSISTED_RESULT_MAX_LEN = 2000
def truncate_tool_result(value: Any) -> Any:
"""Bound a tool result for persistence/streaming; short values pass through unchanged."""
text = value if isinstance(value, str) else str(value)
if len(text) <= PERSISTED_RESULT_MAX_LEN:
return value
return f"{text[:PERSISTED_RESULT_MAX_LEN]}..."
def result_status(result: Any) -> str:
"""Derive the persisted status from a tool's result payload.
Tools report failure in-band (``{"status": "error", ...}`` or an ``error``
key) while the executor used to stamp every returned result ``completed``,
so the stored conversation showed failed calls as successes.
"""
if isinstance(result, dict) and (result.get("status") == "error" or result.get("error")):
return "error"
return "completed"
def _redact_args_for_log(args: Any) -> Any:
"""Truncate long string values so a code/body argument never lands in logs in full."""
if not isinstance(args, dict):
text = str(args)
return text if len(text) <= _LOG_VALUE_PREVIEW_LEN else f"{text[:_LOG_VALUE_PREVIEW_LEN]}...(truncated)"
redacted: Dict[str, Any] = {}
for key, value in args.items():
if isinstance(value, str) and len(value) > _LOG_VALUE_PREVIEW_LEN:
redacted[key] = f"{value[:_LOG_VALUE_PREVIEW_LEN]}...(truncated, {len(value)} chars)"
elif isinstance(value, (dict, list)):
redacted[key] = f"<{type(value).__name__} omitted>"
else:
redacted[key] = value
return redacted
def _journal_key(call_id: str, message_id: Optional[str]) -> str:
"""Namespace the durability-journal key by the per-turn ``message_id``.
``tool_call_attempts.call_id`` is a table-wide primary key, but providers
reuse deterministic ids (e.g. ``functions.create_artifact:0``) across turns
and users, so distinct calls collide on that PK and the later journal rows
are silently dropped (``ON CONFLICT DO NOTHING``). Scoping the key by
``message_id`` (unique per turn) gives each logical call its own row while a
genuine retry of the same call within the same turn still dedupes. The raw
``call_id`` is left untouched for LLM tool-call/tool-result pairing and the
UI. Headless attempts with no ``message_id`` keep the raw key (unchanged
pre-existing behaviour).
"""
return f"{message_id}:{call_id}" if message_id else call_id
def _record_proposed(
call_id: str,
tool_name: str,
@@ -56,18 +120,14 @@ def _record_proposed(
try:
with db_session() as conn:
inserted = ToolCallAttemptsRepository(conn).record_proposed(
call_id,
_journal_key(call_id, message_id),
tool_name,
action_name,
arguments,
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(
@@ -101,12 +161,13 @@ def _mark_executed(
effect would be invisible to the journal. Both paths are scoped to
the owning ``user_id`` so a reused ``call_id`` can't cross tenants.
"""
key = _journal_key(call_id, message_id)
try:
with db_session() as conn:
repo = ToolCallAttemptsRepository(conn)
if proposed_ok:
updated = repo.mark_executed(
call_id,
key,
result,
message_id=message_id,
artifact_id=artifact_id,
@@ -116,7 +177,7 @@ def _mark_executed(
return
# Fallback synthesizes the row so the journal isn't lost.
repo.upsert_executed(
call_id,
key,
tool_name=tool_name or "unknown",
action_name=action_name or "",
arguments=arguments if arguments is not None else {},
@@ -125,21 +186,23 @@ 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)
def _mark_failed(call_id: str, error: str, *, user_id: Optional[str] = None) -> None:
def _mark_failed(
call_id: str,
error: str,
*,
message_id: Optional[str] = None,
user_id: Optional[str] = None,
) -> None:
try:
with db_session() as conn:
ToolCallAttemptsRepository(conn).mark_failed(
call_id, error, user_id=user_id
_journal_key(call_id, message_id), error, user_id=user_id
)
except Exception:
logger.exception("tool_call_attempts failed-write failed for %s", call_id)
@@ -169,13 +232,23 @@ 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] = {}
# Explicit tool-id scope (workflow agent nodes): when set (even empty),
# get_tools() resolves EXACTLY these ids — builtin synthetic ids and
# user_tools rows alike — with no defaults mixed in. None = unscoped.
self.allowed_tool_ids: Optional[List[str]] = None
self.conversation_id: Optional[str] = None
# Set by the workflow engine for agent nodes so run-scoped tools
# (artifact_generator / code_executor) address artifacts by the
# workflow run rather than a conversation.
self.workflow_run_id: Optional[str] = None
self.message_id: Optional[str] = None
# The request's own (already user-scoped) chat attachments, stamped onto
# sandbox tools so a referenced attachment can be lazily bridged to a
# conversation-scoped artifact at tool-use time.
self.attachments: List[Dict] = []
self.client_tools: Optional[List[Dict]] = None
self._name_to_tool: Dict[str, Tuple[str, str]] = {}
self._tool_to_name: Dict[Tuple[str, str], str] = {}
@@ -188,7 +261,9 @@ class ToolExecutor:
If *client_tools* have been set on this executor, they are
automatically merged into the returned dict.
"""
if self.user_api_key:
if self.allowed_tool_ids is not None:
tools = self._get_tools_by_ids(self.allowed_tool_ids)
elif self.user_api_key:
tools = self._get_tools_by_api_key(self.user_api_key)
else:
tools = self._get_user_tools(self.user or "local")
@@ -196,6 +271,39 @@ class ToolExecutor:
self.merge_client_tools(tools, self.client_tools)
return tools
def get_enabled_tool_names(self) -> set:
"""Return the set of tool names enabled for this context.
Authoritative (resolves through :meth:`get_tools`): an agent yields its
configured ``agents.tools``; an agentless chat yields the user's active
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")}
def _get_tools_by_ids(self, tool_ids: List[str]) -> Dict[str, Dict]:
"""Resolve an explicit tool-id scope — exactly these ids, no defaults.
Used by workflow agent nodes: the node's configured tools (builtin
synthetic ids like Artifact/Code Executor/Read Document, or the user's
``user_tools`` rows) are the node's WHOLE toolset. An unresolvable id
is dropped with a warning rather than failing the node.
"""
if not tool_ids:
return {}
with db_readonly() as conn:
tools_repo = UserToolsRepository(conn)
tools: List[Dict] = []
for tid in tool_ids:
row = resolve_tool_by_id(tid, self.user, user_tools_repo=tools_repo)
if row is None:
logger.warning("tool id %s did not resolve; dropped from scoped toolset", tid)
continue
if self.headless and is_headless_excluded_tool(row.get("name")):
continue
tools.append(row)
return {str(tool["id"]): tool for tool in tools}
def _get_tools_by_api_key(self, api_key: str) -> Dict[str, Dict]:
"""Resolve an agent's toolset — exactly ``agents.tools``, no defaults."""
# Per-operation session: the answer pipeline spans a long-lived
@@ -205,16 +313,16 @@ 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)
if row is None:
continue
# Workflow-only builtins (read_document) never resolve for a
# chat/scheduled agent — nodes get them via the scoped-id path.
if row.get("workflow_only"):
continue
# Headless runs (scheduled / webhook) drop chat-only tools
# like ``scheduler`` so a fire-time LLM can't chain schedules.
if self.headless and is_headless_excluded_tool(row.get("name")):
@@ -226,29 +334,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::
@@ -314,18 +416,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
@@ -335,34 +431,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
@@ -377,14 +458,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:
@@ -394,17 +477,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
@@ -434,9 +513,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),
}
@@ -454,9 +531,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),
@@ -473,10 +548,21 @@ 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
)
if require_approval:
@@ -498,10 +584,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),
}
@@ -527,7 +610,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.
@@ -547,11 +633,37 @@ 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,
) -> bool:
"""Live approval decision for a ``code_executor`` invocation.
Honors the deployment-level ``config.require_approval`` even when the
cached action snapshot is stale. Fails closed (require approval) on any
error so a misconfigured tool never silently runs untrusted code.
"""
try:
from application.agents.tools.code_executor import CodeExecutorTool
tool = CodeExecutorTool(
tool_config=tool_data.get("config") or {},
user_id=self.user,
)
requires_approval, _forced = tool.preview_decision(action_name, arguments)
return requires_approval
except Exception:
logger.exception(
"code_executor preview_decision failed; defaulting to a prompt",
)
return True
def execute(self, tools_dict: Dict, call, llm_class_name: str):
"""Execute a tool call. Yields status events, returns (result, call_id)."""
parser = ToolActionParser(llm_class_name, name_mapping=self._name_to_tool)
@@ -577,6 +689,7 @@ class ToolExecutor:
"action_name": llm_name,
"arguments": call_args or {},
"result": f"Failed to parse tool call. Invalid tool name format: {llm_name}",
"status": "error",
}
# Journal the malformed call so it still shows up in tool analytics.
if _record_proposed(
@@ -589,7 +702,10 @@ class ToolExecutor:
agent_id=self.agent_id,
):
_mark_failed(
call_id, tool_call_data["result"], user_id=self.user
call_id,
tool_call_data["result"],
message_id=self.message_id,
user_id=self.user,
)
yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}}
self.tool_calls.append(tool_call_data)
@@ -613,6 +729,7 @@ class ToolExecutor:
"action_name": llm_name,
"arguments": call_args,
"result": f"Tool with ID {tool_id} not found. Available tools: {list(tools_dict.keys())}",
"status": "error",
}
# Journal the unresolvable call so it still shows up in tool analytics.
if _record_proposed(
@@ -627,6 +744,7 @@ class ToolExecutor:
_mark_failed(
call_id,
f"Tool with ID {tool_id} not found.",
message_id=self.message_id,
user_id=self.user,
)
yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}}
@@ -664,14 +782,14 @@ 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"] = {}
tool_call_data["status"] = "error"
if proposed_ok:
_mark_failed(call_id, error_message, user_id=self.user)
_mark_failed(
call_id, error_message, message_id=self.message_id, user_id=self.user
)
yield {
"type": "tool_call",
"data": {**tool_call_data, "status": "error"},
@@ -682,11 +800,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 = {}, {}, {}, {}
@@ -700,29 +814,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",
@@ -734,9 +844,12 @@ class ToolExecutor:
},
)
tool_call_data["result"] = error_message
tool_call_data["status"] = "error"
if proposed_ok:
_mark_failed(call_id, error_message, user_id=self.user)
yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}}
_mark_failed(
call_id, error_message, message_id=self.message_id, user_id=self.user
)
yield {"type": "tool_call", "data": {**tool_call_data}}
self.tool_calls.append(tool_call_data)
return error_message, call_id
@@ -748,22 +861,28 @@ class ToolExecutor:
try:
if tool_data["name"] == "api_tool":
logger.debug(
f"Executing api: {action_name} with query_params: {query_params}, headers: {headers}, body: {body}"
"Executing api: %s with query_params: %s, headers: %s, body: %s",
action_name,
_redact_args_for_log(query_params),
_redact_args_for_log(headers),
_redact_args_for_log(body),
)
result = tool.execute_action(action_name, **body)
else:
logger.debug(f"Executing tool: {action_name} with args: {call_args}")
logger.debug(
"Executing tool: %s with args: %s",
action_name,
_redact_args_for_log(call_args),
)
result = tool.execute_action(action_name, **parameters)
except Exception as exc:
if proposed_ok:
_mark_failed(call_id, str(exc), user_id=self.user)
_mark_failed(
call_id, str(exc), message_id=self.message_id, 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):
@@ -782,9 +901,10 @@ 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"] = truncate_tool_result(result_full)
# A tool that ran but reported failure in-band persists as ``error``,
# not ``completed`` -- the model saw an error and will likely retry.
tool_call_data["status"] = result_status(result)
# Tool side effect has run; flip the journal row so the
# message-finalize path can later confirm it. If the proposed
@@ -805,23 +925,33 @@ 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"}}
yield {"type": "tool_call", "data": {**stream_tool_call_data}}
self.tool_calls.append(tool_call_data)
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 ''}"
if cache_key in self._loaded_tools:
return self._loaded_tools[cache_key]
cached = self._loaded_tools[cache_key]
# A tool cached on an earlier turn carries that turn's attachments;
# refresh them so a chat attachment added this turn is bridgeable.
cached_config = getattr(cached, "config", None)
if isinstance(cached_config, dict) and self.conversation_id:
# Refresh unconditionally so a turn with no attachments clears the
# prior turn's list (no stale carryover within the session).
cached_config["attachments"] = self.attachments or []
return cached
tm = ToolManager(config={})
@@ -834,12 +964,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
@@ -865,9 +991,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)
@@ -885,13 +1009,22 @@ 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).
if self.attachments:
tool_config["attachments"] = self.attachments
# Workflow agent nodes run-scope their artifact tools so a short
# ref (A1) and edit_artifact resolve against the workflow run.
if self.workflow_run_id:
tool_config["workflow_run_id"] = self.workflow_run_id
if tool_data["name"] == "scheduler":
# 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
@@ -915,12 +1048,8 @@ class ToolExecutor:
"action_name": tool_call.get("action_name"),
"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"]
),
"status": "completed",
"result": truncate_tool_result(tool_call.get("result")),
"status": tool_call.get("status", "completed"),
}
for tool_call in self.tool_calls
]
@@ -0,0 +1,796 @@
"""Artifact Generator tool: render editable documents from a JSON spec and version them append-only.
The ``artifact_versions.spec`` JSONB is the source of truth; the rendered
``.pptx``/``.docx``/``.xlsx``/``.pdf``/``.html`` is derived. ``create_artifact`` stores
v1, ``edit_artifact`` applies an RFC 7386 merge-patch to the current spec and
appends a version, ``rewrite_artifact`` replaces the spec wholesale and appends
a version. Rendering runs a FIXED program in the sandbox that reads the spec as
DATA (``json.loads``) — spec values are never interpolated into the program, so
a spec string containing code/quotes is rendered as literal text, not executed.
"""
from __future__ import annotations
import copy
import json
import logging
import uuid
from typing import Any, Dict, List, Optional
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.agents.tools.base import Tool
from application.core.settings import settings
from application.sandbox.artifacts_capture import (
QuotaExceeded,
append_artifact_version,
persist_new_artifact,
)
from application.sandbox.sandbox_creator import SandboxCreator
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
logger = logging.getLogger(__name__)
try:
import jsonschema
except Exception: # pragma: no cover - jsonschema is a declared dependency
jsonschema = None # type: ignore[assignment]
# Per-kind output metadata: artifact ``kind`` + produced file extension + mime.
_KIND_INFO: Dict[str, Dict[str, str]] = {
"presentation": {
"kind": "presentation",
"ext": "pptx",
"mime": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
},
"document": {
"kind": "document",
"ext": "docx",
"mime": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
},
"spreadsheet": {
"kind": "spreadsheet",
"ext": "xlsx",
"mime": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
},
"pdf": {
"kind": "document",
"ext": "pdf",
"mime": "application/pdf",
},
"html": {
"kind": "html",
"ext": "html",
"mime": "text/html",
},
}
# Tight per-kind JSON schemas. ``additionalProperties: false`` keeps specs minimal
# and rejects stray keys before any rendering happens.
_SCHEMAS: Dict[str, Dict[str, Any]] = {
"presentation": {
"type": "object",
"additionalProperties": False,
"properties": {
"title": {"type": "string"},
"slides": {
"type": "array",
"items": {
"type": "object",
"additionalProperties": False,
"properties": {
"title": {"type": "string"},
"bullets": {"type": "array", "items": {"type": "string"}},
"notes": {"type": "string"},
},
},
},
},
"required": ["slides"],
},
"document": {
"type": "object",
"additionalProperties": False,
"properties": {
"title": {"type": "string"},
"sections": {
"type": "array",
"items": {
"type": "object",
"additionalProperties": False,
"properties": {
"heading": {"type": "string"},
"paragraphs": {"type": "array", "items": {"type": "string"}},
},
},
},
},
"required": ["sections"],
},
"spreadsheet": {
"type": "object",
"additionalProperties": False,
"properties": {
"sheets": {
"type": "array",
"items": {
"type": "object",
"additionalProperties": False,
"properties": {
"name": {"type": "string"},
"rows": {
"type": "array",
"items": {"type": "array", "items": {}},
},
},
"required": ["rows"],
},
},
},
"required": ["sheets"],
},
"pdf": {
"type": "object",
"additionalProperties": False,
"properties": {
"title": {"type": "string"},
"blocks": {
"type": "array",
"items": {
"type": "object",
"additionalProperties": False,
"properties": {
"type": {"type": "string", "enum": ["heading", "paragraph"]},
"text": {"type": "string"},
},
"required": ["type", "text"],
},
},
},
"required": ["blocks"],
},
"html": {
"type": "object",
"additionalProperties": False,
"properties": {
"title": {"type": "string"},
"blocks": {
"type": "array",
"items": {
"oneOf": [
{
"type": "object",
"additionalProperties": False,
"properties": {
"type": {"const": "heading"},
"text": {"type": "string"},
"level": {"type": "integer", "minimum": 1, "maximum": 3},
},
"required": ["type", "text"],
},
{
"type": "object",
"additionalProperties": False,
"properties": {
"type": {"const": "paragraph"},
"text": {"type": "string"},
},
"required": ["type", "text"],
},
{
"type": "object",
"additionalProperties": False,
"properties": {
"type": {"const": "list"},
"ordered": {"type": "boolean"},
"items": {"type": "array", "items": {"type": "string"}},
},
"required": ["type", "items"],
},
{
"type": "object",
"additionalProperties": False,
"properties": {
"type": {"const": "table"},
"headers": {"type": "array", "items": {"type": "string"}},
"rows": {
"type": "array",
"items": {"type": "array", "items": {"type": "string"}},
},
},
"required": ["type", "rows"],
},
{
"type": "object",
"additionalProperties": False,
"properties": {
"type": {"const": "code"},
"text": {"type": "string"},
},
"required": ["type", "text"],
},
],
},
},
},
"required": ["blocks"],
},
}
# Compact per-kind spec shapes surfaced in the tool metadata. Without this the
# model has only the validation errors to reverse-engineer the schema from and
# brute-forces spec shapes call after call. Keep in sync with ``_SCHEMAS``
# (guarded by a test that checks every schema key appears here).
_SPEC_SYNOPSIS = (
"Exact spec shape per kind (no other keys are accepted): "
'presentation {"title"?, "slides": [{"title"?, "bullets"?: [str], "notes"?}]} · '
'document {"title"?, "sections": [{"heading"?, "paragraphs"?: [str]}]} · '
'spreadsheet {"sheets": [{"name"?, "rows": [[cell, ...]]}]} · '
'pdf {"title"?, "blocks": [{"type": "heading"|"paragraph", "text"}]} · '
'html {"title"?, "blocks": [...]} where each block is '
'{"type": "heading", "text", "level"?: 1-3} | {"type": "paragraph", "text"} | '
'{"type": "list", "items": [str], "ordered"?: bool} | '
'{"type": "table", "rows": [[str]], "headers"?: [str]} | {"type": "code", "text"}'
)
# FIXED renderer programs. Each reads ``spec.json`` from the workspace as DATA
# and writes ``out.<ext>``. The spec is NEVER string-interpolated into the
# program; ``{spec_path}``/``{out_path}`` are server-controlled path literals.
_RENDERERS: Dict[str, str] = {
"presentation": (
"import json\n"
"from pptx import Presentation\n"
"from pptx.util import Pt\n"
"spec = json.load(open({spec_path!r}))\n"
"prs = Presentation()\n"
"blank = prs.slide_layouts[6]\n"
"title_only = prs.slide_layouts[5]\n"
"for s in spec.get('slides', []):\n"
" slide = prs.slides.add_slide(title_only)\n"
" slide.shapes.title.text = str(s.get('title', '') or '')\n"
" bullets = s.get('bullets') or []\n"
" if bullets:\n"
" left = top = Pt(72)\n"
" width = prs.slide_width - Pt(144)\n"
" height = prs.slide_height - Pt(216)\n"
" box = slide.shapes.add_textbox(left, Pt(150), width, height)\n"
" tf = box.text_frame\n"
" tf.word_wrap = True\n"
" for i, b in enumerate(bullets):\n"
" para = tf.paragraphs[0] if i == 0 else tf.add_paragraph()\n"
" para.text = str(b)\n"
" notes = s.get('notes')\n"
" if notes:\n"
" slide.notes_slide.notes_text_frame.text = str(notes)\n"
"prs.save({out_path!r})\n"
),
"document": (
"import json\n"
"from docx import Document\n"
"spec = json.load(open({spec_path!r}))\n"
"doc = Document()\n"
"title = spec.get('title')\n"
"if title:\n"
" doc.add_heading(str(title), level=0)\n"
"for sec in spec.get('sections', []):\n"
" heading = sec.get('heading')\n"
" if heading:\n"
" doc.add_heading(str(heading), level=1)\n"
" for p in (sec.get('paragraphs') or []):\n"
" doc.add_paragraph(str(p))\n"
"doc.save({out_path!r})\n"
),
"spreadsheet": (
"import json\n"
"from openpyxl import Workbook\n"
"spec = json.load(open({spec_path!r}))\n"
# Formula-injection guard: spec content is model / prompt-injection
# controlled, so neutralize string cells openpyxl would treat as a live
# formula (leading =,+,-,@ or a control char) by quote-prefixing them.
"def _safe_cell(c):\n"
" 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"
"wb = Workbook()\n"
"wb.remove(wb.active)\n"
"for idx, sheet in enumerate(spec.get('sheets', [])):\n"
" name = str(sheet.get('name') or ('Sheet%d' % (idx + 1)))[:31]\n"
" ws = wb.create_sheet(title=name)\n"
" for row in (sheet.get('rows') or []):\n"
" ws.append([_safe_cell(c) for c in row])\n"
"if not wb.sheetnames:\n"
" wb.create_sheet(title='Sheet1')\n"
"wb.save({out_path!r})\n"
),
"pdf": (
"import json\n"
"from reportlab.lib.pagesizes import letter\n"
"from reportlab.lib.styles import getSampleStyleSheet\n"
"from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer\n"
"from xml.sax.saxutils import escape\n"
"spec = json.load(open({spec_path!r}))\n"
"styles = getSampleStyleSheet()\n"
"story = []\n"
"title = spec.get('title')\n"
"if title:\n"
" story.append(Paragraph(escape(str(title)), styles['Title']))\n"
" story.append(Spacer(1, 12))\n"
"for block in spec.get('blocks', []):\n"
" style = styles['Heading1'] if block.get('type') == 'heading' else styles['BodyText']\n"
" story.append(Paragraph(escape(str(block.get('text', ''))), style))\n"
" story.append(Spacer(1, 6))\n"
"SimpleDocTemplate({out_path!r}, pagesize=letter).build(story)\n"
),
"html": (
"import json\n"
"import html\n"
"spec = json.load(open({spec_path!r}))\n"
"def esc(value):\n"
" return html.escape('' if value is None else str(value))\n"
"parts = []\n"
"title = spec.get('title')\n"
"if title:\n"
" parts.append('<h1>' + esc(title) + '</h1>')\n"
"for block in spec.get('blocks', []):\n"
" kind = block.get('type')\n"
" if kind == 'heading':\n"
" level = block.get('level') or 2\n"
" try:\n"
" level = int(level)\n"
" except (TypeError, ValueError):\n"
" level = 2\n"
" level = min(max(level, 1), 3) + 1\n"
" parts.append('<h%d>%s</h%d>' % (level, esc(block.get('text', '')), level))\n"
" elif kind == 'paragraph':\n"
" parts.append('<p>' + esc(block.get('text', '')) + '</p>')\n"
" elif kind == 'list':\n"
" tag = 'ol' if block.get('ordered') else 'ul'\n"
" items = ''.join('<li>' + esc(i) + '</li>' for i in (block.get('items') or []))\n"
" parts.append('<%s>%s</%s>' % (tag, items, tag))\n"
" elif kind == 'table':\n"
" rows_html = []\n"
" headers = block.get('headers')\n"
" if headers:\n"
" cells = ''.join('<th>' + esc(h) + '</th>' for h in headers)\n"
" rows_html.append('<thead><tr>' + cells + '</tr></thead>')\n"
" body = []\n"
" for row in (block.get('rows') or []):\n"
" cells = ''.join('<td>' + esc(c) + '</td>' for c in row)\n"
" body.append('<tr>' + cells + '</tr>')\n"
" rows_html.append('<tbody>' + ''.join(body) + '</tbody>')\n"
" parts.append('<table>' + ''.join(rows_html) + '</table>')\n"
" elif kind == 'code':\n"
" parts.append('<pre><code>' + esc(block.get('text', '')) + '</code></pre>')\n"
# CSS braces are doubled so the outer ``_RENDERERS[kind].format(...)`` leaves them literal.
"css = (\n"
" 'body{{font-family:system-ui,-apple-system,Segoe UI,Roboto,sans-serif;'\n"
" 'line-height:1.6;color:#1a1a1a;max-width:800px;margin:0 auto;padding:24px}}'\n"
" 'h1,h2,h3,h4{{line-height:1.25;margin:1.2em 0 0.5em}}'\n"
" 'table{{border-collapse:collapse;width:100%;margin:1em 0}}'\n"
" 'th,td{{border:1px solid #d0d0d0;padding:6px 10px;text-align:left}}'\n"
" 'th{{background:#f5f5f5}}'\n"
" 'pre{{background:#f5f5f5;padding:12px;border-radius:6px;overflow:auto}}'\n"
" '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'
" '<title>' + esc(title or 'Report') + '</title>'\n"
" '<style>' + css + '</style></head><body>'\n"
" + ''.join(parts) + '</body></html>'\n"
")\n"
"open({out_path!r}, 'w', encoding='utf-8').write(doc)\n"
),
}
def merge_patch(target: Any, patch: Any) -> Any:
"""Apply an RFC 7386 JSON Merge Patch to ``target`` and return the result."""
if not isinstance(patch, dict):
return copy.deepcopy(patch)
if not isinstance(target, dict):
target = {}
result = copy.deepcopy(target)
for key, value in patch.items():
if value is None:
result.pop(key, None)
else:
result[key] = merge_patch(result.get(key), value)
return result
def _apply_spec_append(spec: Dict[str, Any], spec_append: Dict[str, Any]) -> Dict[str, Any]:
"""Concatenate items onto the spec's top-level lists (blocks/sections/slides/sheets).
RFC 7386 merge-patch replaces arrays wholesale, so "add a section" via
spec_patch silently wipes the existing ones unless the model resends the
full array. spec_append is the safe additive path: each key must name a
list (absent counts as empty) and its items are appended in order.
Returns {"spec": merged} or {"error": message}.
"""
result = copy.deepcopy(spec)
for key, items in spec_append.items():
if not isinstance(items, list):
return {"error": f"spec_append[{key!r}] must be a list of items to append."}
current = result.get(key, [])
if not isinstance(current, list):
return {"error": f"spec_append target {key!r} is not a list in the current spec."}
result[key] = current + copy.deepcopy(items)
return {"spec": result}
class ArtifactGeneratorTool(Tool):
"""Artifact
Create, edit, and version documents - slides, docs, sheets, PDF, HTML.
"""
def __init__(self, tool_config: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None) -> None:
"""Bind the tool to the invoker and its conversation/run-scoped sandbox session."""
self.config: Dict[str, Any] = tool_config or {}
self.user_id: Optional[str] = user_id
self.tool_id: Optional[str] = self.config.get("tool_id")
self.conversation_id: Optional[str] = self.config.get("conversation_id")
self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id")
self.message_id: Optional[str] = self.config.get("message_id")
self._last_artifact_id: Optional[str] = None
# ------------------------------------------------------------------
# Tool ABC
# ------------------------------------------------------------------
def get_actions_metadata(self) -> List[Dict[str, Any]]:
"""Return JSON metadata describing the create/edit/rewrite actions for tool schemas."""
kinds = sorted(_KIND_INFO.keys())
return [
{
"name": "create_artifact",
"description": (
"Render a new editable document from a JSON spec and store it as version 1. "
"The spec is the source of truth; the rendered file is derived. The response "
"carries a short ref (like `A1`) you can pass to edit_artifact/rewrite_artifact."
),
"active": True,
"parameters": {
"type": "object",
"properties": {
"kind": {
"type": "string",
"enum": kinds,
"description": (
"Document kind to render; `html` is an inline-rendered, versionable HTML report."
),
},
"title": {"type": "string", "description": "Optional artifact title."},
"spec": {"type": "object", "description": _SPEC_SYNOPSIS},
},
"required": ["kind", "spec"],
},
},
{
"name": "edit_artifact",
"description": (
"Apply a JSON merge-patch (RFC 7386) and/or append items to the current spec, "
"re-render, and append a new version. Preferred for small, targeted changes. "
"CAUTION: an array in spec_patch REPLACES the whole existing array — to add "
"slides/sections/blocks/sheets while keeping the existing ones, use spec_append."
),
"active": True,
"parameters": {
"type": "object",
"properties": {
"id": {
"type": "string",
"description": "Artifact to edit; accepts the short ref like `A1` "
"returned by a previous artifact action, or the full artifact id.",
},
"spec_patch": {
"type": "object",
"description": "RFC 7386 merge-patch; null values delete keys; arrays "
"replace the existing array wholesale.",
},
"spec_append": {
"type": "object",
"description": "Additive edit: {key: [items]} appends items to the "
'spec\'s top-level list, e.g. {"blocks": [{"type": "heading", "text": '
'"Risks"}]} keeps existing blocks and adds these after them.',
},
},
"required": ["id"],
},
},
{
"name": "rewrite_artifact",
"description": "Replace the spec wholesale, re-render, and append a new version.",
"active": True,
"parameters": {
"type": "object",
"properties": {
"id": {
"type": "string",
"description": "Artifact to rewrite; accepts the short ref like `A1` "
"returned by a previous artifact action, or the full artifact id.",
},
"spec": {"type": "object", "description": f"Replacement spec. {_SPEC_SYNOPSIS}"},
},
"required": ["id", "spec"],
},
},
]
def get_config_requirements(self) -> Dict[str, Any]:
"""Return configuration requirements (none beyond the deployment sandbox backend)."""
return {}
def get_artifact_id(self, action_name: str, **kwargs: Any) -> Optional[str]:
"""Return the produced artifact id so the UI artifact rail lights up."""
return self._last_artifact_id
# ------------------------------------------------------------------
# Dispatch
# ------------------------------------------------------------------
def execute_action(self, action_name: str, **kwargs: Any) -> Dict[str, Any]:
"""Dispatch a create/edit/rewrite action."""
self._last_artifact_id = None
if not self.user_id:
return {"status": "error", "error": "artifact_generator requires a valid user_id."}
if self.conversation_id is None and self.workflow_run_id is None:
return {"status": "error", "error": "artifact_generator requires a conversation_id or workflow_run_id."}
if jsonschema is None:
return {"status": "error", "error": "jsonschema is required for spec validation."}
if action_name == "create_artifact":
return self._create(**kwargs)
if action_name == "edit_artifact":
return self._edit(**kwargs)
if action_name == "rewrite_artifact":
return self._rewrite(**kwargs)
return {"status": "error", "error": f"unknown action: {action_name}"}
# ------------------------------------------------------------------
# Actions
# ------------------------------------------------------------------
def _create(self, **kwargs: Any) -> Dict[str, Any]:
"""Validate, render, and persist a new artifact at version 1."""
kind = kwargs.get("kind")
spec = kwargs.get("spec")
title = kwargs.get("title")
if kind not in _KIND_INFO:
return {"status": "error", "error": f"unsupported kind: {kind!r}; expected one of {sorted(_KIND_INFO)}."}
valid = self._validate(kind, spec)
if valid is not None:
return valid
rendered = self._render(kind, spec)
if rendered.get("error"):
return {"status": "error", "error": rendered["error"]}
info = _KIND_INFO[kind]
filename = self._filename(title, kind)
try:
ref = persist_new_artifact(
user_id=self.user_id,
kind=info["kind"],
data=rendered["data"],
filename=filename,
mime_type=info["mime"],
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),
)
except QuotaExceeded as exc:
return {"status": "error", "error": str(exc)}
if ref is None:
return {"status": "error", "error": "failed to persist artifact."}
self._last_artifact_id = ref["artifact_id"]
return {"status": "ok", **ref}
def _edit(self, **kwargs: Any) -> Dict[str, Any]:
"""Merge-patch and/or list-append the current spec, re-render, and append a version."""
spec_patch = kwargs.get("spec_patch")
spec_append = kwargs.get("spec_append")
if spec_patch is None and spec_append is None:
return {"status": "error", "error": "edit_artifact needs spec_patch and/or spec_append."}
if spec_patch is not None and not isinstance(spec_patch, dict):
return {"status": "error", "error": "spec_patch must be a JSON object (merge-patch)."}
if spec_append is not None and not isinstance(spec_append, dict):
return {"status": "error", "error": "spec_append must be a JSON object of {key: [items]}."}
loaded = self._load_current(kwargs.get("id"))
if loaded.get("error"):
return {"status": "error", "error": loaded["error"]}
new_spec = merge_patch(loaded["spec"], spec_patch) if spec_patch else dict(loaded["spec"] or {})
if spec_append:
appended = _apply_spec_append(new_spec, spec_append)
if "error" in appended:
return {"status": "error", "error": appended["error"]}
new_spec = appended["spec"]
return self._reversion(
loaded["artifact_id"], loaded["kind"], new_spec, "edit_artifact", loaded.get("title")
)
def _rewrite(self, **kwargs: Any) -> Dict[str, Any]:
"""Replace the spec wholesale, re-render, and append a version."""
spec = kwargs.get("spec")
loaded = self._load_current(kwargs.get("id"))
if loaded.get("error"):
return {"status": "error", "error": loaded["error"]}
return self._reversion(
loaded["artifact_id"], loaded["kind"], spec, "rewrite_artifact", loaded.get("title")
)
def _reversion(
self, artifact_id: str, kind: str, spec: Any, action: str, title: Optional[str] = None
) -> Dict[str, Any]:
"""Validate the new spec, re-render, and append the next version of an existing artifact."""
valid = self._validate(kind, spec)
if valid is not None:
return valid
rendered = self._render(kind, spec)
if rendered.get("error"):
return {"status": "error", "error": rendered["error"]}
info = _KIND_INFO[kind]
# Keep the original artifact's download name across versions (v2 of "Q3 Deck"
# must stay "Q3 Deck.pptx", not a generic "artifact.pptx").
filename = self._filename(title, kind)
try:
ref = append_artifact_version(
user_id=self.user_id,
artifact_id=artifact_id,
data=rendered["data"],
filename=filename,
mime_type=info["mime"],
spec=spec,
produced_by=self._produced_by(action, kind),
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
except QuotaExceeded as exc:
return {"status": "error", "error": str(exc)}
if ref is None:
return {"status": "error", "error": "failed to persist artifact version."}
self._last_artifact_id = ref["artifact_id"]
return {"status": "ok", **ref}
# ------------------------------------------------------------------
# Spec / render helpers
# ------------------------------------------------------------------
def _validate(self, kind: str, spec: Any) -> Optional[Dict[str, Any]]:
"""Return an error payload when ``spec`` is invalid for ``kind``, else None."""
if not isinstance(spec, dict):
return {"status": "error", "error": "spec must be a JSON object."}
try:
jsonschema.validate(spec, _SCHEMAS[kind])
except jsonschema.ValidationError as exc:
return {"status": "error", "error": f"invalid {kind} spec: {exc.message}"}
return None
def _load_current(self, raw_id: Any) -> Dict[str, Any]:
"""Resolve a short ref/uuid to its parent-scoped artifact and current-version spec for edit/rewrite."""
if not isinstance(raw_id, str) or not raw_id.strip():
return {"error": "id is required."}
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
# A ref (A1/A2/...) resolves to an id within this parent only; the
# resolved id is then re-checked through the parent-scoped gate so a
# ref can never reach another tenant.
artifact_id = resolve_artifact_id(
repo,
raw_id.strip(),
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
if artifact_id is None:
return {"error": f"artifact {raw_id} not found in this conversation/run."}
artifact = repo.get_artifact_in_parent(
artifact_id,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
if artifact is None:
return {"error": f"artifact {raw_id} not found in this conversation/run."}
version = repo.get_version(artifact_id, artifact["current_version"])
except Exception:
logger.exception("artifact_generator: failed to load artifact")
return {"error": f"failed to load artifact {raw_id}."}
if not version or version.get("spec") is None:
return {"error": f"artifact {raw_id} has no editable spec."}
kind = self._kind_for(artifact, version)
if kind is None:
return {"error": f"artifact {raw_id} is not a spec-rendered document."}
return {
"artifact_id": artifact_id,
"kind": kind,
"spec": version["spec"],
"title": artifact.get("title"),
}
@staticmethod
def _kind_for(artifact: Dict[str, Any], version: Dict[str, Any]) -> Optional[str]:
"""Resolve the spec kind from ``produced_by`` (preferred) or the version mime type."""
produced = version.get("produced_by")
if isinstance(produced, dict):
spec_kind = produced.get("spec_kind")
if spec_kind in _KIND_INFO:
return spec_kind
mime = version.get("mime_type") or ""
for spec_kind, info in _KIND_INFO.items():
if info["mime"] == mime:
return spec_kind
return None
def _render(self, kind: str, spec: Any) -> Dict[str, Any]:
"""Run the fixed renderer in the sandbox and return the produced file bytes."""
session_id = self._resolve_session_id()
if session_id is None:
return {"error": "artifact_generator requires a conversation_id or workflow_run_id."}
token = uuid.uuid4().hex
token_dir = f"artifacts/{token}"
spec_path = f"{token_dir}/spec.json"
out_path = f"{token_dir}/out.{_KIND_INFO[kind]['ext']}"
program = _RENDERERS[kind].format(spec_path=spec_path, out_path=out_path)
timeout = float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60))
manager = SandboxCreator.get_manager()
try:
manager.open(session_id, ttl=timeout)
except Exception as exc:
logger.exception("artifact_generator: failed to open sandbox session")
return {"error": f"sandbox unavailable: {type(exc).__name__}: {exc}"}
try:
# The spec rides in as a JSON file the program ``json.load``s; it is
# never interpolated into the program, so its contents stay data.
manager.put_file(session_id, spec_path, json.dumps(spec).encode("utf-8"))
result = manager.exec(session_id, program, timeout=timeout)
if not result.ok:
detail = (
f"{result.error_name}: {result.error_value}"
if result.error_name
else (result.error_value or "render failed")
)
return {"error": f"render failed: {detail}"}
data = manager.get_file(session_id, out_path)
except Exception as exc:
logger.exception("artifact_generator: render failed")
return {"error": f"render failed: {type(exc).__name__}: {exc}"}
finally:
# Drop this render's scratch dir, but do NOT close the session: it is the
# shared conversation/run session that code_executor(persist=True) keeps
# warm. A render is self-contained (it builds a document from the artifact
# spec, not from prior kernel state) and does not own that session -- its
# lifecycle belongs to the manager's TTL reaper / the conversation.
manager.remove_path(session_id, token_dir)
if not data:
return {"error": "renderer produced an empty file."}
return {"data": data}
# ------------------------------------------------------------------
# Misc helpers
# ------------------------------------------------------------------
def _produced_by(self, action: str, kind: str) -> Dict[str, Any]:
"""Build the ``produced_by`` provenance record, carrying the spec kind for re-editing."""
return {
"tool": "artifact_generator",
"action": action,
"spec_kind": kind,
"tool_id": self.tool_id,
}
@staticmethod
def _filename(title: Optional[str], kind: str) -> str:
"""Derive a download filename from a title (or a generic stem) plus the kind extension."""
if kind == "html":
return "report.html"
stem = (title or "artifact").strip() or "artifact"
return f"{stem}.{_KIND_INFO[kind]['ext']}"
def _resolve_session_id(self) -> Optional[str]:
"""Derive the sandbox session id from the bound conversation/run; sanitize to the gateway charset."""
raw = self.conversation_id or self.workflow_run_id
if not raw:
return None
sanitized = "".join(c if c.isalnum() or c in "-_" else "-" for c in str(raw))
return sanitized or None
+67
View File
@@ -0,0 +1,67 @@
"""Virtual short artifact handles (``A1``, ``A2``, ...) the model can type to reference an artifact.
A ref is NOT a stored column: ``A{n}`` is the artifact's STABLE per-parent ``ref_seq``, assigned at
creation and kept in the artifact's ``metadata``, so deleting an earlier artifact no longer
re-points a later ref the model already holds. Artifacts created before ``ref_seq`` existed have
none, so resolution falls back to the legacy positional (n-th by created_at) lookup. Refs resolve
only inside the caller's parent (``conversation_id`` or ``workflow_run_id``), never cross-tenant;
resolution still goes through the parent-scoped authz gate.
"""
from __future__ import annotations
import re
from typing import Any, Optional
from application.storage.db.base_repository import looks_like_uuid
_REF_RE = re.compile(r"^[Aa](\d+)$")
def make_ref(position: int) -> str:
"""Build the short ref string for a 1-based position (``1`` -> ``"A1"``)."""
return f"A{position}"
def parse_ref(value: Any) -> Optional[int]:
"""Parse a short ref like ``A1``/``a2`` into its 1-based position, or None when it is not a ref."""
if not isinstance(value, str):
return None
match = _REF_RE.match(value.strip())
if match is None:
return None
position = int(match.group(1))
return position if position >= 1 else None
def resolve_artifact_id(
repo: Any,
raw: Any,
*,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
) -> Optional[str]:
"""Resolve a short ref or a uuid to an artifact id, scoped to the caller's parent; None otherwise."""
position = parse_ref(raw)
if position is not None:
# A ref is the artifact's stable per-parent ``ref_seq``: resolve by it first so a
# deletion of an earlier artifact never re-points this ref. Legacy rows (created
# before ref_seq) and repos without the newer method fall back to the positional
# (n-th by created_at) lookup.
by_seq = getattr(repo, "resolve_id_by_ref_seq", None)
if callable(by_seq):
resolved = by_seq(
position,
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
)
if resolved is not None:
return resolved
return repo.artifact_id_at_position(
position,
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
)
if looks_like_uuid(raw):
return str(raw).strip()
return None
@@ -0,0 +1,152 @@
"""Lazily bridge a chat attachment into a conversation-scoped artifact when a tool references it.
A chat attachment lives in the ``attachments`` table (parsed to text for the LLM context); it is
not an artifact and so cannot be fed to ``code_executor`` / ``read_document`` directly. When one of
those tools references an attachment by id or filename, this module materializes it into a
conversation-scoped artifact on demand — only the request's own (already user-scoped) attachments are
reachable, and an already-bridged attachment is reused so repeated references never burn extra quota.
"""
from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional
from application.core.settings import settings
from application.sandbox.artifacts_capture import QuotaExceeded, persist_new_artifact
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.repositories.attachments import AttachmentsRepository
from application.storage.db.session import db_readonly
from application.storage.storage_creator import StorageCreator
logger = logging.getLogger(__name__)
class AttachmentBridgeError(Exception):
"""Raised when a matched attachment cannot be bridged (e.g. quota, unreadable bytes)."""
def _normalize_name(value: Any) -> str:
"""Lowercase + strip a filename for tolerant matching."""
return str(value or "").strip().lower()
def match_attachment(
attachments: Optional[List[Dict[str, Any]]], raw_ref: str, user_id: str
) -> Optional[Dict[str, Any]]:
"""Match a model-supplied id/filename against the caller's OWN request attachments; None otherwise.
Matching is confined to ``attachments`` (already user-scoped when loaded) so a forged id/name can
never reach another user's or conversation's attachment. An id match is re-verified against
``AttachmentsRepository.get_any(id, user_id)`` so only the owner's row is ever bridged. When two
attachments share a filename the first is chosen; reference by id to disambiguate.
"""
if not attachments or not raw_ref:
return None
ref = raw_ref.strip()
if not ref:
return None
ref_norm = _normalize_name(ref)
by_filename: Optional[Dict[str, Any]] = None
for attachment in attachments:
if not isinstance(attachment, dict):
continue
ids = {
str(attachment.get(key))
for key in ("id", "_id", "legacy_mongo_id")
if attachment.get(key) is not None
}
if ref in ids:
return _verify_owner(attachment, user_id)
filename = attachment.get("filename")
if by_filename is None and filename and _normalize_name(filename) == ref_norm:
by_filename = attachment
if by_filename is not None:
return _verify_owner(by_filename, user_id)
return None
def _verify_owner(attachment: Dict[str, Any], user_id: str) -> Optional[Dict[str, Any]]:
"""Re-confirm the attachment belongs to ``user_id`` via the user-scoped repo; in-memory dict on hit."""
attachment_id = attachment.get("id") or attachment.get("_id") or attachment.get("legacy_mongo_id")
if attachment_id is None:
return None
try:
with db_readonly() as conn:
owned = AttachmentsRepository(conn).get_any(str(attachment_id), user_id)
except Exception:
logger.exception("attachment_bridge: ownership re-check failed")
return None
# Prefer the DB row (authoritative upload_path/mime) but only when it confirms ownership.
return owned if owned is not None else None
def bridge_attachment(
attachment: Dict[str, Any], *, user_id: str, conversation_id: str
) -> str:
"""Return the conversation artifact id for ``attachment``, reusing an existing bridge or creating one.
Idempotent (best-effort): an artifact already derived from this attachment in this conversation
(matched via its version ``produced_by.attachment_id``) is reused, so a second reference never
consumes a new quota slot. The reuse is a read-then-write across transactions, so two concurrent
references to the same not-yet-bridged attachment may each create one. Otherwise the attachment
bytes are read server-side and persisted as a conversation-scoped ``file`` artifact (server-computed
size/sha256/storage key).
"""
attachment_id = str(attachment.get("id") or attachment.get("_id") or attachment.get("legacy_mongo_id"))
try:
with db_readonly() as conn:
existing = ArtifactsRepository(conn).find_bridged_attachment(
attachment_id, conversation_id=conversation_id
)
except Exception:
logger.exception("attachment_bridge: idempotency lookup failed")
existing = None
if existing is not None:
return str(existing["id"])
upload_path = attachment.get("upload_path") or attachment.get("path")
if not upload_path:
raise AttachmentBridgeError(f"attachment {attachment_id} has no stored content.")
filename = attachment.get("filename") or "attachment"
mime_type = attachment.get("mime_type") or "application/octet-stream"
# Reject oversize attachments BEFORE buffering them: the authoritative ``size``
# column lets us avoid pulling a multi-hundred-MB file fully into worker memory,
# and the bounded read below backstops a missing/lying ``size``.
max_bytes = int(getattr(settings, "ARTIFACT_MAX_BYTES", 0) or 0)
declared_size = attachment.get("size")
if max_bytes and isinstance(declared_size, (int, float)) and declared_size > max_bytes:
raise AttachmentBridgeError(
f"attachment {attachment_id} exceeds the {max_bytes}-byte artifact size limit."
)
try:
file_obj = StorageCreator.get_storage().get_file(upload_path)
try:
data = file_obj.read(max_bytes + 1) if max_bytes else file_obj.read()
finally:
close = getattr(file_obj, "close", None)
if callable(close):
close()
except Exception as exc:
logger.exception("attachment_bridge: failed to read attachment bytes")
raise AttachmentBridgeError(f"failed to read attachment {attachment_id}.") from exc
if max_bytes and len(data) > max_bytes:
raise AttachmentBridgeError(
f"attachment {attachment_id} exceeds the {max_bytes}-byte artifact size limit."
)
try:
ref = persist_new_artifact(
user_id=user_id,
kind="file",
data=data,
filename=filename,
mime_type=mime_type,
title=filename,
conversation_id=conversation_id,
produced_by={"attachment_id": attachment_id, "source": "chat_attachment"},
)
except QuotaExceeded as exc:
raise AttachmentBridgeError(str(exc)) from exc
if ref is None:
raise AttachmentBridgeError(f"failed to bridge attachment {attachment_id}.")
return str(ref["artifact_id"])
+474
View File
@@ -0,0 +1,474 @@
"""Code Executor tool: run sandboxed code in a semi-persistent session and capture produced files as artifacts."""
from __future__ import annotations
import logging
import re
from typing import Any, Dict, List, Optional, Tuple
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.agents.tools.attachment_bridge import (
AttachmentBridgeError,
bridge_attachment,
match_attachment,
)
from application.agents.tools.base import Tool
from application.core.settings import settings
from application.sandbox.artifacts_capture import (
MAX_CAPTURED_FILES,
capture_artifacts,
snapshot_signatures,
unique_input_path,
)
from application.sandbox.artifacts_capture import (
infer_mime as _infer_mime,
)
from application.sandbox.artifacts_capture import (
kind_for_mime as _kind_for_mime,
)
from application.sandbox.base import ExecResult
from application.sandbox.sandbox_creator import SandboxCreator
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
from application.storage.storage_creator import StorageCreator
from application.utils import safe_filename
logger = logging.getLogger(__name__)
# Re-exported for back-compat: callers (and tests) import these mime helpers
# from this module; they now live in the shared capture helper.
__all__ = ["CodeExecutorTool", "_infer_mime", "_kind_for_mime", "_tail", "_OUTPUT_TAIL_BYTES"]
# Maximum bytes of stdout/stderr returned to the LLM. The raw stream is never
# forwarded; only this tail keeps binary/runaway output out of the context.
_OUTPUT_TAIL_BYTES = 4000
# Session ids become a kernel workspace path component; the gateway only accepts
# [A-Za-z0-9_-]+, so any disallowed character is stripped before binding.
_SESSION_ID_RE = re.compile(r"[^A-Za-z0-9_-]+")
def _tail(stream: Optional[str]) -> str:
"""Return the trailing slice of ``stream`` bounded by ``_OUTPUT_TAIL_BYTES``."""
if not stream:
return ""
if len(stream) <= _OUTPUT_TAIL_BYTES:
return stream
return stream[-_OUTPUT_TAIL_BYTES:]
class CodeExecutorTool(Tool):
"""Code Executor
Run code in a sandboxed session; files it writes become downloadable artifacts.
"""
def __init__(self, tool_config: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None) -> None:
"""Bind the tool to the invoker and its conversation/run-scoped sandbox session."""
self.config: Dict[str, Any] = tool_config or {}
self.user_id: Optional[str] = user_id
self.tool_id: Optional[str] = self.config.get("tool_id")
self.conversation_id: Optional[str] = self.config.get("conversation_id")
self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id")
self.message_id: Optional[str] = self.config.get("message_id")
# Static, deployment-level approval gate (mirrors the action metadata flag).
self._require_approval: bool = bool(self.config.get("require_approval", False))
self._last_artifact_id: Optional[str] = None
# ------------------------------------------------------------------
# Tool ABC
# ------------------------------------------------------------------
@staticmethod
def _environment_note() -> str:
"""Backend-specific note on what the sandbox has preinstalled.
Without this the model discovers the environment by failing: importing
pandas on a bare image, or pip-installing libraries that are already
baked in. Keep the package lists in sync with deployment/sandbox/Dockerfile
(jupyter) and scripts/build_daytona_snapshot.py (daytona snapshot).
"""
backend = str(getattr(settings, "SANDBOX_BACKEND", "jupyter") or "jupyter").lower()
if backend == "daytona":
if getattr(settings, "DAYTONA_SNAPSHOT", None):
return (
"Preinstalled beyond the stdlib: python-pptx, python-docx, openpyxl, "
"reportlab, lxml, pillow. pip install anything else from within the code "
"before importing it."
)
return (
"Only the Python stdlib is preinstalled. pip install any third-party "
"package (pandas, python-docx, ...) from within the code before importing it."
)
return (
"Preinstalled beyond the stdlib: pandas, matplotlib, python-pptx, python-docx, "
"openpyxl, reportlab. pip install anything else from within the code before "
"importing it."
)
def get_actions_metadata(self) -> List[Dict[str, Any]]:
"""Return JSON metadata describing the ``run_code`` action for tool schemas."""
return [
{
"name": "run_code",
"description": (
"Execute Python in a sandboxed, stateful session bound to this conversation. "
"Files written by the code are saved as downloadable artifacts (write throwaway "
"files under `tmp/`, or pass `outputs` to save only specific files); only a compact "
"summary (output tail + artifact references) is returned, never raw bytes. "
"Each call is capped at ~60s of wall-clock; for longer work, start it in the "
"background and poll with additional run_code calls (use persist=true to keep state). "
+ self._environment_note()
),
"active": True,
"require_approval": self._require_approval,
"parameters": {
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "Python source to execute in the session. Install packages from "
"within the code itself (e.g. subprocess pip install) if needed.",
},
"inputs": {
"type": "array",
"items": {"type": "string"},
"description": "Files to materialize into the workspace; each accepts the short "
"ref like `A1` returned by a previous artifact action, a full artifact id, or "
"the name/id of a file the user attached to this conversation. Each is staged "
"at `inputs/<filename>` before the code runs — read it from that path (the "
"result's `inputs_loaded` echoes the exact staged paths).",
},
"outputs": {
"type": "array",
"items": {"type": "string"},
"description": "Filenames or globs (e.g. `report.pdf`, `*.csv`) to save as "
"downloadable artifacts. When set, only matching files are saved; when omitted, "
"every produced file is saved except scratch paths under `tmp/`.",
},
"ttl": {
"type": "integer",
"description": "Keep-alive lifetime (seconds) for the session; clamped by SANDBOX_MAX_TTL.",
},
"persist": {
"type": "boolean",
"description": (
"Keep the session warm after the call (state survives the next run). "
"The session is kept alive when this is true or a positive ttl is given "
"(clamped by SANDBOX_MAX_TTL); otherwise it is closed after the run."
),
},
"capture_artifacts": {
"type": "boolean",
"description": "Save produced workspace files as downloadable artifacts "
"(default: true). Set false for setup or install-only steps that write nothing "
"worth keeping.",
},
},
"required": ["code"],
},
}
]
def get_config_requirements(self) -> Dict[str, Any]:
"""Return configuration requirements (none; approval is an action-level flag,
and the sandbox backend is a deployment-level setting)."""
return {}
def get_artifact_id(self, action_name: str, **kwargs: Any) -> Optional[str]:
"""Return the primary produced artifact id so the UI artifact rail lights up."""
return self._last_artifact_id
def preview_decision(self, action_name: str, params: dict) -> Tuple[bool, bool]:
"""Return ``(requires_approval, denylist_forced)`` for the approval gate; never denylist-forced here."""
if action_name != "run_code":
return True, False
return self._require_approval, False
# ------------------------------------------------------------------
# Execution
# ------------------------------------------------------------------
def execute_action(self, action_name: str, **kwargs: Any) -> Dict[str, Any]:
"""Dispatch a tool action; only ``run_code`` is supported."""
if action_name != "run_code":
return {"status": "error", "error": f"unknown action: {action_name}"}
self._last_artifact_id = None
return self._run_code(**kwargs)
def _run_code(self, **kwargs: Any) -> Dict[str, Any]:
"""Bind a session, materialize inputs, execute, and capture produced artifacts."""
if not self.user_id:
return {"status": "error", "error": "code_executor requires a valid user_id."}
session_id = self._resolve_session_id()
if session_id is None:
return {"status": "error", "error": "code_executor requires a conversation_id or workflow_run_id."}
code = kwargs.get("code")
if not isinstance(code, str) or not code.strip():
return {"status": "error", "error": "code is required."}
should_capture = kwargs.get("capture_artifacts", True)
outputs = self._normalize_outputs(kwargs.get("outputs"))
ttl = self._coerce_int(kwargs.get("ttl"))
timeout = self._exec_timeout()
inputs = kwargs.get("inputs") or []
manager = SandboxCreator.get_manager()
try:
manager.open(session_id, ttl=ttl)
except Exception as exc:
logger.exception("code_executor: failed to open sandbox session")
return {"status": "error", "error": f"sandbox unavailable: {type(exc).__name__}: {exc}"}
try:
materialized = self._materialize_inputs(manager, session_id, inputs)
if materialized.get("error"):
return {"status": "error", "error": materialized["error"]}
pre_signatures: Dict[str, Tuple[int, Optional[str]]] = {}
if should_capture:
pre_signatures = self._snapshot_signatures(manager, session_id)
try:
result = manager.exec(session_id, code, timeout=timeout)
except Exception as exc:
logger.exception("code_executor: exec raised")
return {"status": "error", "error": f"execution failed: {type(exc).__name__}: {exc}"}
# Capture even on error/timeout while the runtime remains reachable
# so partial outputs aren't lost; capture never masks the run status.
artifacts: List[Dict[str, Any]] = []
if should_capture and not result.runtime_invalidated:
try:
artifacts = self._capture_artifacts(manager, session_id, pre_signatures, outputs)
except Exception:
logger.exception("code_executor: artifact capture failed")
return self._shape_payload(result, artifacts, materialized.get("loaded", []))
finally:
if not self._keep_alive(kwargs.get("persist"), ttl):
try:
manager.close(session_id)
except Exception:
logger.exception("code_executor: session close failed")
# ------------------------------------------------------------------
# Inputs / outputs
# ------------------------------------------------------------------
def _materialize_inputs(self, manager: Any, session_id: str, inputs: List[Any]) -> Dict[str, Any]:
"""Fetch parent-scoped input artifacts and copy their current-version bytes into the workspace."""
loaded: List[str] = []
if not inputs:
return {"loaded": loaded}
storage = StorageCreator.get_storage()
# Two inputs whose current versions share a filename would clobber each other at
# the same ``inputs/{name}`` path; track used paths and disambiguate deterministically.
used_paths: set = set()
for raw_id in inputs:
raw = str(raw_id).strip()
if not raw:
continue
artifact_id: Optional[str] = raw
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
# A short ref (A1/A2/...) resolves to an id within this parent
# only; the resolved id still passes through the parent-scoped
# gate so a ref can never reach another tenant.
artifact_id = resolve_artifact_id(
repo,
raw,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
artifact = (
repo.get_artifact_in_parent(
artifact_id,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
if artifact_id is not None
else None
)
if artifact is None:
# Conversation scope only: a raw ref that is not an artifact
# may name a chat attachment; bridge it on demand. Workflows
# bridge attachments up front, so never double-bridge there.
bridged_id = self._bridge_chat_attachment(raw)
if isinstance(bridged_id, dict):
return bridged_id # error payload
if bridged_id is None:
return {"error": f"input artifact {raw} not found in this conversation/run."}
artifact_id = bridged_id
artifact = repo.get_artifact_in_parent(artifact_id, conversation_id=self.conversation_id)
if artifact is None:
return {"error": f"input artifact {raw} not found in this conversation/run."}
version = repo.get_version(artifact_id, artifact["current_version"])
except Exception:
logger.exception("code_executor: failed to load input artifact")
return {"error": f"failed to load input artifact {artifact_id}."}
if not version or not version.get("storage_path"):
return {"error": f"input artifact {artifact_id} has no stored content."}
# Reject an oversize input BEFORE buffering it: the declared ``size``
# avoids pulling a huge file into worker memory, and the bounded read
# below backstops a missing/lying size column.
max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 0) or 0)
declared_size = version.get("size")
if max_bytes and isinstance(declared_size, (int, float)) and declared_size > max_bytes:
return {"error": f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit."}
filename = safe_filename(version.get("filename") or artifact_id)
try:
file_obj = storage.get_file(version["storage_path"])
try:
data = file_obj.read(max_bytes + 1) if max_bytes else file_obj.read()
finally:
close = getattr(file_obj, "close", None)
if callable(close):
close()
except Exception:
logger.exception("code_executor: failed to read input artifact bytes")
return {"error": f"failed to read input artifact {artifact_id}."}
if max_bytes and len(data) > max_bytes:
return {"error": f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit."}
rel_path = unique_input_path(f"inputs/{filename}", used_paths)
try:
manager.put_file(session_id, rel_path, data)
except Exception:
logger.exception("code_executor: put_file failed for input artifact")
return {"error": f"failed to stage input artifact {artifact_id} into the workspace."}
loaded.append(rel_path)
return {"loaded": loaded}
def _bridge_chat_attachment(self, raw: str) -> Any:
"""Bridge a referenced chat attachment to a conversation artifact id; None on miss, error dict on failure."""
if not self.conversation_id or not self.user_id:
return None
attachment = match_attachment(self.config.get("attachments"), raw, self.user_id)
if attachment is None:
return None
try:
return bridge_attachment(attachment, user_id=self.user_id, conversation_id=self.conversation_id)
except AttachmentBridgeError as exc:
return {"error": f"failed to attach {raw}: {exc}"}
# Cap the per-run capture work so a workspace full of pre-existing files
# can't turn one exec into an unbounded read+persist sweep.
_MAX_CAPTURED_FILES = MAX_CAPTURED_FILES
def _snapshot_signatures(self, manager: Any, session_id: str) -> Dict[str, Tuple[int, Optional[str]]]:
"""Map each non-input workspace file to a (size, sha256) signature for change detection."""
return snapshot_signatures(manager, session_id)
@staticmethod
def _normalize_outputs(raw: Any) -> Optional[List[str]]:
"""Coerce the ``outputs`` arg to a list of non-empty glob strings, or None.
Tolerates a bare string (some models pass one instead of an array); an empty
or non-list value means "no allow-list" (auto-capture).
"""
if isinstance(raw, str):
raw = [raw]
if not isinstance(raw, list):
return None
patterns = [str(p).strip() for p in raw if isinstance(p, str) and str(p).strip()]
return patterns or None
def _capture_artifacts(
self,
manager: Any,
session_id: str,
pre_signatures: Dict[str, Tuple[int, Optional[str]]],
outputs: Optional[List[str]] = None,
) -> List[Dict[str, Any]]:
"""Persist produced workspace files (only ``outputs`` globs when given)."""
captured = capture_artifacts(
manager,
session_id,
pre_signatures,
user_id=self.user_id,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
message_id=self.message_id,
produced_by={
"tool": "code_executor",
"action": "run_code",
"session_id": session_id,
},
outputs=outputs,
)
if captured:
self._last_artifact_id = captured[0]["artifact_id"]
return captured
def _shape_payload(
self, result: ExecResult, artifacts: List[Dict[str, Any]], inputs_loaded: List[str]
) -> Dict[str, Any]:
"""Build the compact LLM-facing payload; raw bytes never appear here."""
status = "ok" if result.ok else "error"
payload: Dict[str, Any] = {
"status": status,
"stdout_tail": _tail(result.stdout),
"artifacts": artifacts,
}
stderr_tail = _tail(result.stderr)
if stderr_tail:
payload["stderr_tail"] = stderr_tail
if not result.ok:
if self._is_timeout(result):
cap = int(self._exec_timeout())
payload["error"] = (
f"Execution timed out. Each run_code call is capped at {cap}s and the limit "
"cannot be raised. For long-running work, start it in the background (e.g. launch a "
"subprocess or `nohup ... &` and write progress to a file) and return immediately, "
"then poll with additional run_code calls to check on it. Pass persist=true (or a "
"ttl) so the background process and its files survive between calls."
)
else:
payload["error"] = (
f"{result.error_name}: {result.error_value}"
if result.error_name
else (result.error_value or "execution error")
)
if inputs_loaded:
payload["inputs_loaded"] = inputs_loaded
return payload
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
def _resolve_session_id(self) -> Optional[str]:
"""Derive a sandbox session id from the bound conversation/run; sanitize to the gateway charset."""
raw = self.conversation_id or self.workflow_run_id
if not raw:
return None
sanitized = _SESSION_ID_RE.sub("-", str(raw))
return sanitized or None
@staticmethod
def _coerce_int(value: Any) -> Optional[int]:
"""Coerce a value to a positive int, or None when absent/invalid."""
if value is None:
return None
try:
parsed = int(value)
except (TypeError, ValueError):
return None
return parsed if parsed > 0 else None
@staticmethod
def _exec_timeout() -> float:
"""Return the fixed per-run wall-clock cap (SANDBOX_EXEC_TIMEOUT; not caller-adjustable)."""
return float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60))
@staticmethod
def _is_timeout(result: ExecResult) -> bool:
"""True when a failed exec looks like a wall-clock timeout (any backend's naming/message)."""
blob = f"{result.error_name or ''} {result.error_value or ''}".lower()
return "timeout" in blob or "timed out" in blob
@staticmethod
def _keep_alive(persist: Any, ttl: Optional[int]) -> bool:
"""True when the agent asked to keep the session warm after the call."""
return bool(persist) or (ttl is not None and ttl > 0)
+322
View File
@@ -0,0 +1,322 @@
"""Read Document tool: parse an input artifact to text/markdown/structured/chunks via the backend parser.
The ``read_document`` action resolves a parent-scoped input artifact, enqueues a
``parse_document`` task on the dedicated ``parsing`` Celery queue, and awaits the
result with a timeout. The run-scoped authz gate is enforced TWICE — here before
enqueue (reject cross-tenant) and again in the worker (re-resolve, never trusting a
raw path). When a ``json_schema`` is supplied the structured payload is validated
through the existing jsonschema path; the full result may also be persisted as a
``data`` artifact by reference (handled in the worker).
"""
from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional
from celery import current_task
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.agents.tools.attachment_bridge import (
AttachmentBridgeError,
bridge_attachment,
match_attachment,
)
from application.agents.tools.base import Tool
from application.core.json_schema_utils import (
JsonSchemaValidationError,
normalize_json_schema_payload,
)
from application.core.settings import settings
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
logger = logging.getLogger(__name__)
try:
import jsonschema
except Exception: # pragma: no cover - jsonschema is a declared dependency
jsonschema = None # type: ignore[assignment]
class ReadDocumentTool(Tool):
"""Read Document
Parse a document (PDF, Word, PowerPoint, ...) to text, markdown, or structured data.
"""
# Hidden from the Add-Tool catalog; surfaced (workflow-only) via the
# BUILTIN_AGENT_TOOLS synthetic-id path. Does not gate tool_manager loading
# nor synthetic-id execution.
internal: bool = True
def __init__(self, tool_config: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None) -> None:
"""Bind the tool to the invoker and its conversation/run scope."""
self.config: Dict[str, Any] = tool_config or {}
self.user_id: Optional[str] = user_id
self.tool_id: Optional[str] = self.config.get("tool_id")
self.conversation_id: Optional[str] = self.config.get("conversation_id")
self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id")
self.message_id: Optional[str] = self.config.get("message_id")
self._last_artifact_id: Optional[str] = None
# ------------------------------------------------------------------
# Tool ABC
# ------------------------------------------------------------------
def get_actions_metadata(self) -> List[Dict[str, Any]]:
"""Return JSON metadata describing the ``read_document`` action for tool schemas."""
return [
{
"name": "read_document",
"description": (
"Read a document artifact (pdf/docx/pptx/...) and return its parsed content as "
"markdown, plain text, structured JSON (with tables), or chunks. Optionally "
"validate the structured result against a json_schema and persist it as a "
"downloadable data artifact."
),
"active": True,
"parameters": {
"type": "object",
"properties": {
"input": {
"type": "string",
"description": "Document to read; accepts the short ref like `A1` returned by a "
"previous artifact action, a full artifact id, or the name/id of a file the user "
"attached to this conversation.",
},
"output": {
"type": "string",
"enum": ["markdown", "text", "structured", "chunks"],
"description": "Shape of the parsed result (default: markdown). Note: "
"`structured` always uses the Docling engine regardless of `engine` "
"(the `fast` engine is markdown/text only).",
},
"ocr": {
"type": "string",
"enum": ["auto", "on", "off"],
"description": "OCR mode for scanned pages/images (default: auto, follows server config).",
},
"pages": {
"type": "string",
"description": "Optional page range to read, e.g. `1-3` or `2` (best-effort).",
},
"engine": {
"type": "string",
"enum": ["auto", "docling", "fast"],
"description": "Parser engine (default: auto). `fast` is a lighter "
"markdown/text-only engine; it is ignored when `output='structured'`, "
"which always uses Docling.",
},
"max_chars": {
"type": "integer",
"description": "Optional cap on returned characters.",
},
"include_tables": {
"type": "boolean",
"description": "Include extracted tables in the result (default: true).",
},
"json_schema": {
"type": "object",
"description": "Optional JSON schema the structured payload must satisfy.",
},
"persist": {
"type": "boolean",
"description": "Persist the parsed result as a downloadable data artifact (default true).",
},
},
"required": ["input"],
},
}
]
def get_config_requirements(self) -> Dict[str, Any]:
"""Return configuration requirements (none beyond a running parsing worker)."""
return {}
def get_artifact_id(self, action_name: str, **kwargs: Any) -> Optional[str]:
"""Return the persisted parse artifact id so the UI artifact rail lights up."""
return self._last_artifact_id
# ------------------------------------------------------------------
# Dispatch
# ------------------------------------------------------------------
def execute_action(self, action_name: str, **kwargs: Any) -> Dict[str, Any]:
"""Dispatch a tool action; only ``read_document`` is supported."""
self._last_artifact_id = None
if action_name != "read_document":
return {"status": "error", "error": f"unknown action: {action_name}"}
if not self.user_id:
return {"status": "error", "error": "read_document requires a valid user_id."}
if self.conversation_id is None and self.workflow_run_id is None:
return {"status": "error", "error": "read_document requires a conversation_id or workflow_run_id."}
return self._read(**kwargs)
# ------------------------------------------------------------------
# Read
# ------------------------------------------------------------------
def _read(self, **kwargs: Any) -> Dict[str, Any]:
"""Resolve the input run-scoped (reject cross-tenant before enqueue), enqueue+await, validate."""
input_id = kwargs.get("input")
json_schema = kwargs.get("json_schema")
if not isinstance(input_id, str) or not input_id.strip():
return {"status": "error", "error": "input artifact id is required."}
if json_schema is not None:
schema_err = self._check_schema(json_schema)
if schema_err is not None:
return schema_err
artifact_id = self._resolve_input(input_id.strip())
if isinstance(artifact_id, dict):
return artifact_id # error payload
options = {
"output": kwargs.get("output", "markdown"),
"ocr": kwargs.get("ocr", "auto"),
"pages": kwargs.get("pages"),
"engine": kwargs.get("engine", "auto"),
"max_chars": kwargs.get("max_chars"),
"include_tables": kwargs.get("include_tables", True),
"persist": kwargs.get("persist", True),
"tool_id": self.tool_id,
}
result = self._dispatch(artifact_id, options)
if result.get("status") == "error":
return result
if json_schema is not None:
valid = self._validate(json_schema, result.get("structured"))
if valid is not None:
return valid
artifact = result.get("artifact")
if isinstance(artifact, dict) and artifact.get("artifact_id"):
self._last_artifact_id = artifact["artifact_id"]
return result
def _dispatch(self, artifact_id: str, options: Dict[str, Any]) -> Dict[str, Any]:
"""Parse INLINE inside a Celery worker, else dispatch to the parsing queue and await.
This tool runs in the WEB process (/stream) OR inside a Celery worker
(headless/scheduled/workflow agents). When it already runs inside a worker that also
serves the ``parsing`` queue (the shipped default ``-Q docsgpt,parsing``), dispatching
and blocking on ``get()`` would self-deadlock: concurrent agent tasks each hold a pool
slot blocked in ``get()`` so ``parse_document`` never gets a free slot. So parse INLINE
in-process inside a worker; only dispatch+await (degrading on timeout/failure) from web.
"""
parent = self._parent()
# ``current_task`` is a Celery proxy: truthy only while this runs inside a worker task,
# falsy in the web process (the bare proxy is NOT identity-None, so test truthiness).
if current_task:
from application.worker import run_parse_document
try:
result = run_parse_document(artifact_id, parent, self.user_id, options)
except Exception as exc:
logger.exception("read_document: inline parse failed")
return {"status": "error", "error": f"document parsing failed: {type(exc).__name__}: {exc}"}
if not isinstance(result, dict):
return {"status": "error", "error": "document parsing produced an unexpected result."}
return result
from celery.exceptions import TimeoutError as CeleryTimeoutError
from application.api.user.tasks import parse_document
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)
# 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.
result = async_result.get(timeout=timeout, disable_sync_subtasks=False)
except (CeleryTimeoutError, TimeoutError):
return {"status": "error", "error": f"document parsing timed out after {int(timeout)}s."}
except Exception as exc:
logger.exception("read_document: parse task failed")
return {"status": "error", "error": f"document parsing failed: {type(exc).__name__}: {exc}"}
if not isinstance(result, dict):
return {"status": "error", "error": "document parsing produced an unexpected result."}
return result
# ------------------------------------------------------------------
# Input resolution (run-scoped gate, before enqueue)
# ------------------------------------------------------------------
def _resolve_input(self, raw_id: str) -> Any:
"""Resolve a short ref/uuid to a parent-scoped artifact id; an error dict on miss/cross-tenant."""
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
artifact_id = resolve_artifact_id(
repo,
raw_id,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
artifact = (
repo.get_artifact_in_parent(
artifact_id,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
if artifact_id is not None
else None
)
except Exception:
logger.exception("read_document: failed to resolve input artifact")
return {"status": "error", "error": f"failed to resolve input artifact {raw_id}."}
if artifact is None:
# Conversation scope only: a raw ref that is not an artifact may name a
# chat attachment; bridge it on demand. Workflows bridge up front.
bridged_id = self._bridge_chat_attachment(raw_id)
if isinstance(bridged_id, dict):
return bridged_id
if bridged_id is not None:
return bridged_id
return {"status": "error", "error": f"input artifact {raw_id} not found in this conversation/run."}
return str(artifact_id)
def _bridge_chat_attachment(self, raw_id: str) -> Any:
"""Bridge a referenced chat attachment to a conversation artifact id; None on miss, error dict on failure."""
if not self.conversation_id or not self.user_id:
return None
attachment = match_attachment(self.config.get("attachments"), raw_id, self.user_id)
if attachment is None:
return None
try:
return bridge_attachment(attachment, user_id=self.user_id, conversation_id=self.conversation_id)
except AttachmentBridgeError as exc:
return {"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:
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}
# ------------------------------------------------------------------
# Schema validation
# ------------------------------------------------------------------
@staticmethod
def _check_schema(json_schema: Any) -> Optional[Dict[str, Any]]:
"""Return an error payload when ``json_schema`` itself is malformed, else None."""
try:
normalize_json_schema_payload(json_schema)
except JsonSchemaValidationError as exc:
return {"status": "error", "error": f"invalid json_schema: {exc}"}
return None
@staticmethod
def _validate(json_schema: Any, instance: Any) -> Optional[Dict[str, Any]]:
"""Validate ``instance`` against the (already-normalized) json_schema; error payload on mismatch."""
if jsonschema is None:
return {"status": "error", "error": "jsonschema is required for json_schema validation."}
if instance is None:
return {"status": "error", "error": "json_schema validation requires output='structured'."}
schema = normalize_json_schema_payload(json_schema)
try:
jsonschema.validate(instance=instance, schema=schema)
except jsonschema.exceptions.ValidationError as exc:
return {"status": "error", "error": f"parsed structure did not match json_schema: {exc.message}"}
return None
+24 -2
View File
@@ -29,7 +29,18 @@ class ToolManager:
for member_name, obj in inspect.getmembers(module, inspect.isclass):
if issubclass(obj, Tool) and obj is not Tool:
if (
tool_name in {"mcp_tool", "notes", "memory", "todo_list", "scheduler", "remote_device"}
tool_name
in {
"mcp_tool",
"notes",
"memory",
"todo_list",
"scheduler",
"remote_device",
"code_executor",
"artifact_generator",
"read_document",
}
and user_id
):
return obj(tool_config, user_id)
@@ -40,7 +51,18 @@ class ToolManager:
if tool_name not in self.tools:
raise ValueError(f"Tool '{tool_name}' not loaded")
if (
tool_name in {"mcp_tool", "memory", "todo_list", "notes", "scheduler", "remote_device"}
tool_name
in {
"mcp_tool",
"memory",
"todo_list",
"notes",
"scheduler",
"remote_device",
"code_executor",
"artifact_generator",
"read_document",
}
and user_id
):
tool_config = self.config.get(tool_name, {})
+289 -47
View File
@@ -1,6 +1,6 @@
import logging
from datetime import datetime, timezone
from typing import Any, Dict, Generator, Optional
from typing import Any, Dict, Generator, List, Optional, Tuple
from application.agents.base import BaseAgent
from application.agents.workflows.schemas import (
@@ -12,7 +12,9 @@ from application.agents.workflows.schemas import (
WorkflowRun,
)
from application.agents.workflows.workflow_engine import WorkflowEngine
from application.logging import log_activity, LogContext
from application.core.settings import settings
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
from application.storage.db.repositories.workflow_nodes import WorkflowNodesRepository
@@ -22,6 +24,10 @@ from application.storage.db.session import db_readonly, db_session
logger = logging.getLogger(__name__)
# Per-run cap on attachments staged as run-scoped artifacts; the remainder is
# dropped (the per-user artifact quota is only a best-effort soft cap).
_MAX_INPUT_DOCUMENTS = 25
class WorkflowAgent(BaseAgent):
"""A specialized agent that executes predefined workflows."""
@@ -39,23 +45,72 @@ class WorkflowAgent(BaseAgent):
self.workflow_owner = workflow_owner
self._workflow_data = workflow
self._engine: Optional[WorkflowEngine] = None
self._run_persisted = False
# Set to a message when the input-document bridge fails fatally (quota), so
# the run is finalized FAILED instead of running with missing documents.
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."}
return
self._engine = WorkflowEngine(graph, self)
yield from self._engine.execute({}, query)
self._save_workflow_run(query)
# Two distinct identities: the workflow *owner* (A) owns the workflow
# definition and is used to resolve the workflow row; the *runner* (B,
# the caller) owns the run and its artifacts. They are the same user
# except for a shared agent, where B != A. Owning run artifacts by the
# runner means quota is charged to the uploader and the caller can read
# 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)
self._run_persisted = pg_workflow_id is not None
try:
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
# than executing nodes with silently-missing documents.
self._bridge_error = str(exc)
yield {
"type": "error",
"user_facing": True,
"error": (
"This run's input documents exceed your artifact storage quota. "
"Delete some artifacts and try again."
),
}
self._finalize_workflow_run(workflow_owner_id, run_user_id, pg_workflow_id, query)
return
# Non-fatal: some attachments were dropped (oversize / unreadable). Emit a
# ``notice`` -- NOT an ``error`` -- so the client surfaces which were dropped
# without marking the turn failed or ending the stream (an ``error`` event is
# terminal client-side and disables reconnect). The run then still executes
# with the documents that did bridge.
if dropped:
yield {"type": "notice", "notice": " ".join(dropped)}
self._engine.run_persisted = self._run_persisted
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:
@@ -114,9 +169,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:
@@ -126,10 +179,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)
@@ -141,10 +191,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(
@@ -180,50 +232,243 @@ class WorkflowAgent(BaseAgent):
logger.error(f"Failed to load workflow from database: {e}")
return None
def _save_workflow_run(self, query: str) -> None:
if not self._engine:
return
def _resolve_owner_id(self) -> Optional[str]:
"""Resolve the workflow *owner* (used to resolve the owned workflow row)."""
owner_id = self.workflow_owner
if not owner_id and isinstance(self.decoded_token, dict):
owner_id = self.decoded_token.get("sub")
return owner_id
def _resolve_run_user_id(self, workflow_owner_id: Optional[str]) -> Optional[str]:
"""Resolve the *runner* (caller) who owns the run and its artifacts.
Equals the owner for a user running their own workflow (and for external
API-key calls, where the key owner is the caller); for a shared agent it
is the calling user, so their uploads/outputs are owned by and readable
to them rather than silently accruing under the agent owner's account.
"""
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]:
"""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
wf_repo = WorkflowsRepository(conn)
if looks_like_uuid(self.workflow_id):
workflow_row = wf_repo.get(self.workflow_id, owner_id)
else:
workflow_row = wf_repo.get_by_legacy_id(self.workflow_id, owner_id)
return str(workflow_row["id"]) if workflow_row is not None else None
def _precreate_workflow_run(
self,
workflow_owner_id: Optional[str],
run_user_id: Optional[str],
query: str,
) -> Optional[str]:
"""Insert the run row up front so run-scoped artifacts are authz-reachable mid-run.
The workflow row is resolved against its *owner*; the run is owned by the
*runner* so artifact access (``run.user_id == caller``) tracks the caller.
"""
if not self._engine or not self.workflow_id or not workflow_owner_id or not run_user_id:
return None
try:
with db_session() as conn:
pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, workflow_owner_id)
if pg_workflow_id is None:
return None
WorkflowRunsRepository(conn).create(
pg_workflow_id,
run_user_id,
ExecutionStatus.RUNNING.value,
run_id=self._engine.workflow_run_id,
inputs={"query": query},
started_at=datetime.now(timezone.utc),
)
return pg_workflow_id
except Exception as e:
logger.error(f"Failed to pre-create workflow run: {e}")
return None
def _bridge_attachments(
self, run_user_id: Optional[str], *, persisted: bool
) -> Tuple[List[Dict[str, Any]], List[str]]:
"""Stage uploaded attachments as run-scoped artifacts the nodes can read.
Bytes are read server-side from each attachment's ``upload_path`` (bounded
by ``ARTIFACT_MAX_BYTES``, handle always closed) and re-persisted through
``persist_new_artifact`` (size/sha256/storage key all derived server-side);
only the resulting references enter the run state. Artifacts are owned by
the *runner* (the uploader), not the workflow owner.
Returns the bridged references and a list of user-facing notices for
attachments that were dropped (oversize / unreadable / unstorable).
``QuotaExceeded`` is NOT swallowed: it propagates to the caller so the run
fails cleanly instead of running with silently-missing documents.
"""
if not self._engine or not self.attachments or not run_user_id:
return [], []
# Without a persisted run row the artifacts would be orphaned (no authz
# parent), so skip the bridge for unowned/draft ids.
if not persisted:
return [], []
from application.sandbox.artifacts_capture import persist_new_artifact
from application.storage.storage_creator import StorageCreator
storage = StorageCreator.get_storage()
max_bytes = int(getattr(settings, "ARTIFACT_MAX_BYTES", 0) or 0)
dropped: List[str] = []
if len(self.attachments) > _MAX_INPUT_DOCUMENTS:
over = len(self.attachments) - _MAX_INPUT_DOCUMENTS
logger.warning(
"Workflow run input documents exceed cap (%d); dropping %d attachment(s)",
_MAX_INPUT_DOCUMENTS,
over,
)
dropped.append(
f"Only the first {_MAX_INPUT_DOCUMENTS} input document(s) were used; "
f"{over} additional attachment(s) were dropped."
)
refs: List[Dict[str, Any]] = []
for attachment in self.attachments[:_MAX_INPUT_DOCUMENTS]:
upload_path = attachment.get("upload_path") or attachment.get("path")
if not upload_path:
continue
filename = attachment.get("filename") or "attachment"
mime_type = attachment.get("mime_type") or "application/octet-stream"
# Reject oversize attachments via the authoritative ``size`` column
# BEFORE buffering the bytes into worker memory (a memory-DoS guard);
# the bounded read below backstops a missing/lying ``size``.
declared_size = attachment.get("size")
if max_bytes and isinstance(declared_size, (int, float)) and declared_size > max_bytes:
dropped.append(f'Document "{filename}" exceeds the artifact size limit and was skipped.')
continue
data = self._read_attachment_bytes(storage, upload_path, max_bytes)
if data is None:
dropped.append(f'Document "{filename}" could not be read and was skipped.')
continue
if max_bytes and len(data) > max_bytes:
dropped.append(f'Document "{filename}" exceeds the artifact size limit and was skipped.')
continue
# QuotaExceeded propagates (fatal); persist_new_artifact returns None on
# any other error, which we report as a per-attachment drop.
ref = persist_new_artifact(
user_id=run_user_id,
kind="file",
data=data,
filename=filename,
mime_type=mime_type,
title=filename,
workflow_run_id=self._engine.workflow_run_id,
)
if ref is None:
dropped.append(f'Document "{filename}" could not be stored and was skipped.')
continue
refs.append(
{
"artifact_id": ref["artifact_id"],
"ref": ref.get("ref"),
"filename": ref["filename"],
"mime_type": ref["mime_type"],
}
)
return refs, dropped
@staticmethod
def _read_attachment_bytes(storage: Any, upload_path: str, max_bytes: int) -> Optional[bytes]:
"""Read an attachment with a bounded read and a guaranteed handle close; None on failure."""
try:
file_obj = storage.get_file(upload_path)
except Exception as exc:
logger.error("Failed to open attachment for workflow run: %s", type(exc).__name__)
return None
try:
return file_obj.read(max_bytes + 1) if max_bytes else file_obj.read()
except Exception as exc:
logger.error("Failed to read attachment for workflow run: %s", type(exc).__name__)
return None
finally:
close = getattr(file_obj, "close", None)
if callable(close):
close()
def _finalize_workflow_run(
self,
workflow_owner_id: Optional[str],
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=owner_id,
status=self._determine_run_status(),
user=run_user_id,
status=status,
inputs={"query": query},
outputs=self._serialize_state(self._engine.state),
steps=self._engine.get_execution_summary(),
created_at=datetime.now(timezone.utc),
completed_at=datetime.now(timezone.utc),
)
steps_json = [step.model_dump(mode="json") for step in run.steps]
if not self.workflow_id or not owner_id:
if not self.workflow_id or not workflow_owner_id or not run_user_id:
return
with db_session() as conn:
wf_repo = WorkflowsRepository(conn)
if looks_like_uuid(self.workflow_id):
workflow_row = wf_repo.get(self.workflow_id, owner_id)
else:
workflow_row = wf_repo.get_by_legacy_id(
self.workflow_id, owner_id,
if pg_workflow_id is None:
pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, workflow_owner_id)
if pg_workflow_id is None:
return
runs_repo = WorkflowRunsRepository(conn)
updated = False
if self._run_persisted:
updated = runs_repo.finalize(
self._engine.workflow_run_id,
run_user_id,
run.status.value,
result=run.outputs,
steps=steps_json,
ended_at=run.completed_at,
)
if not updated:
logger.warning(
"Workflow run %s finalize matched no row; "
"recovering via insert so terminal data is not lost",
self._engine.workflow_run_id,
)
if not self._run_persisted or not updated:
runs_repo.create(
pg_workflow_id,
run_user_id,
run.status.value,
run_id=self._engine.workflow_run_id,
inputs=run.inputs,
result=run.outputs,
steps=steps_json,
started_at=run.created_at,
ended_at=run.completed_at,
)
if workflow_row is None:
return
WorkflowRunsRepository(conn).create(
str(workflow_row["id"]),
owner_id,
run.status.value,
inputs=run.inputs,
result=run.outputs,
steps=[step.model_dump(mode="json") for step in run.steps],
started_at=run.created_at,
ended_at=run.completed_at,
)
except Exception as e:
logger.error(f"Failed to save workflow run: {e}")
def _determine_run_status(self) -> ExecutionStatus:
# A fatal input-document bridge failure (quota) means the engine never ran;
# the run is FAILED even though there is no per-node failure log entry.
if self._bridge_error is not None:
return ExecutionStatus.FAILED
if not self._engine or not self._engine.execution_log:
return ExecutionStatus.COMPLETED
for log in self._engine.execution_log:
@@ -239,10 +484,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):
+10 -33
View File
@@ -1,6 +1,6 @@
"""Workflow Node Agents - defines specialized agents for workflow nodes."""
from typing import Any, Dict, List, Optional, Type
from typing import Dict, List, Optional, Type
from application.agents.agentic_agent import AgenticAgent
from application.agents.base import BaseAgent
@@ -9,34 +9,6 @@ from application.agents.research_agent import ResearchAgent
from application.agents.workflows.schemas import AgentType
class ToolFilterMixin:
"""Mixin that filters fetched tools to only those specified in tool_ids."""
_allowed_tool_ids: List[str]
def _get_user_tools(self, user: str = "local") -> Dict[str, Dict[str, Any]]:
all_tools = super()._get_user_tools(user)
if not self._allowed_tool_ids:
return {}
filtered_tools = {
tool_id: tool
for tool_id, tool in all_tools.items()
if str(tool.get("_id", "")) in self._allowed_tool_ids
}
return filtered_tools
def _get_tools(self, api_key: str = None) -> Dict[str, Dict[str, Any]]:
all_tools = super()._get_tools(api_key)
if not self._allowed_tool_ids:
return {}
filtered_tools = {
tool_id: tool
for tool_id, tool in all_tools.items()
if str(tool.get("_id", "")) in self._allowed_tool_ids
}
return filtered_tools
class _WorkflowNodeMixin:
"""Common __init__ for all workflow node agents."""
@@ -56,18 +28,23 @@ class _WorkflowNodeMixin:
api_key=api_key,
**kwargs,
)
self._allowed_tool_ids = tool_ids or []
# Scope the executor to exactly the node's configured tools. Agents
# fetch their toolset via ``tool_executor.get_tools()``, so the scope
# must live on the executor — it resolves builtin synthetic ids
# (Artifact / Code Executor / Read Document) and ``user_tools`` rows
# alike, and an empty list means the node's LLM gets no tools.
self.tool_executor.allowed_tool_ids = [str(t) for t in (tool_ids or [])]
class WorkflowNodeClassicAgent(ToolFilterMixin, _WorkflowNodeMixin, ClassicAgent):
class WorkflowNodeClassicAgent(_WorkflowNodeMixin, ClassicAgent):
pass
class WorkflowNodeAgenticAgent(ToolFilterMixin, _WorkflowNodeMixin, AgenticAgent):
class WorkflowNodeAgenticAgent(_WorkflowNodeMixin, AgenticAgent):
pass
class WorkflowNodeResearchAgent(ToolFilterMixin, _WorkflowNodeMixin, ResearchAgent):
class WorkflowNodeResearchAgent(_WorkflowNodeMixin, ResearchAgent):
pass
+26 -1
View File
@@ -12,6 +12,7 @@ class NodeType(str, Enum):
NOTE = "note"
STATE = "state"
CONDITION = "condition"
CODE = "code"
class AgentType(str, Enum):
@@ -48,6 +49,24 @@ class AgentNodeConfig(BaseModel):
retriever: str = ""
model_id: Optional[str] = None
json_schema: Optional[Dict[str, Any]] = None
# Run-scoped documents fed to this node's LLM. Entries are state-var names
# holding artifact refs (single dict or a list of dicts), raw artifact ids,
# short refs (``A1``), or the ``"*"``/``"input_documents"`` token meaning
# "every ref in ``state['input_documents']``".
input_documents: List[str] = Field(default_factory=list)
# How selected documents reach the model: ``auto`` (native when the model
# accepts the mime, else extract to text), ``native`` (force native; raise
# on an unsupported mime), or ``extract`` (always inline extracted text).
file_passing: Literal["auto", "native", "extract"] = "auto"
class CodeNodeConfig(BaseModel):
model_config = ConfigDict(extra="allow")
code: str = ""
inputs: List[str] = Field(default_factory=list)
output_variable: Optional[str] = None
timeout: Optional[int] = None
json_schema: Optional[Dict[str, Any]] = None
class ConditionCase(BaseModel):
@@ -146,8 +165,14 @@ class NodeExecutionLog(BaseModel):
status: ExecutionStatus
started_at: datetime
completed_at: Optional[datetime] = None
duration_ms: Optional[int] = None
error: Optional[str] = None
state_snapshot: Dict[str, Any] = Field(default_factory=dict)
# The node's state DELTA (keys it added or changed), not the full state:
# point-in-time state is the merge of deltas up to this step. Runs
# persisted before the rename carry this as ``state_snapshot``.
state_delta: Dict[str, Any] = Field(default_factory=dict)
# Compact per-node tool-call summary: [{tool_name, action_name, status}].
tool_calls: List[Dict[str, Any]] = Field(default_factory=list)
class WorkflowRunCreate(BaseModel):
+733 -18
View File
@@ -1,5 +1,7 @@
import json
import logging
import re
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, Generator, List, Optional, TYPE_CHECKING
@@ -8,6 +10,7 @@ from application.agents.workflows.node_agent import WorkflowNodeAgentFactory
from application.agents.workflows.schemas import (
AgentNodeConfig,
AgentType,
CodeNodeConfig,
ConditionNodeConfig,
ExecutionStatus,
NodeExecutionLog,
@@ -34,15 +37,45 @@ logger = logging.getLogger(__name__)
StateValue = Any
WorkflowState = Dict[str, StateValue]
TEMPLATE_RESERVED_NAMESPACES = {"agent", "system", "source", "tools", "passthrough"}
TEMPLATE_RESERVED_NAMESPACES = {"agent", "artifacts", "system", "source", "tools", "passthrough"}
# Run ids become a sandbox-session / kernel-workspace path component; the gateway
# only accepts [A-Za-z0-9_-]+, so any disallowed character is stripped before binding.
_SESSION_ID_RE = re.compile(r"[^A-Za-z0-9_-]+")
# State keys never staged into a code node's ``state.json``. ``chat_history`` is
# the caller's conversation and must not be reachable by owner-authored,
# egress-open sandbox code (see ``_json_safe_state``).
_CODE_STATE_EXCLUDED_KEYS = frozenset({"chat_history"})
# Sentinel id stamped on the synthetic text attachment that flags documents a
# node skipped past the per-node blocking-extract cap (callers/tests match on it).
_EXTRACT_TRUNCATION_ID = "workflow-extract-truncated"
class WorkflowEngine:
MAX_EXECUTION_STEPS = 50
def __init__(self, graph: WorkflowGraph, agent: "BaseAgent"):
def __init__(
self,
graph: WorkflowGraph,
agent: "BaseAgent",
workflow_run_id: Optional[str] = None,
):
"""Bind the engine to a graph + agent; mint a run id for run-scoped sandbox/artifacts."""
self.graph = graph
self.agent = agent
# The run id scopes the code-node sandbox session and every produced
# artifact's parent; mint one when the caller has not supplied a
# pre-created ``workflow_runs`` id so the engine is self-contained.
self.workflow_run_id: str = workflow_run_id or str(uuid.uuid4())
# Per-node tool-call summary collected by the node executors and folded
# into the step log; reset before each node runs.
self._last_node_tool_calls: List[Dict[str, Any]] = []
# False when no ``workflow_runs`` row backs ``workflow_run_id`` (unsaved /
# embedded draft), so run-scoped artifacts must NOT be persisted as orphans.
# Defaults True; WorkflowAgent flips it off on the draft path.
self.run_persisted: bool = True
self.state: WorkflowState = {}
self.execution_log: List[Dict[str, Any]] = []
self._condition_result: Optional[str] = None
@@ -51,15 +84,46 @@ class WorkflowEngine:
def execute(
self, initial_inputs: WorkflowState, query: str
) -> Generator[Dict[str, str], None, None]:
"""Run the workflow graph, closing the run-scoped sandbox session once when the run ends."""
try:
yield from self._run_graph(initial_inputs, query)
finally:
# The sandbox session is keyed by the run id and shared by every code
# node and agent-node tool in this run, so it is torn down exactly once
# here rather than per node. peek_manager() never builds the manager, so
# a run that never opened a session closes nothing.
from application.sandbox.sandbox_creator import SandboxCreator
mgr = SandboxCreator.peek_manager()
if mgr is not None:
try:
mgr.close(self._session_id())
except Exception:
logger.exception("Workflow run failed to close its sandbox session")
def _run_graph(
self, initial_inputs: WorkflowState, query: str
) -> Generator[Dict[str, str], None, None]:
self._initialize_state(initial_inputs, query)
# Surface the run id up front so the client can list this run's
# artifacts (GET /api/artifacts?workflow_run_id=) once it has been
# persisted; the same id parents every artifact produced by code nodes.
yield {"type": "workflow_run", "workflow_run_id": self.workflow_run_id}
start_node = self.graph.get_start_node()
if not start_node:
yield {"type": "error", "error": "No start node found in workflow."}
return
current_node_id: Optional[str] = start_node.id
steps = 0
# Snapshots are stored as per-node DELTAS: full state copied into
# every step grows the run row O(n^2) and repeats every upstream
# output verbatim. Point-in-time state = merge of deltas up to a step.
# The empty baseline attributes the pre-loop initialization (query,
# input documents) to the first step so steps alone reconstruct state.
pre_state: Dict[str, Any] = {}
while current_node_id and steps < self.MAX_EXECUTION_STEPS:
node = self.graph.get_node_by_id(current_node_id)
@@ -67,6 +131,7 @@ class WorkflowEngine:
yield {"type": "error", "error": f"Node {current_node_id} not found."}
break
log_entry = self._create_log_entry(node)
self._last_node_tool_calls = []
yield {
"type": "workflow_step",
@@ -79,7 +144,7 @@ class WorkflowEngine:
try:
yield from self._execute_node(node)
log_entry["status"] = ExecutionStatus.COMPLETED.value
log_entry["completed_at"] = datetime.now(timezone.utc)
self._finalize_log_entry(log_entry, pre_state)
output_key = f"node_{node.id}_output"
node_output = self.state.get(output_key)
@@ -90,15 +155,14 @@ class WorkflowEngine:
"node_type": node.type.value,
"node_title": node.title,
"status": "completed",
"state_snapshot": dict(self.state),
"state_delta": log_entry["state_delta"],
"output": node_output,
}
except Exception as e:
logger.error(f"Error executing node {node.id}: {e}", exc_info=True)
log_entry["status"] = ExecutionStatus.FAILED.value
log_entry["error"] = str(e)
log_entry["completed_at"] = datetime.now(timezone.utc)
log_entry["state_snapshot"] = dict(self.state)
self._finalize_log_entry(log_entry, pre_state)
self.execution_log.append(log_entry)
user_friendly_error = sanitize_api_error(e)
@@ -108,13 +172,13 @@ class WorkflowEngine:
"node_type": node.type.value,
"node_title": node.title,
"status": "failed",
"state_snapshot": dict(self.state),
"state_delta": log_entry["state_delta"],
"error": user_friendly_error,
}
yield {"type": "error", "error": user_friendly_error}
break
log_entry["state_snapshot"] = dict(self.state)
self.execution_log.append(log_entry)
pre_state = dict(self.state)
if node.type == NodeType.END:
break
@@ -142,9 +206,40 @@ class WorkflowEngine:
"completed_at": None,
"status": ExecutionStatus.RUNNING.value,
"error": None,
"state_snapshot": {},
"state_delta": {},
}
def _finalize_log_entry(self, log_entry: Dict[str, Any], pre_state: Dict[str, Any]) -> None:
"""Stamp completion time, duration, the node's state delta, and its tool-call summary."""
completed_at = datetime.now(timezone.utc)
log_entry["completed_at"] = completed_at
log_entry["duration_ms"] = int((completed_at - log_entry["started_at"]).total_seconds() * 1000)
log_entry["state_delta"] = self._state_delta(pre_state)
if self._last_node_tool_calls:
log_entry["tool_calls"] = list(self._last_node_tool_calls)
def _state_delta(self, previous: Dict[str, Any]) -> Dict[str, Any]:
"""Keys this node added or changed. Deleted keys are not tracked; nothing deletes state today."""
return {
key: value
for key, value in self.state.items()
if key not in previous or previous[key] != value
}
@staticmethod
def _summarize_tool_calls(node_agent: Any) -> List[Dict[str, Any]]:
"""Compact per-node tool-call summary for the run record (no arguments/results)."""
executor = getattr(node_agent, "tool_executor", None)
calls = getattr(executor, "tool_calls", None) or []
return [
{
"tool_name": call.get("tool_name"),
"action_name": call.get("action_name"),
"status": call.get("status", "completed"),
}
for call in calls
]
def _get_next_node_id(self, current_node_id: str) -> Optional[str]:
node = self.graph.get_node_by_id(current_node_id)
edges = self.graph.get_outgoing_edges(current_node_id)
@@ -170,6 +265,7 @@ class WorkflowEngine:
NodeType.START: self._execute_start_node,
NodeType.NOTE: self._execute_note_node,
NodeType.AGENT: self._execute_agent_node,
NodeType.CODE: self._execute_code_node,
NodeType.STATE: self._execute_state_node,
NodeType.CONDITION: self._execute_condition_node,
NodeType.END: self._execute_end_node,
@@ -227,10 +323,10 @@ class WorkflowEngine:
)
node_api_key = get_api_key_for_provider(node_llm_name) or self.agent.api_key
# Structured output gates on the model's registry capability flags;
# fetch them only when a node json_schema needs the check.
if node_json_schema and node_model_id:
model_capabilities = get_model_capabilities(
node_model_id, user_id=node_user_id
)
model_capabilities = get_model_capabilities(node_model_id, user_id=node_user_id)
if model_capabilities and not model_capabilities.get(
"supports_structured_output", False
):
@@ -238,6 +334,11 @@ class WorkflowEngine:
f'Model "{node_model_id}" does not support structured output for node "{node.title}"'
)
node_prompt = node_config.system_prompt
doc_manifest = self._node_document_manifest(node_config)
if doc_manifest:
node_prompt = f"{node_prompt}\n\n{doc_manifest}" if node_prompt else doc_manifest
factory_kwargs = {
"agent_type": node_config.agent_type,
"endpoint": self.agent.endpoint,
@@ -246,7 +347,7 @@ class WorkflowEngine:
"model_user_id": getattr(self.agent, "model_user_id", None),
"api_key": node_api_key,
"tool_ids": node_config.tools,
"prompt": node_config.system_prompt,
"prompt": node_prompt,
"chat_history": self.agent.chat_history,
"decoded_token": self.agent.decoded_token,
"json_schema": node_json_schema,
@@ -265,12 +366,50 @@ class WorkflowEngine:
}
node_agent = WorkflowNodeAgentFactory.create(**factory_kwargs)
# Attribute this node's tool calls to the run's message and, crucially,
# namespace their durability journal keys by it: node executors are built
# without a message_id, so providers that reuse deterministic call ids
# ("functions.create_artifact:0") would otherwise collide across runs on
# the tool_call_attempts primary key and drop the later journal rows.
node_executor = getattr(node_agent, "tool_executor", None)
if node_executor is not None:
node_executor.message_id = getattr(
getattr(self.agent, "tool_executor", None), "message_id", None
)
# Run-scope the node agent's tools so artifact_generator / code_executor
# address artifacts by this workflow run: a short ref (A1) created by one
# node resolves for edit_artifact in a later node within the same run. Only
# when a workflow_runs row backs the run -- otherwise an artifact parented
# to this run id would be an orphan (403 on get/download); left unset, the
# run-scoped tools persist under a conversation parent or cleanly error.
if self.run_persisted and getattr(node_agent, "tool_executor", None) is not None:
node_agent.tool_executor.workflow_run_id = self.workflow_run_id
# Decide native-eligibility from the SAME supported-types list the provider
# handler filters on at send time, so a mime is never sent native-but-empty
# and then silently dropped. Read post-construction: BaseAgent consumes
# ``self.attachments`` at gen time, so assigning here takes effect.
node_attachments = self._materialize_node_attachments(
node_config, node.title, self._agent_supported_attachment_types(node_agent)
)
if node_attachments:
node_agent.attachments = node_attachments
full_response_parts: List[str] = []
structured_response_parts: List[str] = []
has_structured_response = False
first_chunk = True
for event in node_agent.gen(formatted_prompt):
# A tool that pauses for approval makes the LLM handler yield
# ``tool_calls_pending`` and end. An ephemeral node agent has no resume path,
# so silently continuing would leave the node with empty output (or a confusing
# "Structured output was expected" when it has a json_schema). Fail visibly.
if event.get("type") == "tool_calls_pending":
raise ValueError(
f'Node "{node.title}" uses a tool that requires approval, which is not '
"supported inside a workflow. Disable require_approval for tools used in "
"workflow nodes."
)
if "answer" in event:
chunk = str(event["answer"])
full_response_parts.append(chunk)
@@ -286,6 +425,7 @@ class WorkflowEngine:
if node_config.stream_to_user:
self._has_streamed = True
self._last_node_tool_calls = self._summarize_tool_calls(node_agent)
full_response = "".join(full_response_parts).strip()
output_value: Any = full_response
if has_structured_response:
@@ -314,6 +454,527 @@ class WorkflowEngine:
if node_config.output_variable:
self.state[node_config.output_variable] = output_value
def _execute_code_node(
self, node: WorkflowNode
) -> Generator[Dict[str, str], None, None]:
"""Run code in the run-scoped sandbox, persist produced files, and write an artifact reference."""
from application.sandbox.artifacts_capture import capture_artifacts, snapshot_signatures
from application.sandbox.sandbox_creator import SandboxCreator
config = CodeNodeConfig(**node.config.get("config", node.config))
code = config.code or ""
if not code.strip():
raise ValueError(f'Code node "{node.title}" has no code to execute.')
# Code nodes are NEVER Jinja-rendered: state is untrusted (document-derived)
# so interpolating it into the program would be code injection. Prior state is
# passed as DATA via ``state.json`` (read below), never templated into code.
user_id = self._resolve_user_id()
if not user_id:
raise ValueError(f'Code node "{node.title}" requires an authenticated user.')
node_json_schema = self._normalize_node_json_schema(config.json_schema, node.title)
session_id = self._session_id()
timeout = self._resolve_code_timeout(config.timeout)
manager = SandboxCreator.get_manager()
# The session is keyed by the run id and shared by every code node and every
# agent-node tool in this run, so it is NOT closed here: closing per node
# would cold-drop later nodes' interpreter/filesystem state. The run session
# is torn down once in ``execute``'s finally when the whole run ends.
manager.open(session_id)
loaded = self._materialize_code_inputs(manager, session_id, config.inputs, user_id)
# Stage prior state as DATA the node code reads with
# ``json.load(open("state.json"))`` -- e.g. ``state["decision"]``. The
# file lands at the workspace root, which is the kernel cwd, so a
# relative open resolves it. State is never templated into the program.
state_json = json.dumps(self._json_safe_state(), default=str).encode("utf-8")
manager.put_file(session_id, "state.json", state_json)
pre_signatures = snapshot_signatures(manager, session_id)
result = manager.exec(session_id, code, timeout=timeout)
# Code nodes execute the sandbox directly (no tool_executor); record the
# run as a synthetic tool call so the step log shows it like agent nodes.
self._last_node_tool_calls = [
{
"tool_name": "code_executor",
"action_name": "run_code",
"status": "completed" if result.ok else "error",
}
]
if result.runtime_invalidated:
# A hard timeout destroyed the workspace runtime, so there is
# nothing reachable to capture and the timeout below stays primary.
artifacts = []
elif self.run_persisted:
artifacts = capture_artifacts(
manager,
session_id,
pre_signatures,
user_id=user_id,
workflow_run_id=self.workflow_run_id,
produced_by={"node_id": node.id, "node_type": NodeType.CODE.value},
)
else:
# No workflow_runs row backs this run (unsaved/embedded draft): an
# artifact parented to this run id would be an unreachable orphan
# (403 on get/download). Skip persistence; the sandbox still ran and
# ``_build_code_output`` handles the empty-artifacts case.
artifacts = []
if not result.ok:
error = (
f"{result.error_name}: {result.error_value}"
if result.error_name
else (result.error_value or "execution error")
)
raise ValueError(f'Code node "{node.title}" failed: {error}')
# The primary produced file becomes the node's pass-by-reference output;
# it is JSON primitives only ({artifact_id, version, mime_type, filename})
# so it survives the workflow_runs state-snapshot serialization. Bytes
# never enter state. A structured decision (optional json_schema) is
# parsed from stdout and validated through the existing jsonschema path.
output_value: Any = self._build_code_output(node, result, artifacts, loaded, node_json_schema)
default_output_key = f"node_{node.id}_output"
self.state[default_output_key] = output_value
if config.output_variable:
self.state[config.output_variable] = output_value
yield from ()
def _build_code_output(
self,
node: WorkflowNode,
result: Any,
artifacts: List[Dict[str, Any]],
inputs_loaded: List[str],
node_json_schema: Optional[Dict[str, Any]],
) -> Any:
"""Shape a code node's pass-by-reference output (artifact ref and/or validated decision)."""
if node_json_schema is not None:
parsed_success, decision = self._parse_structured_output(result.stdout or "")
if not parsed_success:
raise ValueError(
f'Code node "{node.title}" must print JSON matching its schema, '
"but stdout was not valid JSON"
)
self._validate_structured_output(node_json_schema, decision)
if artifacts:
# Carry the produced artifact reference alongside the decision so
# downstream nodes can branch on both.
if isinstance(decision, dict) and "artifacts" not in decision:
decision = {**decision, "artifacts": artifacts}
return decision
if artifacts:
return artifacts[0]
return {"artifacts": [], "status": "ok"}
def _materialize_code_inputs(
self, manager: Any, session_id: str, inputs: List[str], user_id: str
) -> List[str]:
"""Stage referenced input artifacts (run-scoped, never cross-tenant) into the workspace."""
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.core.settings import settings
from application.sandbox.artifacts_capture import unique_input_path
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
from application.storage.storage_creator import StorageCreator
from application.utils import safe_filename
loaded: List[str] = []
raw_ids = self._resolve_input_artifact_ids(inputs)
if not raw_ids:
return loaded
max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 0) or 0)
storage = StorageCreator.get_storage()
# Two inputs whose current versions share a filename would clobber each other at the
# same ``inputs/{name}`` path; track used paths and disambiguate deterministically.
used_paths: set = set()
for raw in raw_ids:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
# A short ref (A1/A2/...) resolves to an id within this run only;
# the resolved id is re-checked through the run-scoped gate so a ref
# can never reach another tenant.
artifact_id = resolve_artifact_id(repo, raw, workflow_run_id=self.workflow_run_id)
artifact = (
repo.get_artifact_in_parent(artifact_id, workflow_run_id=self.workflow_run_id)
if artifact_id is not None
else None
)
if artifact is None:
raise ValueError(f"input artifact {raw} not found in this run.")
version = repo.get_version(artifact_id, artifact["current_version"])
if not version or not version.get("storage_path"):
raise ValueError(f"input artifact {artifact_id} has no stored content.")
declared_size = version.get("size")
if max_bytes and isinstance(declared_size, (int, float)) and declared_size > max_bytes:
raise ValueError(
f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit."
)
filename = safe_filename(version.get("filename") or artifact_id)
file_obj = storage.get_file(version["storage_path"])
try:
data = file_obj.read(max_bytes + 1) if max_bytes else file_obj.read()
finally:
close = getattr(file_obj, "close", None)
if callable(close):
close()
if max_bytes and len(data) > max_bytes:
raise ValueError(
f"input artifact {artifact_id} exceeds the {max_bytes}-byte sandbox input limit."
)
rel_path = unique_input_path(f"inputs/{filename}", used_paths)
manager.put_file(session_id, rel_path, data)
loaded.append(rel_path)
return loaded
def _node_document_manifest(self, node_config: AgentNodeConfig) -> str:
"""One-line-per-document manifest of the node's selected input documents.
Appended to the node's system prompt so the model knows the concrete
refs (``A1``) and filenames it can pass to document/code tools —
without it, models guess placeholder names ("attached_file") or paste
file contents inline. Scoped to THIS node's ``input_documents``
selection, so it never widens per-node document access. Best-effort:
a resolution failure drops the manifest, never the node.
"""
from application.agents.tools.artifact_ref import make_ref, resolve_artifact_id
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
try:
raw_ids = self._resolve_input_artifact_ids(node_config.input_documents)
if not raw_ids:
return ""
lines: List[str] = []
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
for raw in raw_ids:
artifact_id = resolve_artifact_id(repo, raw, workflow_run_id=self.workflow_run_id)
artifact = (
repo.get_artifact_in_parent(artifact_id, workflow_run_id=self.workflow_run_id)
if artifact_id is not None
else None
)
if artifact is None:
continue
version = repo.get_version(artifact_id, artifact["current_version"]) or {}
ref_seq = (artifact.get("metadata") or {}).get("ref_seq")
# Legacy artifacts predate ref_seq; the full id works in every tool.
handle = make_ref(int(ref_seq)) if ref_seq else str(artifact_id)
filename = version.get("filename") or artifact.get("title") or str(artifact_id)
mime = version.get("mime_type") or "application/octet-stream"
lines.append(f"- {handle}: {filename} ({mime})")
except Exception:
logger.exception("Node document manifest failed; continuing without it")
return ""
if not lines:
return ""
return (
"## Input documents for this node\n"
"These files are staged for this node. Pass a ref (e.g. A1) or the "
"filename to a document or code tool to read the file:\n" + "\n".join(lines)
)
def _materialize_node_attachments(
self,
node_config: AgentNodeConfig,
node_title: str,
supported_types: List[str],
) -> List[Dict[str, Any]]:
"""Resolve a node's selected documents to native/extracted attachment dicts for its LLM."""
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.core.settings import settings
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
raw_ids = self._resolve_input_artifact_ids(node_config.input_documents)
if not raw_ids:
return []
supported = set(supported_types)
supports_images = any(t.startswith("image/") for t in supported)
max_files = int(getattr(settings, "WORKFLOW_NODE_NATIVE_MAX_FILES", 5))
extract_max = int(getattr(settings, "WORKFLOW_NODE_EXTRACT_MAX_FILES", 5))
max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024))
# One read-only connection for the whole batch; the resolved-version
# rows are collected, then storage reads happen outside the DB context.
resolved: List[tuple] = []
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
for raw in raw_ids:
# A short ref (A1/...) resolves to an id within this run only; the
# resolved id is re-checked through the run-scoped gate so a forged
# or cross-run ref can never reach another tenant's bytes.
artifact_id = resolve_artifact_id(repo, raw, workflow_run_id=self.workflow_run_id)
artifact = (
repo.get_artifact_in_parent(artifact_id, workflow_run_id=self.workflow_run_id)
if artifact_id is not None
else None
)
if artifact is None:
raise ValueError(
f'Document "{raw}" for node "{node_title}" was not found in this run.'
)
version = repo.get_version(artifact_id, artifact["current_version"])
if not version or not version.get("storage_path"):
raise ValueError(f"input document {artifact_id} has no stored content.")
resolved.append((str(artifact_id), artifact, version))
attachments: List[Dict[str, Any]] = []
native_count = 0
extract_count = 0
dropped_for_cap = 0
for artifact_id, artifact, version in resolved:
storage_path = version["storage_path"]
mime_type = version.get("mime_type") or "application/octet-stream"
filename = version.get("filename") or artifact.get("title") or artifact_id
size = version.get("size")
if isinstance(size, int) and size > max_bytes:
logger.warning(
"Workflow node %s: document %s (%d bytes) exceeds the %d-byte cap; skipping",
node_title, artifact_id, size, max_bytes,
)
continue
native_ok = self._is_native_mime(mime_type, supported, supports_images)
policy = node_config.file_passing
if policy == "native" and not native_ok:
raise ValueError(
f'Model "{node_config.model_id or self.agent.model_id}" cannot read '
f'"{mime_type}" natively for node "{node_title}".'
)
go_native = policy == "native" or (policy == "auto" and native_ok)
if go_native and native_count >= max_files:
logger.warning(
"Workflow node %s: native file cap (%d) reached; extracting %s instead",
node_title, max_files, filename,
)
go_native = False
if go_native:
# No bytes copied: the provider reads them from storage via ``path``.
attachments.append({"id": artifact_id, "mime_type": mime_type, "path": storage_path})
native_count += 1
else:
# Inline-text mimes are read directly (cheap); other mimes route
# through the parsing worker -- a blocking, ~120s per-document call.
# Cap how many documents a single node sends down that path so a
# node referencing many non-native documents (e.g. the ``*`` token)
# can't serialize dozens of parses. Inline text is not capped.
needs_parse = not self._is_inline_text_mime(mime_type)
if needs_parse:
# Count (and gate on) the parse ATTEMPT, not the success: a
# timed-out/failed parse is the ~120s worst case we must bound,
# so it has to consume cap budget too. Otherwise a degraded
# parsing backend (every parse fails) never advances the count
# and the node keeps issuing blocking parses without limit.
if extract_count >= extract_max:
dropped_for_cap += 1
continue
extract_count += 1
content = self._extract_attachment_text(
artifact_id, storage_path, mime_type, filename, max_bytes
)
if content is None:
logger.warning(
"Workflow node %s: could not extract text from %s; skipping",
node_title, filename,
)
continue
# A non-native mime routes this through ``_append_unsupported_attachments``,
# which inlines ``content`` as text in the system prompt.
attachments.append(
{"id": artifact_id, "mime_type": "text/plain", "content": content}
)
if dropped_for_cap:
logger.warning(
"Workflow node %s: blocking-extract cap (%d) reached; %d document(s) omitted",
node_title, extract_max, dropped_for_cap,
)
attachments.append(
self._extract_truncation_note(node_title, extract_max, dropped_for_cap)
)
return attachments
@staticmethod
def _extract_truncation_note(node_title: str, cap: int, dropped: int) -> Dict[str, Any]:
"""Build a non-fatal text attachment flagging documents skipped past the per-node extract cap."""
content = (
f'[Notice] Only the first {cap} document(s) for node "{node_title}" were extracted '
f"to text; {dropped} additional document(s) were omitted to bound execution time. "
"Reference fewer documents, or use a model that reads them natively."
)
return {"id": _EXTRACT_TRUNCATION_ID, "mime_type": "text/plain", "content": content}
@staticmethod
def _agent_supported_attachment_types(node_agent: "BaseAgent") -> List[str]:
"""Return the provider's authoritative supported attachment mime types (handler's source)."""
llm = getattr(node_agent, "llm", None)
getter = getattr(llm, "get_supported_attachment_types", None)
if not callable(getter):
return []
types = getter()
return list(types) if isinstance(types, (list, tuple, set)) else []
@staticmethod
def _is_native_mime(mime_type: str, supported_types: set, supports_images: bool) -> bool:
"""A mime is native if the model accepts it, or it is a PDF a vision model renders to images."""
if mime_type in supported_types:
return True
return mime_type == "application/pdf" and supports_images
def _extract_attachment_text(
self, artifact_id: str, storage_path: str, mime_type: str, filename: str, max_bytes: int
) -> Optional[str]:
"""Get an attachment's text: inline already-text formats, else parse via the parsing worker; None on failure."""
from application.parser.document_reader import truncate_text_head_tail
from application.storage.storage_creator import StorageCreator
if self._is_inline_text_mime(mime_type):
try:
data = StorageCreator.get_storage().get_file(storage_path).read()
except Exception:
logger.exception("Workflow node: failed to read document bytes for extraction")
return None
# Defensive size gate: a NULL/missing version size skips the pre-read cap,
# so re-check the actual bytes before inlining.
if len(data) > max_bytes:
logger.warning(
"Workflow node: document at %s (%d bytes) exceeds the %d-byte cap; skipping",
storage_path, len(data), max_bytes,
)
return None
try:
text = data.decode("utf-8", errors="replace")
except Exception:
return None
# Bound the inlined text to a head+tail window so a large-but-under-cap
# text file can't blow the context (the parse branch is already bounded).
return truncate_text_head_tail(text)
# Non-text mimes parse via the dedicated parsing queue (works on any backend,
# no sandbox): the worker re-resolves the artifact run-scoped and reads its bytes.
return self._parse_document_text(artifact_id)
def _parse_document_text(self, artifact_id: str) -> Optional[str]:
"""Enqueue ``parse_document`` for this run and await the bounded markdown; None on failure."""
from celery.exceptions import TimeoutError as CeleryTimeoutError
from application.api.user.tasks import parse_document
from application.core.settings import settings
user_id = self._resolve_user_id()
if not user_id:
return None
options = {"output": "markdown", "include_tables": False, "persist": False}
queue = getattr(settings, "DOCUMENT_PARSE_QUEUE", "parsing")
timeout = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120))
try:
async_result = parse_document.apply_async(
args=[artifact_id, {"workflow_run_id": self.workflow_run_id}, user_id, options],
queue=queue,
)
# A workflow can run inside a Celery worker (scheduled runs / webhooks).
# In a prefork worker ``task_join_will_block()`` is process-wide, so the
# default ``disable_sync_subtasks=True`` makes ``get()`` raise RuntimeError
# ("Never call result.get() within a task!"). The dedicated parsing queue +
# separate workers already avoid the real self-deadlock, so opt out
# explicitly (mirrors application/agents/tools/read_document.py).
result = async_result.get(timeout=timeout, disable_sync_subtasks=False)
except (CeleryTimeoutError, TimeoutError):
logger.warning("Workflow node: document parse timed out for %s", artifact_id)
return None
except Exception:
logger.exception("Workflow node: document parse failed")
return None
if isinstance(result, dict) and result.get("status") == "ok":
content = result.get("content")
return content if isinstance(content, str) else None
return None
@staticmethod
def _is_inline_text_mime(mime_type: str) -> bool:
"""Already-text formats are inlined directly (no Docling round-trip)."""
if mime_type.startswith("text/"):
return True
return mime_type in ("application/json", "application/xml")
def _resolve_input_artifact_ids(self, inputs: List[str]) -> List[str]:
"""Resolve node ``inputs`` (refs/ids, ``*`` token, or state vars holding a ref or a list of refs)."""
resolved: List[str] = []
for raw in inputs or []:
# ``*`` / ``input_documents`` expands to every run input document.
if isinstance(raw, str) and raw.strip() in ("*", "input_documents"):
resolved.extend(self._input_document_ids())
continue
ref = self.state.get(raw) if isinstance(raw, str) else None
if isinstance(ref, dict) and ref.get("artifact_id"):
resolved.append(str(ref["artifact_id"]))
elif isinstance(ref, list):
# A state var holding a list of ref dicts (e.g. input_documents).
for item in ref:
if isinstance(item, dict) and item.get("artifact_id"):
resolved.append(str(item["artifact_id"]))
elif isinstance(raw, str) and raw.strip():
resolved.append(raw.strip())
# Dedup preserving order so ``["*", "A1"]`` / duplicate refs don't attach twice.
return list(dict.fromkeys(resolved))
def _input_document_ids(self) -> List[str]:
"""Return the artifact ids of every ref in ``state['input_documents']``."""
docs = self.state.get("input_documents")
if not isinstance(docs, list):
return []
return [str(d["artifact_id"]) for d in docs if isinstance(d, dict) and d.get("artifact_id")]
def _session_id(self) -> str:
"""Sanitize the run id into the sandbox-gateway charset for the session key."""
return _SESSION_ID_RE.sub("-", str(self.workflow_run_id)) or str(uuid.uuid4())
def _json_safe_state(self) -> Dict[str, Any]:
"""Project ``self.state`` to a JSON-safe dict (the code node reads it from state.json).
``chat_history`` is excluded: it is the *caller's* full conversation, and a
code node (authored by the workflow owner, who may differ from the runner
in a shared agent) has no legitimate need for it. Since sandbox egress is
open by design, staging it would let owner-authored code exfiltrate the
runner's conversation. Node outputs and ``query`` are still exposed.
"""
projection: Dict[str, Any] = {}
for key, value in self.state.items():
if not isinstance(key, str):
continue
normalized_key = key.strip()
if not normalized_key or normalized_key in _CODE_STATE_EXCLUDED_KEYS:
continue
projection[normalized_key] = value
return projection
def _resolve_user_id(self) -> Optional[str]:
"""Resolve the run's owner for artifact ownership/quota accounting."""
user_id = getattr(self.agent, "user", None)
if user_id:
return user_id
token = getattr(self.agent, "decoded_token", None)
if isinstance(token, dict):
return token.get("sub")
return None
def _resolve_code_timeout(self, requested: Optional[int]) -> float:
"""Return the stricter of the node's requested timeout and the sandbox cap."""
from application.core.settings import settings
cap = float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60))
if requested is None:
return cap
try:
parsed = int(requested)
except (TypeError, ValueError):
return cap
return float(min(parsed, cap)) if parsed > 0 else cap
def _execute_state_node(
self, node: WorkflowNode
) -> Generator[Dict[str, str], None, None]:
@@ -351,7 +1012,14 @@ class WorkflowEngine:
output_template = str(config.get("output_template", ""))
if output_template:
formatted_output = self._format_template(output_template)
# A prior streaming node's text otherwise runs straight into the
# end-node output ("...enterprise growth.Sales analysis complete");
# insert the same paragraph break streamed nodes use between each
# other so the segments stay readable.
if getattr(self, "_has_streamed", False):
yield {"answer": "\n\n"}
yield {"answer": formatted_output}
self._has_streamed = True
def _parse_structured_output(self, raw_response: str) -> tuple[bool, Optional[Any]]:
normalized_response = raw_response.strip()
@@ -361,10 +1029,43 @@ class WorkflowEngine:
try:
return True, json.loads(normalized_response)
except json.JSONDecodeError:
logger.warning(
"Workflow agent returned structured output that was not valid JSON"
)
return False, None
pass
# Some models wrap structured output in a ```json ... ``` fence or add
# prose around it; recover the JSON object/array before giving up so a
# well-formed-but-fenced response still validates.
candidate = self._strip_json_fence(normalized_response)
if candidate is not None:
try:
return True, json.loads(candidate)
except json.JSONDecodeError:
pass
logger.warning(
"Workflow agent returned structured output that was not valid JSON"
)
return False, None
@staticmethod
def _strip_json_fence(text: str) -> Optional[str]:
"""Extract the JSON payload from a fenced/prose-wrapped response, or None."""
fence = re.search(r"```(?:json)?\s*(.*?)\s*```", text, re.DOTALL)
if fence:
return fence.group(1).strip()
# Fall back to the outermost {...} or [...] span, choosing whichever
# bracket opens first. Fixing the order to "{" before "[" would slice a
# top-level array (``[{...},{...}]``) from its first "{" to its last "}",
# dropping the array framing (invalid JSON) or extracting an inner object
# that parses as silently-wrong structured data.
best: Optional[str] = None
best_start = -1
for open_ch, close_ch in (("{", "}"), ("[", "]")):
start = text.find(open_ch)
end = text.rfind(close_ch)
if start != -1 and end > start and (best_start == -1 or start < best_start):
best_start = start
best = text[start : end + 1]
return best
def _normalize_node_json_schema(
self, schema: Optional[Dict[str, Any]], node_title: str
@@ -425,6 +1126,8 @@ class WorkflowEngine:
docs=docs,
docs_together=docs_together,
tools_data=tools_data,
artifacts_data=self._collect_artifact_refs(),
artifact_parent={"workflow_run_id": self.workflow_run_id},
)
agent_context: Dict[str, Any] = {}
@@ -448,6 +1151,16 @@ class WorkflowEngine:
return context
def _collect_artifact_refs(self) -> Dict[str, Any]:
"""Collect state variables that hold artifact references, keyed by their state name."""
refs: Dict[str, Any] = {}
for key, value in self.state.items():
if not isinstance(key, str):
continue
if isinstance(value, dict) and value.get("artifact_id"):
refs[key] = value
return refs
def _get_source_template_data(self) -> tuple[Optional[List[Dict[str, Any]]], Optional[str]]:
docs = getattr(self.agent, "retrieved_docs", None)
if not isinstance(docs, list) or len(docs) == 0:
@@ -501,8 +1214,10 @@ class WorkflowEngine:
status=ExecutionStatus(log["status"]),
started_at=log["started_at"],
completed_at=log.get("completed_at"),
duration_ms=log.get("duration_ms"),
error=log.get("error"),
state_snapshot=log.get("state_snapshot", {}),
state_delta=log.get("state_delta", {}),
tool_calls=log.get("tool_calls", []),
)
for log in self.execution_log
]
@@ -0,0 +1,111 @@
"""0025 artifacts — append-only artifact identity + version store.
Adds the long-anticipated ``artifacts`` entity referenced by the schema header
in ``models.py`` but never built. Two tables:
* ``artifacts`` — one identity row per logical artifact. The stable ``id`` is the
handle passed around (chat/workflow ``state``/message bodies carry only this
reference, never bytes). ``current_version`` mirrors
``workflows.current_graph_version`` (atomic increment on each new version).
Authz is parent-derived: whoever can reach the ``conversation_id`` (chat) or
``workflow_run_id`` (run) can reach its artifacts, so a CHECK requires at least
one parent. ``user_id`` is kept for ownership/quota only; ``team_id`` is a
nullable forward-compat hook for the in-flight Teams work (no FK, matching how
``teams`` itself leaves owner/member columns FK-free).
* ``artifact_versions`` — append-only, never mutated. Each edit appends a row and
bumps ``artifacts.current_version``. ``UNIQUE(artifact_id, version)`` enforces
monotonic, gap-checkable versions. ``storage_path`` is the ``BaseStorage`` key
(NULL when the version is spec-only). Pass-by-reference: bytes live in storage,
only metadata + the key live here.
Revision ID: 0025_artifacts
Revises: 0024_wiki_pages_updated_via
"""
from typing import Sequence, Union
from alembic import op
revision: str = "0025_artifacts"
down_revision: Union[str, None] = "0024_wiki_pages_updated_via"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# --- artifacts -----------------------------------------------------------
op.execute(
"""
CREATE TABLE IF NOT EXISTS artifacts (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
user_id TEXT NOT NULL,
conversation_id UUID,
workflow_run_id UUID,
team_id UUID,
message_id UUID,
kind TEXT NOT NULL,
title TEXT,
metadata JSONB,
current_version INTEGER NOT NULL DEFAULT 1,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CONSTRAINT artifacts_parent_present_check
CHECK (conversation_id IS NOT NULL OR workflow_run_id IS NOT NULL)
);
"""
)
op.execute("CREATE INDEX IF NOT EXISTS artifacts_user_idx ON artifacts (user_id);")
op.execute(
"CREATE INDEX IF NOT EXISTS artifacts_conversation_idx "
"ON artifacts (conversation_id) WHERE conversation_id IS NOT NULL;"
)
op.execute(
"CREATE INDEX IF NOT EXISTS artifacts_workflow_run_idx "
"ON artifacts (workflow_run_id) WHERE workflow_run_id IS NOT NULL;"
)
# Reuse the shared set_updated_at() trigger fn defined in 0001.
op.execute("DROP TRIGGER IF EXISTS artifacts_set_updated_at ON artifacts;")
op.execute(
"""
CREATE TRIGGER artifacts_set_updated_at
BEFORE UPDATE ON artifacts
FOR EACH ROW EXECUTE FUNCTION set_updated_at();
"""
)
# --- artifact_versions ---------------------------------------------------
op.execute(
"""
CREATE TABLE IF NOT EXISTS artifact_versions (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
artifact_id UUID NOT NULL REFERENCES artifacts(id) ON DELETE CASCADE,
version INTEGER NOT NULL,
mime_type TEXT,
filename TEXT,
storage_path TEXT,
size BIGINT,
sha256 TEXT,
spec JSONB,
preview_text TEXT,
produced_by JSONB,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CONSTRAINT artifact_versions_artifact_version_uidx
UNIQUE (artifact_id, version)
);
"""
)
op.execute(
"CREATE INDEX IF NOT EXISTS artifact_versions_artifact_idx "
"ON artifact_versions (artifact_id);"
)
def downgrade() -> None:
op.execute("DROP INDEX IF EXISTS artifact_versions_artifact_idx;")
op.execute("DROP TABLE IF EXISTS artifact_versions;")
op.execute("DROP TRIGGER IF EXISTS artifacts_set_updated_at ON artifacts;")
op.execute("DROP INDEX IF EXISTS artifacts_workflow_run_idx;")
op.execute("DROP INDEX IF EXISTS artifacts_conversation_idx;")
op.execute("DROP INDEX IF EXISTS artifacts_user_idx;")
op.execute("DROP TABLE IF EXISTS artifacts;")
+27 -8
View File
@@ -237,6 +237,10 @@ class BaseAnswerResource:
Server-sent event strings
"""
response_full, thought, source_log_docs, tool_calls = "", "", [], []
# Set when a workflow agent run emits its ``workflow_run`` event; persisted
# onto the message metadata so the chat can render the run's produced
# artifacts on reload.
workflow_run_id: Optional[str] = None
is_structured = False
schema_info = None
structured_chunks = []
@@ -543,14 +547,29 @@ class BaseAnswerResource:
paused = True
yield _emit(line)
elif line.get("type") == "error":
yield _emit(
{
"type": "error",
"error": sanitize_api_error(
line.get("error", "An error occurred")
),
}
)
# An event flagged ``user_facing`` already carries a curated,
# actionable message (e.g. an artifact-quota notice). Passing it
# through sanitize_api_error would substring-match words like
# "quota" and rewrite it into a misleading rate-limit message, so
# emit it verbatim; sanitize only raw/technical errors.
error_text = line.get("error", "An error occurred")
if not line.get("user_facing"):
error_text = sanitize_api_error(error_text)
yield _emit({"type": "error", "error": error_text})
elif line.get("type") == "notice":
# Non-fatal, non-terminal notice (e.g. some workflow input
# documents were dropped). Forwarded verbatim so the client can
# surface it without failing the turn; never sanitized as an error.
yield _emit({"type": "notice", "notice": line.get("notice", "")})
elif line.get("type") == "workflow_run":
# Stash the run id in the message metadata so every
# persistence path (finalize / save / abort / error) records
# it — the chat renders the run's produced artifacts from it
# on reload. Still forwarded so the live client captures it.
workflow_run_id = line.get("workflow_run_id")
if workflow_run_id:
query_metadata["workflow_run_id"] = workflow_run_id
yield _emit(line)
else:
yield _emit(line)
if is_structured and structured_chunks:
@@ -154,7 +154,12 @@ class ConversationService:
completion = llm.gen(
model=getattr(llm, "model_id", None) or model_id,
messages=messages_summary,
max_tokens=500,
# Reasoning-capable default models spend the whole budget inside
# reasoning_content before emitting any title, so 500 came back
# empty (finish_reason=length). Give enough room to finish
# thinking and still produce the 3-word title; non-reasoning
# models stop far short of this cap.
max_tokens=2000,
)
if not completion or not completion.strip():
@@ -418,7 +423,11 @@ class ConversationService:
completion = llm.gen(
model=getattr(llm, "model_id", None) or title_inputs.get("model_id"),
messages=messages_summary,
max_tokens=500,
# Reasoning-capable default models spend the whole budget inside
# reasoning_content before emitting any title, so 500 came back empty
# (finish_reason=length). Give room to finish and still produce the
# 3-word title; non-reasoning models stop far short of this cap.
max_tokens=2000,
)
if not completion or not completion.strip():
completion = fallback_name or "New Conversation"
@@ -806,6 +806,13 @@ class StreamProcessor:
self.agent_config["workflow"] = self.data["workflow"]
if isinstance(self.decoded_token, dict):
self.agent_config["workflow_owner"] = self.decoded_token.get("sub")
# A saved workflow id alongside the embedded graph (builder
# Preview) lets the run persist a ``workflow_runs`` row so its
# artifacts are listable + authz'd; ownership is re-checked on
# save, so a forged id for another user's workflow never persists.
preview_workflow_id = self.data.get("workflow_id")
if preview_workflow_id:
self.agent_config["workflow_id"] = str(preview_workflow_id)
self.agent_config.update(
{
@@ -1125,6 +1132,32 @@ class StreamProcessor:
logger.warning(f"Failed to pre-fetch tools: {type(e).__name__}")
return None
def _enabled_tool_names(self) -> Optional[set]:
"""Resolve the tool names enabled for this turn, for ``tools.enabled`` gating.
Mirrors the executor the agent will use (same user/agent context), so an
agent yields its configured tools and an agentless chat yields user tools
plus defaults. Returns None on failure so the prompt gate fails open
(keeps the section) rather than hiding guidance when resolution breaks.
"""
try:
from application.agents.tool_executor import ToolExecutor
user = self.decoded_token.get("sub") if self.decoded_token else None
tool_executor = ToolExecutor(
user_api_key=self.agent_config.get("user_api_key"),
user=user,
decoded_token=self.decoded_token,
agent_id=self.agent_id,
)
client_tools = self.data.get("client_tools")
if client_tools:
tool_executor.client_tools = client_tools
return tool_executor.get_enabled_tool_names()
except Exception:
logger.warning("Failed to resolve enabled tool names for prompt gating")
return None
def _fetch_tool_data(
self,
tool_doc: Dict[str, Any],
@@ -1510,6 +1543,9 @@ class StreamProcessor:
docs=docs,
docs_together=docs_together,
tools_data=tools_data,
attachments=self.attachments,
enabled_tools=self._enabled_tool_names(),
artifact_parent={"conversation_id": self.conversation_id},
)
# Use the user_id that resolved the model so owner-scoped BYOM
@@ -1658,6 +1694,12 @@ class StreamProcessor:
agent_kwargs["workflow_id"] = workflow_config
elif isinstance(workflow_config, dict):
agent_kwargs["workflow"] = workflow_config
# Embedded-graph Preview run that names a saved workflow: run the
# canvas graph but persist the run under the saved id so artifacts
# parent to a real, ownership-checked ``workflow_runs`` row.
saved_workflow_id = self.agent_config.get("workflow_id")
if saved_workflow_id:
agent_kwargs["workflow_id"] = saved_workflow_id
workflow_owner = self.agent_config.get("workflow_owner")
if workflow_owner:
agent_kwargs["workflow_owner"] = workflow_owner
@@ -0,0 +1,5 @@
"""Artifacts module."""
from .routes import artifacts_ns
__all__ = ["artifacts_ns"]
+168
View File
@@ -0,0 +1,168 @@
"""Parent-derived authorization helpers for artifact access."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
from flask import request
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.repositories.conversations import ConversationsRepository
from application.storage.db.repositories.shared_conversations import (
SharedConversationsRepository,
)
from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository
from application.storage.db.session import db_readonly
@dataclass(frozen=True)
class Principal:
"""The resolved caller of an artifact route.
A decoded JWT is a full *owner* session (``agent_id`` is ``None``). An
``api_key`` is a low-trust, publicly-distributed *agent* credential -- it is
embedded client-side in the widget and accepted from the query string -- so
it resolves to the owning ``user_id`` but stays **agent-scoped**: it may reach
an artifact only when the request also carries that artifact's parent
``conversation_id`` (the unguessable per-visitor secret), confining it to a
single conversation rather than the agent's whole corpus, and may not perform
mutating operations.
"""
user_id: Optional[str] = None
agent_id: Optional[str] = None
@property
def is_agent_scoped(self) -> bool:
"""True for an api_key (agent) principal that must be confined to its agent."""
return self.agent_id is not None
def resolve_principal() -> Principal:
"""Resolve the caller to a :class:`Principal`.
Priority: a decoded JWT (owner session) first, then an ``api_key`` query/form
param (agent-scoped). An unresolvable caller yields an anonymous principal
(both fields ``None``) that only a valid ``share_token`` can authorize.
"""
decoded_token = getattr(request, "decoded_token", None)
if decoded_token:
return Principal(user_id=decoded_token.get("sub"))
api_key = request.args.get("api_key") or request.form.get("api_key")
if api_key:
with db_readonly() as conn:
agent = AgentsRepository(conn).find_by_key(api_key)
if agent and agent.get("user_id") and agent.get("id"):
return Principal(user_id=str(agent["user_id"]), agent_id=str(agent["id"]))
return Principal()
def _shared_row_for(conn, conversation_id, share_token):
"""Return the shared_conversations row iff share_token grants this conversation, else None."""
if not share_token:
return None
shared = SharedConversationsRepository(conn).find_by_uuid(share_token)
if shared and str(shared.get("conversation_id")) == str(conversation_id):
return shared
return None
def user_can_access_conversation(
conn, conversation_id: str, user_id: Optional[str], share_token: Optional[str]
) -> bool:
"""Allow if the caller owns/shares the conversation, or holds a valid share token.
Conversation-level gate only. Reuses ``ConversationsRepository.get`` (owner OR
``shared_with``) so artifact access tracks message access, and honours a valid
share token. A share-token caller reaches the conversation here but must still
be snapshot-scoped per-artifact by the caller (see ``authorize_artifact``): a
valid token does NOT imply access to every artifact in the conversation.
"""
if user_id:
if ConversationsRepository(conn).get(conversation_id, user_id) is not None:
return True
return _shared_row_for(conn, conversation_id, share_token) is not None
def authorize_artifact(conn, artifact: dict, principal: Principal) -> bool:
"""Authorize a READ of ``artifact`` for ``principal``; missing parent fails closed.
A low-trust agent api_key is confined to a single conversation it proves by
carrying that artifact's parent ``conversation_id`` query param (owner match +
matching conversation + agent scope) and never inherits share-link access. A
JWT owner or ``shared_with`` collaborator gets full access to every artifact of
the parent; a share-token holder is confined to the shared ``first_n_queries``
snapshot (an artifact whose ``message_id`` is outside it, or NULL, is denied).
"""
if principal.is_agent_scoped:
# An agent key is not the owner's session and is embedded in public widget
# JS, so require the artifact to belong to the owner AND the request to
# carry the parent conversation_id (the unguessable per-visitor secret):
# the key alone, without that id, cannot download a known artifact.
if str(artifact.get("user_id")) != str(principal.user_id):
return False
req_conv = request.args.get("conversation_id")
parent = artifact.get("conversation_id")
if not req_conv or parent is None or str(parent) != str(req_conv):
return False
return ArtifactsRepository(conn).artifact_in_agent_scope(
str(artifact.get("id")), str(principal.agent_id)
)
conversation_id = artifact.get("conversation_id")
workflow_run_id = artifact.get("workflow_run_id")
share_token = request.args.get("share_token")
if conversation_id is not None:
# Owner or shared_with collaborator: full access to every artifact.
if principal.user_id and ConversationsRepository(conn).get(
str(conversation_id), principal.user_id
) is not None:
return True
# Share-token holder: only artifacts inside the shared first_n_queries snapshot.
shared = _shared_row_for(conn, str(conversation_id), share_token)
if shared is None:
return False
message_id = artifact.get("message_id")
if not message_id:
return False
first_n = int(shared.get("first_n_queries") or 0)
return ConversationsRepository(conn).message_in_first_n(
str(conversation_id), str(message_id), first_n
)
if workflow_run_id is not None:
if not principal.user_id:
return False
run = WorkflowRunsRepository(conn).get(str(workflow_run_id))
return run is not None and run.get("user_id") == principal.user_id
# No parent row reachable -> deny (e.g. deleted conversation/run).
return False
def authorize_artifact_write(conn, artifact: dict, principal: Principal) -> bool:
"""Authorize a *mutating* artifact operation (e.g. delete / restore).
Stricter than :func:`authorize_artifact`: a write requires an authenticated
*owner* of the parent. Low-trust agent api_keys, share links, and
``shared_with`` collaborators inherit read/download access only, so an
agent-scoped, anonymous, or share-token-only caller is denied -- they can read
an artifact but never mutate it.
"""
if principal.is_agent_scoped or not principal.user_id:
return False
conversation_id = artifact.get("conversation_id")
workflow_run_id = artifact.get("workflow_run_id")
if conversation_id is not None:
return (
ConversationsRepository(conn).get_owned(str(conversation_id), principal.user_id)
is not None
)
if workflow_run_id is not None:
run = WorkflowRunsRepository(conn).get(str(workflow_run_id))
return run is not None and run.get("user_id") == principal.user_id
# No parent row reachable -> deny (e.g. deleted conversation/run).
return False
+402
View File
@@ -0,0 +1,402 @@
"""Artifact metadata and download routes (parent-derived authz)."""
from __future__ import annotations
import re
from typing import Optional
from flask import (
Response,
current_app,
jsonify,
make_response,
redirect,
request,
stream_with_context,
)
from flask_restx import Namespace, Resource
from application.api import api
from application.api.user.artifacts.authz import (
_shared_row_for,
authorize_artifact,
authorize_artifact_write,
resolve_principal,
user_can_access_conversation,
)
from application.core.settings import settings
from application.storage.db.base_repository import looks_like_uuid
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.repositories.conversations import ConversationsRepository
from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository
from application.storage.db.session import db_readonly, db_session
from application.storage.storage_creator import StorageCreator
artifacts_ns = Namespace("artifacts", description="Artifact operations", path="/api")
# Presigned-URL TTL for private S3 artifact downloads (seconds).
_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:
return fallback
cleaned = re.sub(r'[\r\n"]', "", str(filename)).strip()
return cleaned or fallback
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),
"kind": artifact.get("kind"),
"title": artifact.get("title"),
"metadata": artifact.get("metadata"),
"current_version": artifact.get("current_version"),
"created_at": _iso(artifact.get("created_at")),
"updated_at": _iso(artifact.get("updated_at")),
}
def _version_summary(version: dict, *, include_spec: bool = False) -> dict:
"""Project a version row to its API metadata shape (storage_path withheld)."""
out = {
"version": version.get("version"),
"mime_type": version.get("mime_type"),
"filename": version.get("filename"),
"size": version.get("size"),
"sha256": version.get("sha256"),
"preview_text": version.get("preview_text"),
"produced_by": version.get("produced_by"),
"created_at": _iso(version.get("created_at")),
}
if include_spec:
out["spec"] = version.get("spec")
return out
def _iso(value):
"""ISO-format a datetime, passing through other values unchanged."""
return value.isoformat() if hasattr(value, "isoformat") else value
@artifacts_ns.route("/artifacts")
class ListArtifacts(Resource):
@api.doc(description="List artifacts for a conversation, workflow run, or the caller")
def get(self):
principal = resolve_principal()
user_id = principal.user_id
conversation_id = request.args.get("conversation_id")
workflow_run_id = request.args.get("workflow_run_id")
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)
# 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)
if workflow_run_id and not looks_like_uuid(workflow_run_id):
return make_response(jsonify({"success": False, "message": "Invalid workflow_run_id"}), 400)
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
if principal.is_agent_scoped:
# A low-trust agent api_key is embedded in public widget JS and
# shared by every anonymous visitor, so it may NOT enumerate the
# agent's artifacts. The unguessable conversation_id (UUIDv4) is
# the only per-visitor secret: require it as a bearer capability
# 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)
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)
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)
rows = repo.list_artifacts(conversation_id=conversation_id)
# Owner / shared_with collaborator sees every artifact; a
# share-token caller is confined to the shared first_n_queries
# snapshot (drop artifacts whose message is outside it or NULL).
conv_repo = ConversationsRepository(conn)
is_owner = bool(user_id and conv_repo.get(conversation_id, user_id) is not None)
if not is_owner:
shared = _shared_row_for(conn, conversation_id, share_token)
first_n = int(shared.get("first_n_queries") or 0) if shared else 0
snapshot_ids = conv_repo.first_n_message_ids(conversation_id, first_n)
rows = [
r
for r in rows
if r.get("message_id") is not None and str(r.get("message_id")) in snapshot_ids
]
elif workflow_run_id:
run = WorkflowRunsRepository(conn).get(workflow_run_id)
if run is None or run.get("user_id") != user_id:
return make_response(jsonify({"success": False, "message": "Forbidden"}), 403)
rows = repo.list_artifacts(workflow_run_id=workflow_run_id)
else:
if not user_id:
return make_response(
jsonify({"success": False, "message": "Authentication required"}),
401,
)
rows = repo.list_artifacts(user_id=user_id)
return make_response(
jsonify({"success": True, "artifacts": [_artifact_summary(r) for r in rows]}),
200,
)
except Exception as err:
current_app.logger.error(f"Error listing artifacts: {err}", exc_info=True)
return make_response(jsonify({"success": False}), 400)
@artifacts_ns.route("/artifacts/<artifact_id>")
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)
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)
if not authorize_artifact(conn, artifact, principal):
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"))
payload = _artifact_summary(artifact)
payload["versions"] = [_version_summary(v) for v in versions]
payload["spec"] = current.get("spec") if current else None
return make_response(jsonify({"success": True, "artifact": payload}), 200)
except Exception as err:
current_app.logger.error(f"Error retrieving artifact: {err}", exc_info=True)
return make_response(jsonify({"success": False}), 400)
@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)
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)
# 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)
storage_paths = repo.delete_artifact(artifact_id)
# Reap the bytes best-effort AFTER the row delete commits.
_reap_storage(storage_paths)
return make_response(jsonify({"success": True}), 200)
except Exception as err:
current_app.logger.error(f"Error deleting artifact: {err}", exc_info=True)
return make_response(jsonify({"success": False}), 500)
def _reap_storage(paths: list) -> None:
"""Best-effort delete artifact bytes after the DB rows are gone; never raises."""
if not paths:
return
try:
storage = StorageCreator.get_storage()
except Exception:
current_app.logger.warning("artifact delete: storage unavailable", exc_info=True)
return
for path in paths:
try:
storage.delete_file(path)
except Exception:
current_app.logger.warning("artifact delete: failed to delete bytes %s", path, exc_info=True)
@artifacts_ns.route("/artifacts/<artifact_id>/versions/<int:version>")
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)
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)
if not authorize_artifact(conn, artifact, principal):
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": 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)
return make_response(jsonify({"success": False}), 400)
@artifacts_ns.route("/artifacts/<artifact_id>/download")
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)
principal = resolve_principal()
version_arg = request.args.get("version")
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)
if not authorize_artifact(conn, artifact, principal):
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)
version_row = repo.get_version(artifact_id, version)
if version_row is None:
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)
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()
# With URL_STRATEGY=="s3" the contract is to hand back a presigned
# URL. If the active backend can't mint one, that's a config error:
# surface a 500 rather than silently proxying bytes from a backend
# 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)
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"}),
500,
)
# A 302 to a cross-origin S3 URL can't be read by the app's authed
# fetch (the bucket has no CORS grant for the app origin). When the
# client opts in via ?disposition=url (or Accept: application/json)
# hand the presigned URL back as JSON so it can navigate to it
# top-level (no CORS). Default stays a 302 so nothing else breaks.
if request.args.get("disposition") == "url" or (
"application/json" in request.headers.get("Accept", "")
):
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
# worker memory (artifacts can be many MB); close the handle when done.
file_obj = storage.get_file(storage_path)
def _stream():
try:
for chunk in iter(lambda: file_obj.read(65536), b""):
yield chunk
finally:
close = getattr(file_obj, "close", None)
if callable(close):
close()
return Response(
stream_with_context(_stream()),
mimetype=mime_type,
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
)
except FileNotFoundError:
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)
return make_response(jsonify({"success": False}), 400)
@artifacts_ns.route("/artifacts/<artifact_id>/restore")
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)
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)
try:
target_version = int(target_version)
except (ValueError, TypeError):
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)
# 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)
source = repo.get_version(artifact_id, target_version)
if source is None:
return make_response(jsonify({"success": False, "message": "Version not found"}), 404)
new_version = repo.append_version(
artifact_id,
mime_type=source.get("mime_type"),
filename=source.get("filename"),
storage_path=source.get("storage_path"),
size=source.get("size"),
sha256=source.get("sha256"),
spec=source.get("spec"),
preview_text=source.get("preview_text"),
produced_by=source.get("produced_by"),
)
return make_response(
jsonify({"success": True, "version": _version_summary(new_version, include_spec=True)}),
200,
)
except Exception as err:
current_app.logger.error(f"Error restoring artifact: {err}", exc_info=True)
return make_response(jsonify({"success": False}), 400)
+21 -5
View File
@@ -47,8 +47,14 @@ def _resolve_authenticated_user():
decoded_token = getattr(request, "decoded_token", None)
api_key = request.form.get("api_key") or request.args.get("api_key")
# Return the RAW user identity (not safe_filename'd): the attachment row's
# ``user_id`` is an identity key that must match the raw ``sub`` used by
# /stream (initial_user_id) and the artifact authz gate. safe_filename is
# applied only to the storage path component at write time. (An email-style
# sub like ``a@b.com`` was previously sanitized to ``abcom``, so the
# uploaded attachment became unreadable by its own owner on /stream.)
if decoded_token:
return safe_filename(decoded_token.get("sub"))
return decoded_token.get("sub")
if api_key:
with db_readonly() as conn:
@@ -57,7 +63,7 @@ def _resolve_authenticated_user():
return make_response(
jsonify({"success": False, "message": "Invalid API key"}), 401
)
return safe_filename(agent.get("user_id"))
return agent.get("user_id")
return None
@@ -162,7 +168,12 @@ class StoreAttachment(Resource):
attachment_id = uuid.uuid4()
original_filename = safe_filename(os.path.basename(file.filename))
_enforce_uploaded_audio_size_limit(file, original_filename)
relative_path = f"{settings.UPLOAD_FOLDER}/{user}/attachments/{str(attachment_id)}/{original_filename}"
# safe_filename only the path component; the DB user_id stays raw.
path_user = safe_filename(user)
relative_path = (
f"{settings.UPLOAD_FOLDER}/{path_user}/attachments/"
f"{str(attachment_id)}/{original_filename}"
)
metadata = storage.save_file(file, relative_path)
file_info = {
@@ -422,7 +433,10 @@ class LiveSpeechToTextChunk(Resource):
404,
)
if safe_filename(str(session_state.get("user", ""))) != auth_user:
# The stored ``user`` is the RAW sub (see _resolve_authenticated_user), and
# auth_user is also raw, so compare raw-vs-raw -- never safe_filename'd, or an
# email-style sub (a@b.com) would 403 the owner on their own session.
if str(session_state.get("user", "")) != auth_user:
return make_response(
jsonify({"success": False, "message": "Forbidden"}),
403,
@@ -592,7 +606,9 @@ class LiveSpeechToTextFinish(Resource):
404,
)
if safe_filename(str(session_state.get("user", ""))) != auth_user:
# Stored ``user`` and auth_user are both RAW subs; compare raw-vs-raw so an
# email-style sub owner is not 403'd on their own finish call.
if str(session_state.get("user", "")) != auth_user:
return make_response(
jsonify({"success": False, "message": "Forbidden"}),
403,
@@ -233,6 +233,9 @@ class GetSingleConversation(Resource):
"status": msg.get("status"),
"request_id": msg.get("request_id"),
"last_heartbeat_at": metadata.get("last_heartbeat_at"),
# Surfaced from metadata so the chat can render a
# workflow run's produced artifacts on reload.
"workflow_run_id": metadata.get("workflow_run_id"),
}
if metadata:
query["metadata"] = metadata
+4
View File
@@ -13,6 +13,7 @@ from .agents import (
agents_webhooks_ns,
)
from .analytics import analytics_ns
from .artifacts import artifacts_ns
from .attachments import attachments_ns
from .conversations import conversations_ns
from .me import me_ns
@@ -31,6 +32,9 @@ user = Blueprint("user", __name__)
# Analytics
api.add_namespace(analytics_ns)
# Artifacts
api.add_namespace(artifacts_ns)
# Attachments
api.add_namespace(attachments_ns)
+95
View File
@@ -1,3 +1,4 @@
import logging
from datetime import timedelta
from application.api.user.idempotency import with_idempotency
@@ -7,6 +8,7 @@ from application.worker import (
attachment_worker,
ingest_worker,
mcp_oauth,
parse_document_worker,
reembed_wiki_page_worker,
remote_worker,
sync,
@@ -245,6 +247,38 @@ def process_agent_webhook(self, agent_id, payload, idempotency_key=None):
return resp
# Not DURABLE: the read_document tool awaits this synchronously with a timeout, so a
# blind autoretry would double-parse and the caller would already have degraded. The
# task is routed to the dedicated ``parsing`` queue (celeryconfig task_routes) so a
# parse enqueued from inside a Celery worker (headless/scheduled agent) is served by a
# separate parsing worker and never self-deadlocks the awaiting worker.
@celery.task(bind=True, acks_late=False, autoretry_for=())
def parse_document(self, artifact_id, parent, user_id, options=None):
"""Parse an input artifact on the parsing queue; self-terminate at the soft time limit."""
from celery.exceptions import SoftTimeLimitExceeded
try:
return parse_document_worker(self, artifact_id, parent, user_id, options or {})
except SoftTimeLimitExceeded:
# A pathological/malicious document must not pin a parsing-worker slot past the
# window the caller already abandoned. Return the worker's clean error shape so
# the slot frees and the Redis result backend still gets a terminal result.
limit = getattr(self, "soft_time_limit", None)
suffix = f" after {int(limit)}s" if limit else ""
return {"status": "error", "error": f"document parsing timed out{suffix}."}
# Bind the soft limit to DOCUMENT_PARSE_TIMEOUT (the same window read_document awaits) so
# the prefork worker self-terminates a runaway parse instead of pinning the slot; the hard
# limit is the SIGKILL backstop if the soft handler can't unwind in time.
try:
from application.core.settings import settings as _parse_settings
parse_document.soft_time_limit = int(_parse_settings.DOCUMENT_PARSE_TIMEOUT)
parse_document.time_limit = parse_document.soft_time_limit + 30
except Exception:
pass
@celery.task(**DURABLE_TASK)
@with_idempotency(
task_name="ingest_connector_task", on_poison=_emit_ingest_poison_event,
@@ -342,6 +376,52 @@ def cleanup_schedule_runs(self):
return {"deleted": deleted, "ttl_days": ttl_days}
@celery.task(bind=True, acks_late=False)
def reap_sandbox_sessions(self):
"""Close sandbox sessions idle past their TTL in this worker process.
The SandboxManager registry is per-process, so this reaps only sessions
bound in THIS worker; the API processes reap their own opportunistically on
``open``. Artifacts are persisted eagerly, so reaping only closes idle
kernels and never loses a user-facing artifact.
"""
try:
from application.sandbox.sandbox_creator import SandboxCreator
reaped = SandboxCreator.get_manager().reap_expired()
except Exception: # noqa: BLE001 - housekeeping must never crash the beat loop
logging.getLogger(__name__).exception("reap_sandbox_sessions failed")
return {"reaped": 0, "error": True}
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
@@ -407,6 +487,21 @@ def setup_periodic_tasks(sender, **kwargs):
cleanup_schedule_runs.s(),
name="cleanup-schedule-runs",
)
# Close idle-past-TTL sandbox sessions roughly every minute. The on-open
# opportunistic reap still runs in the API processes; this covers worker
# processes (and quiet periods where no new session is opened).
sender.add_periodic_task(
timedelta(seconds=60),
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
+50
View File
@@ -11,13 +11,17 @@ from application.agents.default_tools import (
is_builtin_agent_tool_id,
is_default_tool_id,
is_synthesized_tool_id,
WORKFLOW_ONLY_BUILTINS,
)
from application.agents.tools.spec_parser import parse_spec
from application.agents.tools.tool_manager import ToolManager
from application.api import api
from application.api.user.artifacts.authz import Principal, authorize_artifact
from application.api.user.team_sharing import effective_write_owner, visible_with_access
from application.core.url_validation import SSRFError, validate_url
from application.security.encryption import decrypt_credentials, encrypt_credentials
from application.storage.db.base_repository import looks_like_uuid
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.repositories.notes import NotesRepository
from application.storage.db.repositories.todos import TodosRepository
from application.storage.db.repositories.user_tools import UserToolsRepository
@@ -271,12 +275,18 @@ class GetTools(Resource):
# Builtins (e.g. scheduler) hidden from Add-Tool catalog, visible
# to the agent picker. Skip ones already added via the default
# path — both registries share ``_DEFAULT_TOOL_NAMESPACE``.
# ``workflow_only`` builtins (e.g. ``read_document``) carry that
# flag so the classic picker can hide them and the workflow node
# picker can keep them.
for builtin_row in builtin_agent_tools_for_management():
builtin_copy = _row_to_api(builtin_row)
if str(builtin_copy["id"]) in seen_ids:
continue
builtin_copy["builtin"] = True
builtin_copy["default"] = False
builtin_copy["workflow_only"] = (
builtin_copy.get("name") in WORKFLOW_ONLY_BUILTINS
)
user_tools.append(builtin_copy)
except Exception as err:
current_app.logger.error(f"Error getting user tools: {err}", exc_info=True)
@@ -934,6 +944,46 @@ class GetArtifact(Resource):
return make_response(
jsonify({"success": True, "artifact": artifact}), 200
)
# Generalized document/file artifacts live in the
# ``artifacts`` store, authorized by their parent (not user_id).
# Shape-gate the UUID lookup so a non-UUID id (legacy note/todo
# ObjectId already handled above) never reaches the CAST that
# would poison this read-only transaction.
artifact_doc = None
if looks_like_uuid(artifact_id):
artifacts_repo = ArtifactsRepository(conn)
artifact_doc = artifacts_repo.get_artifact(artifact_id)
if artifact_doc and authorize_artifact(
conn, artifact_doc, Principal(user_id=user_id)
):
current = artifacts_repo.get_version(
artifact_id, artifact_doc.get("current_version")
)
artifact = {
"artifact_type": (
"document"
if artifact_doc.get("kind") != "file"
else "file"
),
"data": {
"id": str(artifact_doc.get("id")),
"kind": artifact_doc.get("kind"),
"title": artifact_doc.get("title"),
"current_version": artifact_doc.get("current_version"),
"mime_type": current.get("mime_type") if current else None,
"filename": current.get("filename") if current else None,
"size": current.get("size") if current else None,
"download_url": (
f"/api/artifacts/{artifact_id}/download"
if current and current.get("storage_path")
else None
),
},
}
return make_response(
jsonify({"success": True, "artifact": artifact}), 200
)
except Exception as err:
current_app.logger.error(
f"Error retrieving artifact: {err}", exc_info=True
+13
View File
@@ -357,6 +357,19 @@ def validate_workflow_structure(
if schema_error:
errors.append(f"Agent node '{agent_title}' JSON schema {schema_error}")
code_nodes = [n for n in nodes if n.get("type") == "code"]
for code_node in code_nodes:
code_title = code_node.get("title", code_node.get("id", "unknown"))
raw_config = code_node.get("data", {}) or {}
if not isinstance(raw_config, dict):
errors.append(f"Code node '{code_title}' has invalid configuration")
continue
if not str(raw_config.get("code", "")).strip():
errors.append(f"Code node '{code_title}' must have code to execute")
_, schema_error = validate_json_schema_payload(raw_config.get("json_schema"))
if schema_error:
errors.append(f"Code node '{code_title}' JSON schema {schema_error}")
for node in nodes:
if not node.get("id"):
errors.append("All nodes must have an id")
+18
View File
@@ -1,3 +1,5 @@
from kombu import Queue
from application.core.settings import settings
# Pydantic loads .env into ``settings`` but does not inject values into
@@ -19,6 +21,22 @@ task_default_queue = "docsgpt"
task_default_exchange = "docsgpt"
task_default_routing_key = "docsgpt"
# Route document parsing to a dedicated queue so a parse enqueued from inside a
# Celery worker (headless/scheduled agent) is served by a separate parsing worker
# and never self-deadlocks the awaiting worker. The tool also passes the queue at
# apply_async time, so this routing is the default for any other enqueuer.
task_routes = {
"application.api.user.tasks.parse_document": {"queue": settings.DOCUMENT_PARSE_QUEUE},
}
# Declare every queue so a bare ``celery worker`` (no -Q) consumes ALL of them —
# the default worker does the whole job, parsing included. Operators who want
# heavy OCR isolated run one worker with ``-Q docsgpt`` and another with
# ``-Q parsing``. (dict.fromkeys dedupes if DOCUMENT_PARSE_QUEUE == "docsgpt".)
task_queues = tuple(
Queue(name) for name in dict.fromkeys(["docsgpt", settings.DOCUMENT_PARSE_QUEUE])
)
beat_scheduler = "redbeat.RedBeatScheduler"
redbeat_redis_url = broker_url
redbeat_key_prefix = "redbeat:docsgpt:"
+71 -2
View File
@@ -251,8 +251,11 @@ 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.
DEFAULT_CHAT_TOOLS: list = ["memory", "read_webpage", "scheduler"]
DEFAULT_CHAT_TOOLS: list = [
"memory",
"read_webpage",
"scheduler",
]
# Conversation Compression Settings
ENABLE_CONVERSATION_COMPRESSION: bool = True
@@ -326,6 +329,72 @@ class Settings(BaseSettings):
SCHEDULE_ONCE_MAX_HORIZON: int = 31_536_000
SCHEDULE_RUN_OUTPUT_RETENTION_DAYS: int = 90
# Code-execution sandbox (see artifacts-code-execution-spec.md §4 C2).
# The app is a CLIENT of an always-on runner; defaults are safe so app
# import never fails when the sandbox is unconfigured.
SANDBOX_BACKEND: str = "jupyter" # "jupyter" (self-host) | "daytona" (Daytona Cloud)
# URL of the Jupyter Kernel Gateway runner (the docsgpt-sandbox service).
SANDBOX_GATEWAY_URL: str = "http://localhost:8888"
SANDBOX_GATEWAY_AUTH_TOKEN: Optional[str] = None # gateway auth token, if set
# Kernelspec launched per session. Defaults to the env-scrubbing "docsgpt-python"
# spec (shipped by the docsgpt-sandbox runner) so kernel code cannot read the
# gateway auth token or operator secrets from os.environ. The stock "python3"
# spec inherits the gateway env verbatim and must not be used with untrusted code.
SANDBOX_KERNEL_NAME: str = "docsgpt-python"
SANDBOX_MAX_TTL: int = 1200 # hard cap (s) on agent-selectable keep-alive TTL
# Per-process/worker cap on concurrent live sandbox sessions. Backend-agnostic
# (complements DAYTONA_MAX_SANDBOXES); when reached, an LRU-idle session is
# evicted to make room. This bound is local to each app/worker process.
# 0 (or any non-positive value) disables the cap (unlimited sessions).
SANDBOX_MAX_SESSIONS: int = 32
SANDBOX_EXEC_TIMEOUT: int = 60 # default wall-clock cap (s) per exec call
SANDBOX_HTTP_TIMEOUT: int = 10 # fixed cap (s) for REST control calls (create/delete/alive/interrupt)
SANDBOX_MAX_OUTPUT_BYTES: int = 8 * 1024 * 1024 # cap on buffered stdout+stderr per exec
SANDBOX_MAX_FILE_BYTES: int = 10 * 1024 * 1024 # cap on get_file size routed through stdout
SANDBOX_MAX_INPUT_BYTES: int = 25 * 1024 * 1024 # cap on an input document staged into a sandbox session
# ``read_document`` parsing on a dedicated Celery ``parsing`` queue (backend parser).
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.
WORKFLOW_NODE_NATIVE_MAX_FILES: int = 5
# Per-agent-node cap on documents extracted to text via the parsing worker.
# Each non-native, non-text document issues a separate blocking parse, so a
# node referencing many documents (e.g. the ``*`` token) is bounded here to
# 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.
SANDBOX_MEMORY: str = "1g" # docker mem_limit for the runner container
SANDBOX_CPUS: str = "1.0" # docker cpu quota for the runner container
# Daytona Cloud managed backend (used only when SANDBOX_BACKEND="daytona").
# The app is a REST client of Daytona Cloud authenticated by DAYTONA_API_KEY;
# all knobs are optional so app import never fails when the backend is unused.
DAYTONA_API_KEY: Optional[str] = None # Daytona Cloud API key (secret)
DAYTONA_API_URL: Optional[str] = None # override Daytona API base URL, if self-targeting
DAYTONA_TARGET: Optional[str] = None # Daytona region/target, e.g. "us"
DAYTONA_SNAPSHOT: Optional[str] = None # image for new sandboxes; render libs via scripts/build_daytona_snapshot.py
DAYTONA_LANGUAGE: str = "python" # default runtime language for created sandboxes
DAYTONA_AUTO_STOP_INTERVAL: int = 15 # minutes idle before Daytona auto-stops a sandbox (0 disables)
DAYTONA_AUTO_DELETE_INTERVAL: int = 60 # minutes after stop before Daytona auto-deletes (-1 disables)
DAYTONA_MAX_SANDBOXES: int = 50 # cap on concurrent live Daytona sandboxes (cost-DoS guard)
# Per-user artifact quotas (generous defaults; enforced at persistence time).
# For all three, 0 (or any non-positive value) disables that quota (unlimited).
ARTIFACT_MAX_BYTES: int = 50 * 1024 * 1024 # cap on a single stored artifact version's bytes
ARTIFACT_MAX_COUNT_PER_USER: int = 5000 # cap on artifacts a user may own
ARTIFACT_MAX_TOTAL_BYTES_PER_USER: int = 5 * 1024 * 1024 * 1024 # cap on a user's total stored bytes
@field_validator("POSTGRES_URI", mode="before")
@classmethod
def _normalize_postgres_uri_validator(cls, v):
+58 -2
View File
@@ -963,6 +963,7 @@ class LLMHandler(ABC):
_mark_failed(
pause_info["call_id"],
f"headless: {deny_reason}",
message_id=agent.tool_executor.message_id,
user_id=agent.tool_executor.user,
)
yield {
@@ -1000,6 +1001,10 @@ class LLMHandler(ABC):
# They will be added on resume to keep call/result pairs together.
continue
# One assistant(tool_calls) message per call: track whether the
# success path already appended it so the except below doesn't
# add a second one when create_tool_message fails post-append.
assistant_appended = False
try:
self.tool_calls.append(call)
tool_executor_gen = agent._execute_tool_action(tools_dict, call)
@@ -1041,14 +1046,65 @@ class LLMHandler(ABC):
if reasoning_content:
assistant_msg["reasoning_content"] = reasoning_content
updated_messages.append(assistant_msg)
assistant_appended = True
updated_messages.append(self.create_tool_message(call, tool_response))
# The tool result's tool_call_id must match the id put on the
# assistant tool_call above (``call_id`` — a synthesized UUID
# when the provider omitted an id), not the raw ``call.id`` which
# may be empty. A mismatch orphans the tool message and 400s the
# next completion ("'tool' must be a response to a preceding
# message with 'tool_calls'").
resolved_call = ToolCall(
id=call_id, name=call.name, arguments=call.arguments
)
updated_messages.append(
self.create_tool_message(resolved_call, tool_response)
)
except Exception as e:
logger.error(f"Error executing tool: {str(e)}", exc_info=True)
# The error tool message's tool_call_id must match the assistant
# tool_call id that precedes it. When the success path already
# appended one (create_tool_message failed afterwards) that id is
# the executor-returned ``call_id``; otherwise the except builds
# its own assistant message below from ``call.id``.
error_id = call_id if assistant_appended else call.id
error_call = ToolCall(
id=call.id, name=call.name, arguments=call.arguments
id=error_id, name=call.name, arguments=call.arguments
)
error_response = f"Error executing tool: {str(e)}"
# Mirror the success path: a role:"tool" message must follow
# an assistant message carrying its tool_calls, or the next
# provider completion 400s ("'tool' must be a response to a
# preceding message with 'tool_calls'"). Skip it when the
# success path already appended one for this call — a
# create_tool_message failure after that append would
# otherwise duplicate the assistant message and 400 the
# same way an orphan tool message does.
if not assistant_appended:
args_str = (
json.dumps(call.arguments)
if isinstance(call.arguments, dict)
else call.arguments
)
tool_call_obj = {
"id": call.id,
"type": "function",
"function": {
"name": call.name,
"arguments": args_str,
},
}
if call.thought_signature:
tool_call_obj["thought_signature"] = call.thought_signature
assistant_msg: Dict[str, Any] = {
"role": "assistant",
"content": None,
"tool_calls": [tool_call_obj],
}
if reasoning_content:
assistant_msg["reasoning_content"] = reasoning_content
updated_messages.append(assistant_msg)
error_message = self.create_tool_message(error_call, error_response)
updated_messages.append(error_message)
+5 -3
View File
@@ -1,7 +1,8 @@
"""FastMCP server exposing DocsGPT retrieval over streamable HTTP.
Mounted at ``/mcp`` by ``application/asgi.py``. Bearer tokens are the
existing DocsGPT agent API keys — no new credential surface.
existing DocsGPT agent API keys — no new credential surface. The
``search_docs`` tool searches the caller's knowledge base.
The tool reads the ``Authorization`` header directly via
``get_http_headers(include={"authorization"})``. The ``include`` kwarg
@@ -27,8 +28,6 @@ from application.services.search_service import (
logger = logging.getLogger(__name__)
mcp = FastMCP("docsgpt")
def _extract_bearer_token() -> str | None:
auth = get_http_headers(include={"authorization"}).get("authorization", "")
@@ -38,6 +37,9 @@ def _extract_bearer_token() -> str | None:
return parts[1]
mcp = FastMCP("docsgpt")
@mcp.tool
async def search_docs(query: str, chunks: int = 5) -> list[dict]:
"""Search the caller's DocsGPT knowledge base.
+508
View File
@@ -0,0 +1,508 @@
"""In-process document parsing for the ``read_document`` tool, run on the Celery parsing worker.
``parse_document_bytes`` turns untrusted document bytes into a bounded, shaped
result (markdown/text/structured/chunks) using the BACKEND parsers (Docling by
default). It applies the same untrusted-content safeguards as uploads — an
extension whitelist, a byte cap, ``safe_filename`` staging into a temp file, and
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, Iterator, List, Optional
from application.core.settings import settings
from application.parser.file.bulk import get_default_file_extractor
from application.parser.file.constants import SUPPORTED_SOURCE_EXTENSIONS
from application.utils import safe_filename
logger = logging.getLogger(__name__)
# Default cap for the LLM-facing VIEW of the text (applied in ``bound_parse_payload``,
# NOT during parsing) so a huge document can't flood context; the full result is still
# persisted as a ``data`` artifact. When the text exceeds the cap a head+tail window keeps
# both the document's beginning AND end (e.g. totals/signatures) within the byte budget.
_TEXT_MAX_BYTES = 8000
_MAX_TABLES_RETURNED = 20
_MAX_TABLE_ROWS = 50
_MAX_CELL_CHARS = 200
# Caps applied to the bounded view that rides back through the Redis result
# backend (the full result still lives in the persisted artifact).
_MAX_CHUNKS_RETURNED = 50
_MAX_PAGE_SELECTIONS = 10_000
_MAX_PAGE_SELECTOR_TOKENS = 20_000
_VALID_OUTPUTS = ("markdown", "text", "structured", "chunks")
_VALID_OCR = ("auto", "on", "off")
_VALID_ENGINES = ("auto", "docling", "fast")
def truncate_text_head_tail(text: str, max_bytes: Optional[int] = None) -> str:
"""Bound text to a head+tail byte window so a large file can't flood context."""
cap = int(max_bytes or _TEXT_MAX_BYTES)
if cap <= 0:
return text
encoded = text.encode("utf-8")
if len(encoded) <= cap:
return text
head = cap // 2
tail = cap - head
dropped = len(encoded) - head - tail
head_text = encoded[:head].decode("utf-8", errors="ignore")
tail_text = encoded[-tail:].decode("utf-8", errors="ignore")
return f"{head_text}\n\n...[truncated {dropped} bytes]...\n\n{tail_text}"
def bound_parse_payload(payload: Dict[str, Any], max_chars: Optional[int] = None) -> Dict[str, Any]:
"""Bound every shape of a parse payload so the Redis-backed result stays small.
This is where ALL view-bounding happens: parsing now returns the FULL content and the
persisted ``data`` artifact keeps it, while the view ridden back through the Redis result
backend is bounded here. ``content`` is capped to ``max_chars`` when given, else re-windowed
to a head+tail byte window; ``chunks`` is capped in count and per-chunk length. ``structured``
is left as-is: it rides back so json_schema validation in the tool can run against it, and it
is already bounded by the input byte cap plus the table caps (``_compact_table`` /
``summary``). ``payload['truncated']`` is set when the content view actually cut. The dict is
mutated in place.
"""
content = payload.get("content")
if isinstance(content, str):
if max_chars and int(max_chars) > 0:
capped = content[: int(max_chars)]
else:
capped = truncate_text_head_tail(content)
payload["content"] = capped
payload["truncated"] = capped != content
chunks = payload.get("chunks")
if isinstance(chunks, list):
bounded = [
truncate_text_head_tail(chunk) if isinstance(chunk, str) else chunk
for chunk in chunks[:_MAX_CHUNKS_RETURNED]
]
if len(chunks) > _MAX_CHUNKS_RETURNED:
payload["chunks_truncated"] = True
payload["total_chunks"] = len(chunks)
payload["chunks"] = bounded
return payload
def _max_input_bytes() -> int:
"""Return the size cap for a parsed document (its own setting, else the sandbox cap)."""
explicit = int(getattr(settings, "DOCUMENT_PARSE_MAX_BYTES", 0) or 0)
if explicit > 0:
return explicit
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":
return True
if ocr == "off":
return False
return bool(getattr(settings, "DOCLING_OCR_ENABLED", False))
def _pick_parser(suffix: str, *, ocr_enabled: bool, engine: str):
"""Select the parser for ``suffix`` honoring the requested engine; None when unsupported."""
if engine == "fast":
legacy = _legacy_parser_for(suffix)
if legacy is not None:
return legacy
extractor = get_default_file_extractor(ocr_enabled=ocr_enabled)
return extractor.get(suffix)
def _legacy_parser_for(suffix: str):
"""Return a non-Docling parser for ``suffix`` (the ``fast`` engine), or None."""
from application.parser.file.docs_parser import DocxParser, PDFParser
from application.parser.file.html_parser import HTMLParser
from application.parser.file.markdown_parser import MarkdownParser
from application.parser.file.tabular_parser import ExcelParser, PandasCSVParser
legacy = {
".pdf": PDFParser,
".docx": DocxParser,
".csv": PandasCSVParser,
".xlsx": ExcelParser,
".html": HTMLParser,
".md": MarkdownParser,
".mdx": MarkdownParser,
}
cls = legacy.get(suffix)
return cls() if cls is not None else None
def _parse_to_text(parser: Any, path: Path) -> str:
"""Run a parser and coerce its ``str | List[str]`` result to a single text blob."""
if not parser.parser_config_set:
parser.init_parser()
parsed = parser.parse_file(path, errors="ignore")
if isinstance(parsed, list):
return "\n\n".join(str(part) for part in parsed)
return str(parsed)
def _is_docling_parser(parser: Any) -> bool:
"""True when ``parser`` is Docling-backed (collecting tables would otherwise re-convert)."""
if parser is None:
return False
try:
from application.parser.file.docling_parser import DoclingParser
except Exception:
return False
return isinstance(parser, DoclingParser)
def _compact_table(table: Dict[str, Any]) -> Dict[str, Any]:
"""Bound a single table's rows and cell sizes so one giant table can't bloat context."""
def _cell(value: Any) -> Any:
if isinstance(value, str) and len(value) > _MAX_CELL_CHARS:
return value[:_MAX_CELL_CHARS] + "...[truncated]"
return value
rows = table.get("rows")
if not isinstance(rows, list):
return table
capped = [[_cell(c) for c in row] if isinstance(row, list) else _cell(row) for row in rows[:_MAX_TABLE_ROWS]]
compact = dict(table)
compact["rows"] = capped
if len(rows) > _MAX_TABLE_ROWS:
compact["rows_truncated"] = True
compact["total_rows"] = len(rows)
return compact
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
path), reuse ITS converter + export so OCR/pipeline options and the export fallback
are honored and the content matches the legacy single-parse output exactly; otherwise
fall back to a vanilla converter.
"""
if _is_docling_parser(parser):
if getattr(parser, "_converter", None) is None:
parser._init_parser()
doc = parser._converter.convert(str(path)).document
markdown = parser._export_content(doc)
else:
from docling.document_converter import DocumentConverter
converter = DocumentConverter()
doc = converter.convert(str(path)).document
markdown = doc.export_to_markdown()
structured = doc.export_to_dict()
tables: List[Dict[str, Any]] = []
if include_tables:
for tbl in getattr(doc, "tables", []) or []:
try:
df = tbl.export_to_dataframe()
tables.append({"columns": [str(c) for c in df.columns], "rows": df.astype(str).values.tolist()})
except Exception:
try:
tables.append({"markdown": tbl.export_to_markdown()})
except Exception:
continue
if len(tables) >= _MAX_TABLES_RETURNED:
break
page_count = len(getattr(doc, "pages", {}) or {})
return {"markdown": markdown, "structured": structured, "tables": tables, "page_count": page_count}
def _structure_summary(structured: Any) -> Dict[str, int]:
"""Summarize the Docling structured dict by top-level element counts (keeps context compact)."""
if not isinstance(structured, dict):
return {}
counts: Dict[str, int] = {}
for key in ("texts", "tables", "pictures", "groups", "pages"):
value = structured.get(key)
if isinstance(value, (list, dict)):
counts[key] = len(value)
return counts
def _apply_pages(text: str, pages: Any) -> str:
"""Best-effort page-range slice on a page-delimited markdown blob (``\\f`` separated)."""
if not pages or "\f" not in text:
return text
total = text.count("\f") + 1
selected = _selected_page_indices(pages, total)
if not selected:
return text
# Splitting a hostile form-feed blob could allocate millions of strings.
# Walk boundaries and retain only the bounded set of requested pages.
wanted = set(selected)
found: Dict[int, tuple[int, int]] = {}
page_index = 0
start = 0
while wanted:
end = text.find("\f", start)
if end < 0:
end = len(text)
if page_index in wanted:
found[page_index] = (start, end)
wanted.remove(page_index)
if end == len(text):
break
page_index += 1
start = end + 1
# Preserve requested order and ordinary duplicates, but never let repeated
# selectors amplify the output beyond the source text's resident size.
output = io.StringIO()
output_chars = 0
wrote_page = False
for index in selected:
span = found.get(index)
if span is None:
continue
page_start, page_end = span
added = (page_end - page_start) + (1 if wrote_page else 0)
if output_chars + added > len(text):
break
if wrote_page:
output.write("\f")
output.write(text[page_start:page_end])
output_chars += added
wrote_page = True
return output.getvalue() if wrote_page else text
def _iter_page_tokens(pages: Any) -> Iterator[Any]:
"""Yield bounded selector tokens without materializing a comma-split list."""
if isinstance(pages, list):
for position, token in enumerate(pages):
if position >= _MAX_PAGE_SELECTOR_TOKENS:
break
yield token
return
raw = str(pages)
start = 0
emitted = 0
while emitted < _MAX_PAGE_SELECTOR_TOKENS:
end = raw.find(",", start)
if end < 0:
yield raw[start:]
return
yield raw[start:end]
emitted += 1
start = end + 1
def _selected_page_indices(pages: Any, total: int) -> List[int]:
"""Parse ``pages`` into a bounded list of valid 0-based page occurrences."""
if total <= 0:
return []
indices: List[int] = []
for token in _iter_page_tokens(pages):
if len(indices) >= _MAX_PAGE_SELECTIONS:
break
token = str(token).strip()
if "-" in token:
try:
lo, hi = (int(p) for p in token.split("-", 1))
except ValueError:
continue
start = max(lo - 1, 0)
stop = min(hi, total, start + (_MAX_PAGE_SELECTIONS - len(indices)))
indices.extend(range(start, stop))
else:
try:
index = int(token) - 1
except ValueError:
continue
if 0 <= index < total:
indices.append(index)
return indices
def _to_chunks(text: str, max_chars: Optional[int]) -> List[str]:
"""Chunk parsed text via the ingestion chunker; bounded and JSON-safe for the result."""
from application.parser.chunking_creator import ChunkerCreator
from application.parser.schema.base import Document
chunker = ChunkerCreator.create_chunker("classic_chunk")
chunks = chunker.chunk([Document(text=text)])
cap = int(max_chars or 0)
out: List[str] = []
for chunk in chunks:
body = getattr(chunk, "text", str(chunk))
out.append(body[:cap] if cap > 0 else body)
if len(out) >= 200:
break
return out
def parse_document_bytes(
data: bytes,
filename: str,
*,
output: str = "markdown",
ocr: str = "auto",
pages: Any = None,
engine: str = "auto",
max_chars: Optional[int] = None,
include_tables: bool = True,
) -> Dict[str, Any]:
"""Parse untrusted document bytes into a bounded shaped result; whitelist + size + cleanup guarded."""
if output not in _VALID_OUTPUTS:
return {"error": f"unsupported output '{output}'; expected one of {_VALID_OUTPUTS}."}
if ocr not in _VALID_OCR:
return {"error": f"unsupported ocr '{ocr}'; expected one of {_VALID_OCR}."}
if engine not in _VALID_ENGINES:
return {"error": f"unsupported engine '{engine}'; expected one of {_VALID_ENGINES}."}
safe_name = safe_filename(filename) or "document"
suffix = os.path.splitext(safe_name)[1].lower()
if suffix not in SUPPORTED_SOURCE_EXTENSIONS:
return {"error": f"unsupported file type '{suffix or filename}'."}
cap = _max_input_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
try:
tmp_path.write_bytes(data)
return _shape(tmp_path, suffix, output, ocr_enabled, engine, pages, max_chars, include_tables)
except Exception as exc:
logger.exception("parse_document_bytes: parsing failed")
return {"error": f"parsing failed: {type(exc).__name__}: {exc}"}
finally:
try:
tmp_path.unlink(missing_ok=True)
os.rmdir(tmp_dir)
except OSError:
logger.warning("parse_document_bytes: temp cleanup failed for %s", tmp_dir, exc_info=True)
def _shape(
path: Path,
suffix: str,
output: str,
ocr_enabled: bool,
engine: str,
pages: Any,
max_chars: Optional[int],
include_tables: bool,
) -> Dict[str, Any]:
"""Run the selected parser/engine and shape the result per ``output``; bounded throughout.
``output='structured'`` always uses Docling regardless of ``engine`` — the ``fast``
engine is markdown/text only and cannot produce the structured dict.
"""
if output == "structured":
try:
extracted = _docling_structured(path, ocr_enabled=ocr_enabled, include_tables=include_tables)
except Exception as exc:
return {"error": f"structured parsing requires Docling: {type(exc).__name__}: {exc}"}
bounded, truncated = _bounded(extracted["markdown"])
return {
"output": "structured",
"content": bounded,
"truncated": truncated,
"tables": [_compact_table(t) for t in extracted["tables"]],
"structured": extracted["structured"],
"summary": _structure_summary(extracted["structured"]),
"page_count": extracted["page_count"],
}
parser = _pick_parser(suffix, ocr_enabled=ocr_enabled, engine=engine)
wants_tables = include_tables and engine != "fast" and output != "chunks"
# A Docling-backed parser already converts the whole document to produce its text.
# When tables are also requested, reuse that single conversion for both the markdown
# content and the tables instead of converting a second time just to collect tables
# (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)
text = extracted["markdown"]
tables: List[Dict[str, Any]] = [_compact_table(t) for t in extracted["tables"]]
except Exception:
text, tables = _parse_to_text(parser, path), []
text = _apply_pages(text, pages)
bounded, truncated = _bounded(text)
payload: Dict[str, Any] = {"output": output, "content": bounded, "truncated": truncated}
if tables:
payload["tables"] = tables
return payload
if parser is None:
# A whitelisted extension with no dedicated parser (e.g. .txt) reads as plain
# text, matching SimpleDirectoryReader's standard-read fallback.
text = path.read_text(errors="ignore")
else:
text = _parse_to_text(parser, path)
text = _apply_pages(text, pages)
if output == "chunks":
return {"output": "chunks", "chunks": _to_chunks(text, max_chars), "truncated": False}
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"]
]
except Exception:
tables = []
bounded, truncated = _bounded(text)
payload: Dict[str, Any] = {"output": output, "content": bounded, "truncated": truncated}
if tables:
payload["tables"] = tables
return payload
def _bounded(text: str) -> tuple[str, bool]:
"""Return the FULL extracted text (never truncated here); the view is bounded in ``bound_parse_payload``.
Parsing keeps the complete text so the persisted ``data`` artifact is the full parse;
``max_chars`` and the default head+tail window now bound only the LLM-facing view.
"""
return text, False
+19
View File
@@ -8,6 +8,13 @@ You are DocsGPT, an AI assistant that answers questions using the user's documen
- Use other available tools to fill gaps instead of guessing: read shared URLs with a webpage tool, and use a memory tool (when available) to recall earlier context or save durable facts and preferences. Never invent data, sources, or tool results.
- Respond in the same language as the user's message.
{% if tools.enabled is not defined or 'artifact_generator' in tools.enabled or 'code_executor' in tools.enabled %}
## Producing documents and running code
- When the user wants a document, slide deck, spreadsheet, PDF, or other file and a document or artifact tool is available, create it as an artifact instead of pasting the full file into the chat. The user gets a downloadable, versioned file they can reopen and edit.
- For follow-up changes to a file you already produced, edit that artifact with a targeted change rather than regenerating it from scratch, so its version history stays clean.
- When a code-execution tool is available, run code for real computation, data processing, file parsing or conversion, and charts instead of estimating or writing results by hand. Files the code writes are saved as downloadable artifacts, so write scratch or intermediate files under `tmp/` (not saved), pass `outputs` to save only specific files, or set `capture_artifacts` to false for setup-only steps; never paste raw file contents. Each run is time-limited, so start long work in the background and check on it with another run.
{% endif %}
## Formatting
- Use markdown. Put code in fenced blocks with a language tag.
- Only use a mermaid diagram when the user asks for one or a diagram is clearly the best way to answer, and make sure the syntax is valid.
@@ -23,3 +30,15 @@ Your memory directory (saved in previous conversations):
</memory_directory>
Read relevant memory files with memory_view before answering questions that may depend on them; save new durable facts and preferences with memory_create.
{% endif %}
{% if attachments.files %}
## Attached files
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 %}
{% 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 %}
+19
View File
@@ -7,6 +7,13 @@ You are DocsGPT, an AI assistant that answers questions using the user's documen
- Use other available tools to fill gaps instead of guessing: read shared URLs with a webpage tool, and use a memory tool (when available) to recall earlier context or save durable facts and preferences. Never invent data, sources, or tool results.
- Respond in the same language as the user's message.
{% if tools.enabled is not defined or 'artifact_generator' in tools.enabled or 'code_executor' in tools.enabled %}
## Producing documents and running code
- When the user wants a document, slide deck, spreadsheet, PDF, or other file and a document or artifact tool is available, create it as an artifact instead of pasting the full file into the chat. The user gets a downloadable, versioned file they can reopen and edit.
- For follow-up changes to a file you already produced, edit that artifact with a targeted change rather than regenerating it from scratch, so its version history stays clean.
- When a code-execution tool is available, run code for real computation, data processing, file parsing or conversion, and charts instead of estimating or writing results by hand. Files the code writes are saved as downloadable artifacts, so write scratch or intermediate files under `tmp/` (not saved), pass `outputs` to save only specific files, or set `capture_artifacts` to false for setup-only steps; never paste raw file contents. Each run is time-limited, so start long work in the background and check on it with another run.
{% endif %}
## Formatting
- Use markdown. Put code in fenced blocks with a language tag.
- Only use a mermaid diagram when the user asks for one or a diagram is clearly the best way to answer, and make sure the syntax is valid.
@@ -22,3 +29,15 @@ Your memory directory (saved in previous conversations):
</memory_directory>
Read relevant memory files with memory_view before answering questions that may depend on them; save new durable facts and preferences with memory_create.
{% endif %}
{% if attachments.files %}
## Attached files
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 %}
{% 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 %}
+19
View File
@@ -7,6 +7,13 @@ You are DocsGPT, an AI assistant that answers questions using the user's documen
- If searches and tools do not yield enough information, reply that you do not have the information needed to answer and name what is missing. Never invent information.
- Respond in the same language as the user's message.
{% if tools.enabled is not defined or 'artifact_generator' in tools.enabled or 'code_executor' in tools.enabled %}
## Producing documents and running code
- When the user wants a document, slide deck, spreadsheet, PDF, or other file and a document or artifact tool is available, create it as an artifact instead of pasting the full file into the chat. The user gets a downloadable, versioned file they can reopen and edit.
- For follow-up changes to a file you already produced, edit that artifact with a targeted change rather than regenerating it from scratch, so its version history stays clean.
- When a code-execution tool is available, run code for real computation, data processing, file parsing or conversion, and charts instead of estimating or writing results by hand. Files the code writes are saved as downloadable artifacts, so write scratch or intermediate files under `tmp/` (not saved), pass `outputs` to save only specific files, or set `capture_artifacts` to false for setup-only steps; never paste raw file contents. Each run is time-limited, so start long work in the background and check on it with another run.
{% endif %}
## Formatting
- Use markdown. Put code in fenced blocks with a language tag.
- Only use a mermaid diagram when the user asks for one or a diagram is clearly the best way to answer, and make sure the syntax is valid.
@@ -22,3 +29,15 @@ Your memory directory (saved in previous conversations):
</memory_directory>
Read relevant memory files with memory_view before answering questions that may depend on them; save new durable facts and preferences with memory_create.
{% endif %}
{% if attachments.files %}
## Attached files
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 %}
{% 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 %}
@@ -7,6 +7,13 @@ You are DocsGPT, an AI assistant that answers questions using the user's documen
- Use available tools to fill gaps instead of guessing: read shared URLs with a webpage tool, and use a memory tool (when available) to recall earlier context or save durable facts and preferences. Never invent data, sources, or tool results.
- Respond in the same language as the user's message.
{% if tools.enabled is not defined or 'artifact_generator' in tools.enabled or 'code_executor' in tools.enabled %}
## Producing documents and running code
- When the user wants a document, slide deck, spreadsheet, PDF, or other file and a document or artifact tool is available, create it as an artifact instead of pasting the full file into the chat. The user gets a downloadable, versioned file they can reopen and edit.
- For follow-up changes to a file you already produced, edit that artifact with a targeted change rather than regenerating it from scratch, so its version history stays clean.
- When a code-execution tool is available, run code for real computation, data processing, file parsing or conversion, and charts instead of estimating or writing results by hand. Files the code writes are saved as downloadable artifacts, so write scratch or intermediate files under `tmp/` (not saved), pass `outputs` to save only specific files, or set `capture_artifacts` to false for setup-only steps; never paste raw file contents. Each run is time-limited, so start long work in the background and check on it with another run.
{% endif %}
## Formatting
- Use markdown. Put code in fenced blocks with a language tag.
- Only use a mermaid diagram when the user asks for one or a diagram is clearly the best way to answer, and make sure the syntax is valid.
@@ -33,3 +40,15 @@ Ground your answer in these documents and cite their titles. If they do not answ
No document context was retrieved for this message. Rely on the conversation, tools, and your general knowledge.
{% endif %}
{% if attachments.files %}
## Attached files
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 %}
{% 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 %}
@@ -6,6 +6,13 @@ You are DocsGPT, an AI assistant that answers questions using the user's documen
- Use available tools to fill gaps instead of guessing: read shared URLs with a webpage tool, and use a memory tool (when available) to recall earlier context or save durable facts and preferences. Never invent data, sources, or tool results.
- Respond in the same language as the user's message.
{% if tools.enabled is not defined or 'artifact_generator' in tools.enabled or 'code_executor' in tools.enabled %}
## Producing documents and running code
- When the user wants a document, slide deck, spreadsheet, PDF, or other file and a document or artifact tool is available, create it as an artifact instead of pasting the full file into the chat. The user gets a downloadable, versioned file they can reopen and edit.
- For follow-up changes to a file you already produced, edit that artifact with a targeted change rather than regenerating it from scratch, so its version history stays clean.
- When a code-execution tool is available, run code for real computation, data processing, file parsing or conversion, and charts instead of estimating or writing results by hand. Files the code writes are saved as downloadable artifacts, so write scratch or intermediate files under `tmp/` (not saved), pass `outputs` to save only specific files, or set `capture_artifacts` to false for setup-only steps; never paste raw file contents. Each run is time-limited, so start long work in the background and check on it with another run.
{% endif %}
## Formatting
- Use markdown. Put code in fenced blocks with a language tag.
- Only use a mermaid diagram when the user asks for one or a diagram is clearly the best way to answer, and make sure the syntax is valid.
@@ -32,3 +39,15 @@ Ground your answer in these documents and cite their titles. If they do not answ
No document context was retrieved for this message. Rely on the conversation, tools, and your general knowledge.
{% endif %}
{% if attachments.files %}
## Attached files
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 %}
{% 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 %}
@@ -6,6 +6,13 @@ You are DocsGPT, an AI assistant that answers questions using the user's documen
- Use available tools to fill gaps instead of guessing: read shared URLs with a webpage tool, and use a memory tool (when available) to recall earlier context or save durable facts and preferences. Never invent data, sources, or tool results.
- Respond in the same language as the user's message.
{% if tools.enabled is not defined or 'artifact_generator' in tools.enabled or 'code_executor' in tools.enabled %}
## Producing documents and running code
- When the user wants a document, slide deck, spreadsheet, PDF, or other file and a document or artifact tool is available, create it as an artifact instead of pasting the full file into the chat. The user gets a downloadable, versioned file they can reopen and edit.
- For follow-up changes to a file you already produced, edit that artifact with a targeted change rather than regenerating it from scratch, so its version history stays clean.
- When a code-execution tool is available, run code for real computation, data processing, file parsing or conversion, and charts instead of estimating or writing results by hand. Files the code writes are saved as downloadable artifacts, so write scratch or intermediate files under `tmp/` (not saved), pass `outputs` to save only specific files, or set `capture_artifacts` to false for setup-only steps; never paste raw file contents. Each run is time-limited, so start long work in the background and check on it with another run.
{% endif %}
## Formatting
- Use markdown. Put code in fenced blocks with a language tag.
- Only use a mermaid diagram when the user asks for one or a diagram is clearly the best way to answer, and make sure the syntax is valid.
@@ -32,3 +39,15 @@ Ground your answer strictly in these documents and cite their titles. If they do
No document context was retrieved for this message. If tools cannot provide the answer, say you do not have the information needed.
{% endif %}
{% if attachments.files %}
## Attached files
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 %}
{% 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 %}
+3 -1
View File
@@ -10,6 +10,7 @@ celery-redbeat==2.3.3
croniter==6.2.2
cryptography==46.0.7
dataclasses-json==0.6.7
daytona==0.190.1
defusedxml==0.7.1
docling>=2.16.0
rapidocr>=1.4.0
@@ -110,4 +111,5 @@ werkzeug>=3.1.0
yarl==1.23.0
markdownify==1.2.2
tldextract==5.3.1
websockets==16.0
websockets==16.0
websocket-client==1.9.0
+1
View File
@@ -0,0 +1 @@
"""Pluggable code-execution sandbox abstraction for running untrusted LLM code."""
+483
View File
@@ -0,0 +1,483 @@
"""Shared sandbox->artifact persistence: detect produced files and store them by reference.
Both the ``code_executor`` chat tool and the workflow ``code`` node run code in a
run/conversation-scoped sandbox session and must turn newly written workspace files
into ``artifacts`` rows. Only metadata (size/sha256/mime + the storage key) is stored
here; binary bytes live in ``BaseStorage`` and never enter LLM context or workflow state.
"""
from __future__ import annotations
import fnmatch
import hashlib
import io
import logging
import mimetypes
import os
from typing import Any, Dict, List, Optional, Set, Tuple
from sqlalchemy import text
from application.agents.tools.artifact_ref import make_ref
from application.core.settings import settings
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_session
from application.storage.storage_creator import StorageCreator
from application.utils import safe_filename
logger = logging.getLogger(__name__)
class QuotaExceeded(Exception):
"""Raised when persisting an artifact would breach a per-user size/count quota."""
# Cap the per-run capture work so a workspace full of pre-existing files can't turn
# one exec into an unbounded read+persist sweep.
MAX_CAPTURED_FILES = 64
# Cap the per-pass READ sweep independently of the persist cap: unchanged files are
# read (to detect a change) then skipped without counting toward MAX_CAPTURED_FILES,
# so a workspace full of unchanged files could otherwise be re-read in full every exec.
MAX_SCANNED_FILES = 256
# Auto-capture skips scratch/intermediate workspace paths so install steps, extracted
# archives, and temp files don't each become a downloadable artifact. Agents write
# throwaway files under ``tmp/``; an explicit ``outputs`` list bypasses this skip (the
# agent is then naming exactly what to keep).
_SCRATCH_PREFIXES = ("tmp/",)
_SCRATCH_SUFFIXES = (".tmp", ".lock", ".pyc", ".pyo")
_DEFAULT_KIND = "file"
# Coarse mime -> artifact kind mapping for the UI rail; defaults to "file".
_KIND_BY_MIME_PREFIX: Dict[str, str] = {
"image/": "image",
"text/html": "html",
"text/csv": "data",
"application/json": "data",
"application/vnd.openxmlformats-officedocument.presentationml": "presentation",
"application/vnd.openxmlformats-officedocument.spreadsheetml": "spreadsheet",
"application/vnd.ms-excel": "spreadsheet",
"application/vnd.openxmlformats-officedocument.wordprocessingml": "document",
"application/msword": "document",
"application/pdf": "document",
}
def unique_input_path(rel_path: str, used: Set[str]) -> str:
"""Return ``rel_path`` (reserving it) or a ``-2``/``-3``... suffixed variant when already staged.
Two inputs whose current versions share a filename would otherwise clobber each other at the
same ``inputs/{name}`` path; the numeric suffix is inserted before the extension. Mutates
``used`` to record the returned path.
"""
if rel_path not in used:
used.add(rel_path)
return rel_path
base, ext = os.path.splitext(rel_path)
n = 2
while f"{base}-{n}{ext}" in used:
n += 1
unique = f"{base}-{n}{ext}"
used.add(unique)
return unique
def infer_mime(filename: str) -> str:
"""Infer a mime type from a filename, falling back to a generic binary type."""
mime, _ = mimetypes.guess_type(filename)
return mime or "application/octet-stream"
def kind_for_mime(mime: str) -> str:
"""Map a mime type to a coarse artifact ``kind`` for the artifact rail."""
for prefix, kind in _KIND_BY_MIME_PREFIX.items():
if mime.startswith(prefix):
return kind
return _DEFAULT_KIND
def _is_scratch(rel_path: str) -> bool:
"""True for scratch/junk workspace paths excluded from auto-capture."""
if rel_path == "tmp" or rel_path.startswith(_SCRATCH_PREFIXES):
return True
if any(part == "__pycache__" or part.startswith(".") for part in rel_path.split("/")):
return True
return rel_path.endswith(_SCRATCH_SUFFIXES)
def _matches_outputs(rel_path: str, outputs: List[str]) -> bool:
"""True when a workspace path matches any caller-supplied ``outputs`` glob.
Each pattern is matched against the full workspace-relative path and the bare
filename, so ``report.pdf``, ``*.csv``, and ``out/*.json`` all work.
"""
name = rel_path.rsplit("/", 1)[-1]
return any(fnmatch.fnmatch(rel_path, pat) or fnmatch.fnmatch(name, pat) for pat in outputs)
def snapshot_signatures(manager: Any, session_id: str) -> Dict[str, Tuple[int, Optional[str]]]:
"""Map each non-input, non-scratch workspace file to a (size, sha256) signature."""
signatures: Dict[str, Tuple[int, Optional[str]]] = {}
try:
files = manager.list_files(session_id)
except Exception:
logger.exception("artifacts_capture: pre-exec listing failed")
return signatures
candidates = sorted(f for f in files if not f.startswith("inputs/") and not _is_scratch(f))
if len(candidates) > MAX_SCANNED_FILES:
logger.warning(
"artifacts_capture: pre-exec signature scan capped at %d of %d files",
MAX_SCANNED_FILES,
len(candidates),
)
for rel_path in candidates[:MAX_SCANNED_FILES]:
try:
data = manager.get_file(session_id, rel_path)
except Exception:
logger.exception("artifacts_capture: pre-exec signature read failed")
continue
signatures[rel_path] = (len(data), hashlib.sha256(data).hexdigest())
return signatures
def capture_artifacts(
manager: Any,
session_id: str,
pre_signatures: Dict[str, Tuple[int, Optional[str]]],
*,
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]]:
"""Persist workspace files that are new or whose content changed.
When ``outputs`` is given, capture only files matching those globs (the caller is
naming exactly what to keep). Otherwise auto-capture every produced file except
scratch/intermediate paths (``tmp/`` and obvious junk; see ``_is_scratch``).
Returns one artifact reference per captured file: ``{artifact_id, version,
filename, mime_type, size}`` (JSON primitives only; never bytes).
"""
try:
post_files = set(manager.list_files(session_id))
except Exception:
logger.exception("artifacts_capture: post-exec listing failed")
return []
produced = (f for f in post_files if not f.startswith("inputs/"))
if outputs:
candidates = sorted(f for f in produced if _matches_outputs(f, outputs))
else:
candidates = sorted(f for f in produced if not _is_scratch(f))
captured: List[Dict[str, Any]] = []
scanned = 0
for rel_path in candidates:
if len(captured) >= MAX_CAPTURED_FILES:
logger.warning("artifacts_capture: capture cap reached; remaining files skipped")
break
if scanned >= MAX_SCANNED_FILES:
logger.warning("artifacts_capture: read-scan cap reached; remaining files skipped")
break
scanned += 1
try:
data = manager.get_file(session_id, rel_path)
except Exception:
logger.exception("artifacts_capture: get_file failed during capture")
continue
# A pre-existing file is only captured when its content changed; an
# unchanged file is skipped so re-runs don't re-persist stale inputs.
signature = (len(data), hashlib.sha256(data).hexdigest())
if pre_signatures.get(rel_path) == signature:
continue
try:
ref = persist_artifact(
rel_path,
data,
user_id=user_id,
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
message_id=message_id,
produced_by=produced_by,
)
except QuotaExceeded:
# Out of quota: stop capturing the remaining files rather than retry
# each one. Files already captured this run are returned as-is.
logger.warning("artifacts_capture: per-user quota reached; remaining files not captured")
break
if ref is not None:
captured.append(ref)
return captured
def persist_artifact(
rel_path: str,
data: bytes,
*,
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)."""
# The sandbox filename is display-only; mime/kind are inferred from it, and
# the storage key is derived from server-controlled values downstream.
display_name = rel_path.rsplit("/", 1)[-1]
filename = safe_filename(display_name)
mime_type = infer_mime(filename)
return persist_new_artifact(
user_id=user_id,
kind=kind_for_mime(mime_type),
data=data,
filename=filename,
mime_type=mime_type,
title=display_name,
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
message_id=message_id,
produced_by=produced_by,
)
def _storage_key(user_id: str, artifact_id: str, version: int, filename: str) -> str:
"""Build the server-controlled storage key (``inputs/{user}/artifacts/{id}/v{n}/{file}``)."""
# ``inputs/{user}/artifacts/...`` is the project storage-namespace convention
# (matches attachments + spec §4); ``inputs/`` is the user's namespace root.
return f"inputs/{user_id}/artifacts/{artifact_id}/v{version}/{filename}"
def _set_version_storage_path(conn: Any, artifact_id: str, version: int, storage_path: str) -> None:
"""Set the server-derived storage key on an existing version row."""
conn.execute(
text(
"UPDATE artifact_versions SET storage_path = :p "
"WHERE artifact_id = CAST(:aid AS uuid) AND version = :v"
),
{"p": storage_path, "aid": artifact_id, "v": version},
)
def _cleanup_orphan(storage: Any, saved_key: Optional[str]) -> None:
"""Best-effort delete of a stored key whose owning transaction failed to commit."""
if saved_key is None:
return
try:
storage.delete_file(saved_key)
except Exception:
logger.exception("artifacts_capture: orphaned-key cleanup failed for %s", saved_key)
def _check_single_artifact_size(size: int) -> None:
"""Reject a single artifact version whose byte size exceeds ``ARTIFACT_MAX_BYTES``."""
max_bytes = int(getattr(settings, "ARTIFACT_MAX_BYTES", 0) or 0)
if max_bytes > 0 and size > max_bytes:
raise QuotaExceeded(f"artifact is too large: {size} bytes exceeds the {max_bytes}-byte per-file cap")
def _enforce_user_quota(repo: ArtifactsRepository, user_id: str, added_bytes: int, *, new_artifact: bool) -> None:
"""Reject the write when ``user_id`` is at/over their count or total-bytes quota.
Best-effort SOFT cap (not hard under concurrency): the count/total-bytes reads run
on the same connection as the pending insert, but under READ COMMITTED two concurrent
persists can each read a pre-insert total and both pass, briefly overshooting the cap.
A NEW artifact also consumes one count slot; appending a version only adds bytes to an
existing identity.
"""
_check_single_artifact_size(added_bytes)
max_count = int(getattr(settings, "ARTIFACT_MAX_COUNT_PER_USER", 0) or 0)
if new_artifact and max_count > 0 and repo.count_for_user(user_id) >= max_count:
raise QuotaExceeded(f"artifact count quota reached ({max_count}); delete artifacts to free space")
max_total = int(getattr(settings, "ARTIFACT_MAX_TOTAL_BYTES_PER_USER", 0) or 0)
if max_total > 0 and repo.total_bytes_for_user(user_id) + added_bytes > max_total:
raise QuotaExceeded(f"artifact storage quota reached ({max_total} bytes); delete artifacts to free space")
def _ref_from_metadata(metadata: Any) -> Optional[str]:
"""Return the short ref (``A{ref_seq}``) from a stored metadata dict, or None when absent/non-numeric."""
if not isinstance(metadata, dict):
return None
try:
seq = int(metadata.get("ref_seq"))
except (TypeError, ValueError):
return None
return make_ref(seq) if seq >= 1 else None
def _ref_for(
repo: ArtifactsRepository,
artifact_id: str,
*,
conversation_id: Optional[str],
workflow_run_id: Optional[str],
metadata: Any = None,
) -> Optional[str]:
"""Compute the short ref for an artifact: its STABLE stored ``ref_seq``, else positional fallback.
The ref_seq is assigned at creation and kept in ``metadata`` so it survives deletions of
earlier artifacts. Legacy rows created before ref_seq existed have none, so they fall back to
the (mutable) position-in-parent; ``resolve_artifact_id`` mirrors this fallback.
"""
if conversation_id is None and workflow_run_id is None:
return None
seq_ref = _ref_from_metadata(metadata)
if seq_ref is not None:
return seq_ref
try:
position = repo.position_in_parent(
artifact_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
)
except Exception:
logger.exception("artifacts_capture: failed to compute artifact ref")
return None
return make_ref(position) if position >= 1 else None
def persist_new_artifact(
*,
user_id: str,
kind: str,
data: bytes,
filename: str,
mime_type: str,
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,
) -> Optional[Dict[str, Any]]:
"""Create an artifact + version 1 and write its bytes (storage-write-last); return its reference.
The storage write is the last statement before commit, so a failed write rolls
the row back (bytes are never orphaned). The only remaining window is a commit
that fails after a successful write; that key is deleted best-effort.
"""
safe_name = safe_filename(filename)
size = len(data)
sha256 = hashlib.sha256(data).hexdigest()
storage = StorageCreator.get_storage()
saved_key: Optional[str] = None
ref: Optional[str] = None
try:
with db_session() as conn:
repo = ArtifactsRepository(conn)
_enforce_user_quota(repo, user_id, size, new_artifact=True)
artifact = repo.create_artifact(
user_id,
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,
storage_path=None,
size=size,
sha256=sha256,
spec=spec,
preview_text=preview_text,
produced_by=produced_by,
)
artifact_id = str(artifact["id"])
storage_path = _storage_key(user_id, artifact_id, 1, safe_name)
_set_version_storage_path(conn, artifact_id, 1, storage_path)
# The freshly-created row carries the stable ``ref_seq`` in its metadata.
ref = _ref_for(
repo,
artifact_id,
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
metadata=artifact.get("metadata"),
)
storage.save_file(io.BytesIO(data), storage_path)
saved_key = storage_path
except QuotaExceeded:
# Quota check runs before any storage write, so nothing to clean up;
# surface a clean error the caller can render.
raise
except Exception:
logger.exception("artifacts_capture: failed to persist new artifact")
_cleanup_orphan(storage, saved_key)
return None
payload: Dict[str, Any] = {
"artifact_id": artifact_id,
"version": 1,
"filename": safe_name,
"mime_type": mime_type,
"size": size,
}
if ref is not None:
payload["ref"] = ref
return payload
def append_artifact_version(
*,
user_id: str,
artifact_id: str,
data: bytes,
filename: str,
mime_type: str,
spec: Any = None,
preview_text: Optional[str] = None,
produced_by: Any = None,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
"""Append a new version (new spec + new bytes) to an existing artifact; return its reference."""
safe_name = safe_filename(filename)
size = len(data)
sha256 = hashlib.sha256(data).hexdigest()
storage = StorageCreator.get_storage()
saved_key: Optional[str] = None
ref: Optional[str] = None
try:
with db_session() as conn:
repo = ArtifactsRepository(conn)
_enforce_user_quota(repo, user_id, size, new_artifact=False)
version = repo.append_version(
artifact_id,
mime_type=mime_type,
filename=safe_name,
storage_path=None,
size=size,
sha256=sha256,
spec=spec,
preview_text=preview_text,
produced_by=produced_by,
)
version_number = int(version["version"])
storage_path = _storage_key(user_id, artifact_id, version_number, safe_name)
_set_version_storage_path(conn, artifact_id, version_number, storage_path)
# The identity's stable ``ref_seq`` was assigned at creation; read it so an
# edit hands back the SAME ref the model already holds.
existing = repo.get_artifact(str(artifact_id))
ref = _ref_for(
repo,
str(artifact_id),
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
metadata=(existing or {}).get("metadata"),
)
storage.save_file(io.BytesIO(data), storage_path)
saved_key = storage_path
except QuotaExceeded:
raise
except Exception:
logger.exception("artifacts_capture: failed to append artifact version")
_cleanup_orphan(storage, saved_key)
return None
payload: Dict[str, Any] = {
"artifact_id": str(artifact_id),
"version": version_number,
"filename": safe_name,
"mime_type": mime_type,
"size": size,
}
if ref is not None:
payload["ref"] = ref
return payload
+79
View File
@@ -0,0 +1,79 @@
"""Backend-agnostic code-execution sandbox interface and result types."""
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
@dataclass
class DisplayData:
"""A single rich (non-stream) kernel output keyed by MIME bundle."""
data: Dict[str, Any] = field(default_factory=dict)
metadata: Dict[str, Any] = field(default_factory=dict)
@dataclass
class Plot:
"""An image/plot output captured from execution, base64-encoded by format."""
format: str
content_base64: str
@dataclass
class ExecResult:
"""Outcome of one ``exec`` call against a sandbox session."""
status: str = "ok" # "ok" | "error"
stdout: str = ""
stderr: str = ""
exit_code: int = 0
execution_count: Optional[int] = None
error_name: Optional[str] = None
error_value: Optional[str] = None
traceback: List[str] = field(default_factory=list)
results: List[DisplayData] = field(default_factory=list)
display_data: List[DisplayData] = field(default_factory=list)
plots: List[Plot] = field(default_factory=list)
truncated: bool = False # output exceeded the budget and was cut; status stays "ok"
# The backend invalidated the runtime while producing this result. Managers
# must discard their cached handle so the next open performs a cold start.
runtime_invalidated: bool = False
@property
def ok(self) -> bool:
"""True when the execution completed without raising."""
return self.status == "ok"
class CodeSandbox(ABC):
"""Common interface every sandbox backend (Jupyter, Daytona, ...) implements."""
@abstractmethod
def open(self, session_id: str) -> str:
"""Create the underlying runtime for ``session_id`` and return its handle id."""
@abstractmethod
def attach(self, session_id: str) -> str:
"""Reattach to ``session_id``'s runtime; MAY return a cold kernel (state not guaranteed)."""
@abstractmethod
def close(self, session_id: str) -> None:
"""Tear down the runtime bound to ``session_id``, freeing all resources."""
@abstractmethod
def exec(self, session_id: str, code: str, timeout: Optional[float] = None) -> ExecResult:
"""Execute ``code`` in ``session_id``'s stateful runtime and return its result."""
@abstractmethod
def put_file(self, session_id: str, dest_path: str, data: bytes) -> None:
"""Write ``data`` to ``dest_path``; precondition: a relative path contained in the workspace."""
@abstractmethod
def get_file(self, session_id: str, path: str) -> bytes:
"""Read ``path`` from the workspace; precondition: a relative path contained in the workspace."""
@abstractmethod
def list_files(self, session_id: str) -> List[str]:
"""List workspace-relative file paths for ``session_id`` (never escapes the workspace)."""
+570
View File
@@ -0,0 +1,570 @@
"""Daytona Cloud sandbox: managed, strongly isolated runtimes via the Apache-2.0 Daytona SDK."""
import logging
import posixpath
import re
import threading
from typing import Dict, List, Optional
from application.sandbox.base import (
CodeSandbox,
ExecResult,
Plot,
)
logger = logging.getLogger(__name__)
# Per-session workspace root inside the Daytona sandbox. Relative file paths
# from LLM code and from put_file/get_file share this single directory.
_WORKSPACE_ROOT = "/home/daytona/docsgpt-sandbox"
# Session ids become filesystem path components and Daytona labels; restrict them.
_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9_-]+$")
# Label key used to find/reattach the Daytona sandbox bound to a DocsGPT session.
_SESSION_LABEL = "docsgpt_session_id"
class _Handle:
"""Tracks the Daytona sandbox object plus the workspace it executes in."""
def __init__(self, sandbox: object, sandbox_id: str, workspace: str) -> None:
self.sandbox = sandbox
self.sandbox_id = sandbox_id
self.workspace = workspace
class DaytonaSandbox(CodeSandbox):
"""Drives Daytona Cloud as a REST client; one managed sandbox per session.
Exec is STATELESS per call: Daytona's ``process.code_run`` runs each snippet
in a fresh Python interpreter, so Python variables and imports do NOT persist
across ``exec`` calls (unlike the Jupyter backend's stateful kernel). What
DOES persist is the sandbox FILESYSTEM: the per-session workspace and any
files written there survive between calls, so ``put_file``/``get_file`` and
files produced by one ``exec`` are visible to the next. ``attach`` therefore
always returns the same warm sandbox (filesystem state intact, interpreter
state lost). Matplotlib charts emitted by ``code_run`` are captured as PNG
plots. The app is a CLIENT of Daytona Cloud, authenticated by an API key.
Each cloud sandbox carries a ``docsgpt_session_id`` label. ``open``/``attach``
read that label to reattach to a still-live sandbox after a process restart
rather than creating a duplicate and orphaning the old (paid) one. A
``max_sandboxes`` cap bounds concurrent live sandboxes as a cost-DoS guard.
"""
def __init__(
self,
api_key: str,
api_url: Optional[str] = None,
target: Optional[str] = None,
snapshot: Optional[str] = None,
language: str = "python",
default_timeout: float = 60.0,
create_timeout: float = 60.0,
auto_stop_interval: int = 15,
auto_delete_interval: int = 60,
max_output_bytes: int = 0,
max_file_bytes: int = 10 * 1024 * 1024,
max_sandboxes: int = 50,
) -> None:
"""Configure the Daytona client; no cloud sandbox is created until ``open``."""
if not api_key:
raise ValueError("DAYTONA_API_KEY is required for the daytona sandbox backend")
# Imported lazily so app import never depends on the optional SDK.
from daytona import Daytona, DaytonaConfig
config_kwargs: Dict[str, object] = {"api_key": api_key}
if api_url:
config_kwargs["api_url"] = api_url
if target:
config_kwargs["target"] = target
self._client = Daytona(DaytonaConfig(**config_kwargs))
self._snapshot = snapshot
self._language = language
self._default_timeout = default_timeout
self._create_timeout = create_timeout
self._auto_stop_interval = auto_stop_interval
self._auto_delete_interval = auto_delete_interval
self._max_output_bytes = max_output_bytes
self._max_file_bytes = max_file_bytes
self._max_sandboxes = max_sandboxes
self._handles: Dict[str, _Handle] = {}
self._lock = threading.Lock()
# Session ids with a create/reattach in flight; a concurrent open() for the
# same id waits here and reuses the result, so two threads never create (and
# pay for) two cloud sandboxes for one session.
self._creating: set = set()
self._create_cv = threading.Condition(self._lock)
# -- Helpers ---------------------------------------------------------
@staticmethod
def _validate_session_id(session_id: str) -> None:
"""Reject session ids that would not be a safe path component or label value."""
if not _SESSION_ID_RE.match(session_id):
raise ValueError(f"Invalid session id {session_id!r}: expected [A-Za-z0-9_-]+")
def _get_handle(self, session_id: str) -> _Handle:
with self._lock:
handle = self._handles.get(session_id)
if handle is None:
raise KeyError(f"No sandbox session open for {session_id!r}")
return handle
@staticmethod
def _remote_path(workspace: str, rel_path: str) -> str:
"""Join a workspace-relative path, rejecting absolute paths, NUL/control chars, or traversal."""
if any(ord(ch) < 0x20 or ch == "\x7f" for ch in rel_path):
raise ValueError("path contains NUL or control characters")
if rel_path.startswith("/"):
raise ValueError("absolute paths are not allowed")
root = posixpath.normpath(workspace)
resolved = posixpath.normpath(posixpath.join(root, rel_path))
if resolved != root and not resolved.startswith(root + "/"):
raise ValueError("path escapes the session workspace")
return resolved
# -- Lifecycle -------------------------------------------------------
def open(self, session_id: str) -> str:
"""Reattach to or create the Daytona sandbox for ``session_id`` and prime its workspace.
Reattach order: an in-memory handle, then (across process restarts) a live cloud
sandbox carrying the session label, and only then a fresh ``create``. Enforces
``max_sandboxes`` so a flood of sessions cannot run up unbounded paid resources.
"""
self._validate_session_id(session_id)
# Wait out any in-flight create/reattach for this session, then reuse its
# handle; otherwise claim the in-flight slot so we are the sole creator.
with self._create_cv:
while session_id in self._creating:
self._create_cv.wait()
existing = self._handles.get(session_id)
if existing is not None:
return existing.sandbox_id
self._creating.add(session_id)
try:
# Cross-restart reattach: an earlier process may have created (and labelled)
# a sandbox for this session that is still live in the cloud. Reuse it
# instead of leaking it behind a brand-new create.
reattached = self._reattach_existing(session_id)
if reattached is not None:
self._prime(reattached)
return reattached.sandbox_id
with self._lock:
if len(self._handles) >= self._max_sandboxes:
raise RuntimeError(
f"Daytona sandbox cap reached ({self._max_sandboxes} live); refusing to create another"
)
sandbox = self._create_sandbox(session_id)
# Crash-safe register: if anything between create and registration raises,
# delete the just-created sandbox so it cannot orphan as a paid resource.
try:
sandbox_id = sandbox.id
handle = _Handle(sandbox, sandbox_id, _WORKSPACE_ROOT)
with self._lock:
self._handles[session_id] = handle
except Exception:
try:
self._client.delete(sandbox)
except Exception as del_exc: # noqa: BLE001 - cleanup is best-effort
logger.warning("Failed to delete orphaned Daytona sandbox during open: %s", del_exc)
raise
self._prime(handle)
return sandbox_id
finally:
with self._create_cv:
self._creating.discard(session_id)
self._create_cv.notify_all()
def _reattach_existing(self, session_id: str) -> Optional["_Handle"]:
"""Find a live cloud sandbox labelled for ``session_id`` and rebuild a handle from it.
Reads the ``docsgpt_session_id`` label written at create time so a process
restart reuses the existing sandbox rather than orphaning it. Returns ``None``
when no live sandbox matches.
"""
try:
from daytona import ListSandboxesQuery, SandboxState
query = ListSandboxesQuery(
labels={_SESSION_LABEL: session_id},
states=[SandboxState.STARTED, SandboxState.STOPPED],
)
matches = list(self._client.list(query))
except Exception as exc: # noqa: BLE001 - listing must never block opening a session
logger.warning("Daytona list for session %s failed; will create fresh: %s", session_id, exc)
return None
for sandbox in matches:
if getattr(sandbox, "labels", {}).get(_SESSION_LABEL) != session_id:
continue
if self._wake_if_stopped(sandbox) is None:
continue
handle = _Handle(sandbox, sandbox.id, _WORKSPACE_ROOT)
with self._lock:
self._handles[session_id] = handle
logger.info("Reattached to existing Daytona sandbox %s for session %s", sandbox.id, session_id)
return handle
return None
def _wake_if_stopped(self, sandbox: object) -> Optional[object]:
"""Start a stopped sandbox so it can serve exec/file ops; return None if it can't be woken."""
state = getattr(sandbox, "state", None)
state_value = getattr(state, "value", state)
if state_value == "started":
return sandbox
try:
self._client.start(sandbox, timeout=self._create_timeout)
return sandbox
except Exception as exc: # noqa: BLE001 - a sandbox we can't start is unusable
logger.warning("Failed to start stopped Daytona sandbox %s: %s", getattr(sandbox, "id", "?"), exc)
return None
def _ensure_started(self, handle: "_Handle") -> bool:
"""Refresh the handle's sandbox; start it if auto-stopped. True only if it was woken."""
try:
fresh = self._client.get(handle.sandbox_id)
except Exception as exc: # noqa: BLE001 - a failed refresh just means we don't retry
logger.warning("Daytona get for %s failed while ensuring started: %s", handle.sandbox_id, exc)
return False
state = getattr(fresh, "state", None)
state_value = getattr(state, "value", state)
if state_value == "started":
handle.sandbox = fresh
return False
if self._wake_if_stopped(fresh) is None:
return False
handle.sandbox = fresh
return True
def _create_sandbox(self, session_id: str):
"""Create a fresh Daytona sandbox labelled for ``session_id``."""
from daytona import CreateSandboxFromSnapshotParams
params_kwargs: Dict[str, object] = {
"language": self._language,
"labels": {_SESSION_LABEL: session_id},
"auto_stop_interval": self._auto_stop_interval,
"auto_delete_interval": self._auto_delete_interval,
}
if self._snapshot:
params_kwargs["snapshot"] = self._snapshot
params = CreateSandboxFromSnapshotParams(**params_kwargs)
return self._client.create(params, timeout=self._create_timeout)
def attach(self, session_id: str) -> str:
"""Reattach to the sandbox for ``session_id``; filesystem state is preserved.
Prefers the in-memory handle, then a live labelled cloud sandbox (so a process
restart does not orphan it), and only opens a fresh one as a last resort.
"""
self._validate_session_id(session_id)
with self._lock:
existing = self._handles.get(session_id)
if existing is not None:
return existing.sandbox_id
logger.warning("No live handle for session %s; reattaching or opening a Daytona sandbox", session_id)
return self.open(session_id)
def close(self, session_id: str) -> None:
"""Delete the Daytona sandbox for ``session_id`` so no cloud resource leaks."""
with self._lock:
handle = self._handles.pop(session_id, None)
if handle is None:
return
self._delete_sandbox(handle)
def close_handle(self, session_id: str, sandbox_id: str) -> None:
"""Delete the SPECIFIC sandbox captured at eviction time, never a re-opened one.
Used when the manager evicts a session and a concurrent ``open`` of the same id
may have already created a fresh sandbox: this tears down only the sandbox whose
id was captured, leaving any newer registered handle untouched. Falls back to the
registry only when nothing matches the captured id.
"""
with self._lock:
current = self._handles.get(session_id)
if current is not None and current.sandbox_id == sandbox_id:
# Captured handle is still the registered one: pop and delete it.
self._handles.pop(session_id, None)
victim = current
else:
# A concurrent re-open replaced the handle (or it is already gone); do
# NOT touch the registry. Delete the captured sandbox by id directly so
# the freshly created one survives.
victim = None
if victim is not None:
self._delete_sandbox(victim)
else:
self._delete_sandbox_by_id(sandbox_id)
def _delete_sandbox(self, handle: "_Handle") -> None:
"""Best-effort delete of the sandbox behind ``handle`` (never raises)."""
try:
self._client.delete(handle.sandbox)
except Exception as exc: # noqa: BLE001 - teardown is best-effort, never raise
logger.warning("Failed to delete Daytona sandbox %s: %s", handle.sandbox_id, exc)
def _delete_sandbox_by_id(self, sandbox_id: str) -> None:
"""Best-effort delete a sandbox by its id when no live handle is held (never raises).
``client.delete`` needs a ``Sandbox`` object, so fetch it by id first; this is the
evict-then-concurrent-reopen path where only the captured id is still known.
"""
try:
sandbox = self._client.get(sandbox_id)
self._client.delete(sandbox)
except Exception as exc: # noqa: BLE001 - teardown is best-effort, never raise
logger.warning("Failed to delete Daytona sandbox by id %s: %s", sandbox_id, exc)
def remove_path(self, session_id: str, path: str) -> None:
"""Best-effort delete a workspace-relative path (per-render scratch dir); never raises.
The tools pass a token DIRECTORY (e.g. ``artifacts/{token}``), so the delete is
recursive — ``fs.delete_file`` with ``recursive=True`` removes a folder and its
contents (daytona==0.190.1).
"""
try:
handle = self._get_handle(session_id)
remote = self._remote_path(handle.workspace, path)
if remote == posixpath.normpath(handle.workspace):
return # refuse to delete the workspace root itself
self._delete_remote(handle, remote)
except Exception as exc: # noqa: BLE001 - cleanup is best-effort, never raise
logger.debug("remove_path best-effort delete returned for %r: %s", path, exc)
def _delete_remote(self, handle: "_Handle", remote: str) -> None:
"""Recursively delete ``remote`` inside the sandbox, with a contained-command fallback."""
try:
handle.sandbox.fs.delete_file(remote, recursive=True)
return
except TypeError:
# Older SDK without the recursive kwarg: fall back to a contained rm -rf.
pass
try:
handle.sandbox.fs.delete_file(remote)
except Exception: # noqa: BLE001 - directory delete may need a recursive remove
quoted = "'" + remote.replace("'", "'\\''") + "'"
handle.sandbox.process.exec(f"rm -rf {quoted}", timeout=int(self._default_timeout))
def _prime(self, handle: _Handle) -> None:
"""Create the per-session workspace directory inside the sandbox."""
try:
handle.sandbox.fs.create_folder(handle.workspace, "755")
except Exception as exc: # noqa: BLE001 - folder may already exist on a reused snapshot
logger.debug("Workspace folder prime returned: %s", exc)
# -- Execution -------------------------------------------------------
def exec(self, session_id: str, code: str, timeout: Optional[float] = None) -> ExecResult:
"""Run ``code`` via Daytona ``code_run``; per-call interpreter, persistent filesystem."""
handle = self._get_handle(session_id)
wall = int(timeout or self._default_timeout)
wrapped = self._with_workspace_cwd(handle.workspace, code)
try:
response = handle.sandbox.process.code_run(wrapped, timeout=wall)
except Exception as exc: # noqa: BLE001 - any SDK/cloud error -> error result, never raise
# A cached handle may point at a sandbox Daytona auto-stopped; wake it and
# retry once. Genuine code errors return a nonzero-exit response (they do
# NOT raise), so this only retries transport/stopped faults.
if self._ensure_started(handle):
try:
return self._to_result(handle.sandbox.process.code_run(wrapped, timeout=wall))
except Exception: # noqa: BLE001 - second failure -> error result below
pass
return ExecResult(
status="error",
error_name=type(exc).__name__,
error_value=str(exc) or "code_run failed",
exit_code=-1,
)
return self._to_result(response)
@staticmethod
def _with_workspace_cwd(workspace: str, code: str) -> str:
"""Prepend a chdir into the session workspace so relative paths resolve there.
Leading ``from __future__`` imports are hoisted above the prelude so they stay
the first statements of the module (Python rejects them anywhere else).
"""
hoisted, rest = DaytonaSandbox._split_leading_future_imports(code)
prelude = (
"import os as _os\n"
f"_os.makedirs({workspace!r}, exist_ok=True)\n"
f"_os.chdir({workspace!r})\n"
)
return hoisted + prelude + rest
@staticmethod
def _split_leading_future_imports(code: str) -> tuple[str, str]:
"""Split leading ``from __future__`` imports (and the blank/comment lines around them) from the rest.
A module docstring appearing BEFORE a future import is an unsupported edge (rare in
generated snippets); the common ``from __future__ import annotations`` first-line case
is handled. Everything from the first real statement onward stays in ``rest``.
"""
lines = code.splitlines(keepends=True)
saw_future = False
split_at = 0
for i, line in enumerate(lines):
stripped = line.strip()
if stripped == "" or stripped.startswith("#"):
continue # blank/comment: part of the leading run
if stripped.startswith("from __future__ import"):
saw_future = True
continue
split_at = i # first real statement: the leading run ends here
break
else:
split_at = len(lines) # whole snippet was blank/comment/future
if not saw_future:
return "", code
hoisted = "".join(lines[:split_at])
if hoisted and not hoisted.endswith("\n"):
hoisted += "\n"
return hoisted, "".join(lines[split_at:])
def _to_result(self, response) -> ExecResult:
"""Map a Daytona ``ExecuteResponse`` into the shared ``ExecResult`` shape.
Caps ``stdout`` at ``max_output_bytes`` (0 = disabled) BEFORE it is also reused as
``error_value`` so a huge buffered response cannot propagate unbounded downstream.
"""
exit_code = getattr(response, "exit_code", 0) or 0
artifacts = getattr(response, "artifacts", None)
stdout = ""
if artifacts is not None and getattr(artifacts, "stdout", None) is not None:
stdout = artifacts.stdout
else:
stdout = getattr(response, "result", "") or ""
truncated = False
if self._max_output_bytes and len(stdout.encode("utf-8", "ignore")) > self._max_output_bytes:
stdout = stdout.encode("utf-8", "ignore")[: self._max_output_bytes].decode("utf-8", "ignore")
stdout += f"\n[output truncated at {self._max_output_bytes} bytes]"
truncated = True
result = ExecResult(
status="ok" if exit_code == 0 else "error",
stdout=stdout,
exit_code=exit_code,
truncated=truncated,
)
if exit_code != 0:
result.error_name = "ExecutionError"
result.error_value = stdout or f"exited with code {exit_code}"
if artifacts is not None:
for chart in getattr(artifacts, "charts", None) or []:
png = getattr(chart, "png", None)
if png:
result.plots.append(Plot(format="png", content_base64=png))
return result
# -- File transfer ---------------------------------------------------
def put_file(self, session_id: str, dest_path: str, data: bytes) -> None:
"""Upload ``data`` to ``dest_path`` under the session workspace, creating parent dirs."""
handle = self._get_handle(session_id)
remote = self._remote_path(handle.workspace, dest_path)
parent = posixpath.dirname(remote)
try:
self._upload(handle, remote, parent, data)
except Exception as exc: # noqa: BLE001 - log detail server-side, return a generic error
# A cached handle may point at an auto-stopped sandbox; wake it and retry once.
if self._ensure_started(handle):
try:
self._upload(handle, remote, parent, data)
return
except Exception: # noqa: BLE001 - second failure -> generic IOError below
pass
logger.warning("put_file failed for %r: %s", dest_path, exc)
raise IOError(f"put_file failed: {type(exc).__name__}") from exc
def _upload(self, handle: "_Handle", remote: str, parent: str, data: bytes) -> None:
"""Create the parent folder (best-effort) and upload ``data`` to ``remote``."""
if parent and parent != handle.workspace:
try:
handle.sandbox.fs.create_folder(parent, "755")
except Exception as folder_exc: # noqa: BLE001 - folder may already exist
logger.debug("put_file parent folder create returned: %s", folder_exc)
handle.sandbox.fs.upload_file(data, remote)
def get_file(self, session_id: str, path: str) -> bytes:
"""Download ``path`` from the session workspace as bytes, capped at ``max_file_bytes``."""
handle = self._get_handle(session_id)
remote = self._remote_path(handle.workspace, path)
try:
data = self._download(handle, remote)
except IOError:
raise
except Exception as exc: # noqa: BLE001 - log detail server-side, return a generic error
# A cached handle may point at an auto-stopped sandbox; wake it and retry once.
if not self._ensure_started(handle):
logger.warning("get_file failed for %r: %s", path, exc)
raise IOError(f"get_file failed: {type(exc).__name__}") from exc
try:
data = self._download(handle, remote)
except IOError:
raise
except Exception: # noqa: BLE001 - second failure -> generic IOError
logger.warning("get_file failed for %r: %s", path, exc)
raise IOError(f"get_file failed: {type(exc).__name__}") from exc
if data is None:
raise IOError(f"get_file produced no payload for {path!r}")
data = data if isinstance(data, bytes) else bytes(data)
# The pre-download size guard may be skipped when get_file_info has no size;
# enforce the cap against the actual payload so an oversized file never slips through.
if len(data) > self._max_file_bytes:
raise IOError(f"file too large: {len(data)} > {self._max_file_bytes} bytes")
return data
def _download(self, handle: "_Handle", remote: str) -> object:
"""Fetch ``remote``'s bytes, rejecting a file whose declared size exceeds the cap."""
info = handle.sandbox.fs.get_file_info(remote)
size = getattr(info, "size", None)
if size is not None and size > self._max_file_bytes:
raise IOError(f"file too large: {size} > {self._max_file_bytes} bytes")
return handle.sandbox.fs.download_file(remote)
def list_files(self, session_id: str) -> List[str]:
"""List workspace-relative file paths for ``session_id`` (recursive, never escapes the workspace)."""
handle = self._get_handle(session_id)
try:
return self._list_all(handle)
except Exception: # noqa: BLE001 - transport/stopped fault: wake and retry once, else re-raise
# A cached handle may point at an auto-stopped sandbox; wake it and retry once.
if self._ensure_started(handle):
return self._list_all(handle)
raise
def _list_all(self, handle: "_Handle") -> List[str]:
"""Walk the workspace subtree for ``handle`` and return workspace-relative file paths."""
out: List[str] = []
self._walk(handle.sandbox, handle.workspace, "", out)
return out
def _walk(self, sandbox: object, root: str, rel_dir: str, out: List[str]) -> None:
"""Recurse one workspace subtree, appending workspace-relative file paths to ``out``."""
abs_dir = posixpath.join(root, rel_dir) if rel_dir else root
try:
entries = sandbox.fs.list_files(abs_dir)
except Exception as exc: # noqa: BLE001 - log detail server-side, return a generic error
logger.warning("list_files failed for %r: %s", rel_dir or ".", exc)
raise IOError(f"list_files failed: {type(exc).__name__}") from exc
for entry in entries or []:
name = getattr(entry, "name", None)
if not name:
continue
child_rel = posixpath.join(rel_dir, name) if rel_dir else name
# Defend against a backend returning ".."/absolute entries that would escape root.
resolved = posixpath.normpath(posixpath.join(root, child_rel))
if resolved != root and not resolved.startswith(root + "/"):
continue
if getattr(entry, "is_dir", False):
self._walk(sandbox, root, child_rel, out)
else:
out.append(child_rel)
+676
View File
@@ -0,0 +1,676 @@
"""Jupyter Kernel Gateway sandbox: stateful in-process kernels over REST + WebSocket."""
import base64
import hashlib
import json
import logging
import re
import threading
import time
import uuid
from typing import Dict, List, Optional
from urllib.parse import urlparse, urlunparse
import requests
import websocket
from application.sandbox.base import (
CodeSandbox,
DisplayData,
ExecResult,
Plot,
)
logger = logging.getLogger(__name__)
# Per-session workspace root inside the runner container. The kernel sets its
# cwd here so relative paths from LLM code and file in/out share one directory.
_WORKSPACE_ROOT = "/tmp/docsgpt-sandbox" # nosec B108 - controlled per-session sandbox workspace dir
# Marker the get_file helper prints around the base64 payload so we can extract
# it from interleaved stdout without a contents API (the gateway has none).
_FILE_BEGIN = "<<<DOCSGPT_FILE_BEGIN>>>"
_FILE_END = "<<<DOCSGPT_FILE_END>>>"
# Session ids become filesystem path components; allow only safe characters.
_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9_-]+$")
# Kernel-side helper that resolves a workspace-relative path and rejects any
# absolute path or one that escapes the per-session workspace (path traversal).
_CONTAINMENT_SNIPPET = (
"def _resolve(_base, _rel):\n"
" import os as _os\n"
" if _os.path.isabs(_rel):\n"
" raise ValueError('absolute paths are not allowed')\n"
" _root = _os.path.realpath(_base)\n"
" _rp = _os.path.realpath(_os.path.join(_root, _rel))\n"
" if _rp != _root and not _rp.startswith(_root + _os.sep):\n"
" raise ValueError('path escapes the session workspace')\n"
" return _rp\n"
)
class _Kernel:
"""Tracks one gateway kernel plus the per-session workspace it executes in."""
def __init__(
self, kernel_id: str, workspace: str, session_id: Optional[str] = None
) -> None:
self.kernel_id = kernel_id
self.workspace = workspace
self.session_id = session_id
self.initialized = False
class JupyterKernelGatewaySandbox(CodeSandbox):
"""Drives one always-on Jupyter Kernel Gateway, one stateful kernel per session."""
# Raw-byte chunk size for a staged upload: each execute_request carries one
# chunk's base64 (~4 MB at 3 MB raw), well under the gateway's 10 MiB websocket
# frame cap. The 3-byte boundary keeps every base64 block self-contained.
_PUT_CHUNK_BYTES = 3 * 1024 * 1024
def __init__(
self,
gateway_url: str,
auth_token: Optional[str] = None,
kernel_name: str = "python3",
default_timeout: float = 60.0,
http_timeout: float = 10.0,
max_output_bytes: int = 8 * 1024 * 1024,
max_file_bytes: int = 10 * 1024 * 1024,
) -> None:
"""Configure the client; no kernel is created until ``open`` is called."""
self._base_url = gateway_url.rstrip("/")
self._auth_token = auth_token
self._kernel_name = kernel_name
self._default_timeout = default_timeout
self._http_timeout = http_timeout
self._max_output_bytes = max_output_bytes
self._max_file_bytes = max_file_bytes
self._kernels: Dict[str, _Kernel] = {}
# Timed-out kernels whose DELETE could not yet be confirmed. Immutable
# ids stay retryable here and are never reused as active runtimes.
self._quarantined_kernels: Dict[str, Optional[str]] = {}
self._lock = threading.Lock()
# Session ids with a create in flight; a second open() for the same id
# waits on this CV and reuses the result instead of double-creating a
# kernel (the idempotency contract SandboxManager.open relies on).
self._creating: set = set()
self._create_cv = threading.Condition(self._lock)
# -- HTTP helpers ----------------------------------------------------
def _headers(self) -> Dict[str, str]:
headers = {"Content-Type": "application/json"}
if self._auth_token:
headers["Authorization"] = f"token {self._auth_token}"
return headers
def _ws_url(self, kernel_id: str) -> str:
# The token is sent only via the Authorization header (_ws_headers); keeping it
# out of the URL query avoids leaking it into gateway/proxy access logs.
parsed = urlparse(self._base_url)
scheme = "wss" if parsed.scheme == "https" else "ws"
path = f"/api/kernels/{kernel_id}/channels"
return urlunparse((scheme, parsed.netloc, path, "", "", ""))
def _get_kernel(self, session_id: str) -> _Kernel:
# A manager may cache a replacement handle, so normal traffic is also
# an opportunity to retire an older quarantined kernel.
self._retry_quarantined_kernels(session_id)
with self._lock:
kernel = self._kernels.get(session_id)
if kernel is None:
raise KeyError(f"No sandbox session open for {session_id!r}")
return kernel
@staticmethod
def _validate_session_id(session_id: str) -> None:
"""Reject session ids that would not be a safe filesystem path component."""
if not _SESSION_ID_RE.match(session_id):
raise ValueError(f"Invalid session id {session_id!r}: expected [A-Za-z0-9_-]+")
# -- Lifecycle -------------------------------------------------------
def open(self, session_id: str) -> str:
"""Start a fresh kernel for ``session_id`` and prime its workspace cwd.
Idempotent under concurrency: if another thread is already creating a kernel
for this session, wait for it and reuse the result rather than POSTing a
second kernel that would orphan on the gateway.
"""
self._validate_session_id(session_id)
with self._create_cv:
while session_id in self._creating:
self._create_cv.wait()
existing = self._kernels.get(session_id)
has_quarantine = any(
owner == session_id for owner in self._quarantined_kernels.values()
)
if existing is not None and not has_quarantine:
return existing.kernel_id
self._creating.add(session_id)
try:
unresolved = self._retry_quarantined_kernels(session_id)
with self._lock:
# A replacement may already exist. It is safe to use even when
# deletion of an older exact id remains pending.
existing = self._kernels.get(session_id)
if existing is not None:
return existing.kernel_id
if unresolved:
raise RuntimeError(
f"Previous timed-out kernel for {session_id!r} could not be terminated"
)
resp = requests.post(
f"{self._base_url}/api/kernels",
headers=self._headers(),
data=json.dumps({"name": self._kernel_name}),
timeout=self._http_timeout,
)
resp.raise_for_status()
kernel_id = resp.json()["id"]
workspace = f"{_WORKSPACE_ROOT}/{session_id}"
kernel = _Kernel(kernel_id, workspace, session_id)
with self._lock:
self._kernels[session_id] = kernel
self._prime(kernel)
return kernel_id
finally:
with self._create_cv:
self._creating.discard(session_id)
self._create_cv.notify_all()
def attach(self, session_id: str) -> str:
"""Reattach to a still-running kernel for ``session_id``; open a cold one if gone."""
self._validate_session_id(session_id)
unresolved = self._retry_quarantined_kernels(session_id)
with self._lock:
existing = self._kernels.get(session_id)
if existing is not None and self._kernel_alive(existing.kernel_id):
return existing.kernel_id
if existing is not None:
with self._lock:
self._kernels.pop(session_id, None)
if unresolved:
raise RuntimeError(
f"Previous timed-out kernel for {session_id!r} could not be terminated"
)
logger.warning("Re-attaching session %s to a cold kernel; previous state is lost", session_id)
return self.open(session_id)
def close(self, session_id: str) -> None:
"""Sweep the session workspace (best-effort) then delete its kernel and registry entry."""
with self._lock:
kernel = self._kernels.pop(session_id, None)
if kernel is None:
return
self._cleanup_workspace(kernel)
self._delete_kernel(kernel.kernel_id)
def _cleanup_workspace(self, kernel: _Kernel) -> None:
"""Best-effort rmtree of the per-session workspace while the kernel is still alive."""
if not kernel.initialized:
return
code = "import shutil as _sh\n" f"_sh.rmtree({kernel.workspace!r}, ignore_errors=True)\n"
try:
self._run(kernel, code, self._http_timeout)
except Exception: # noqa: BLE001 - teardown is best-effort and must never raise
logger.warning("Failed to sweep workspace for kernel %s", kernel.kernel_id, exc_info=True)
def close_handle(self, session_id: str, kernel_id: str) -> None:
"""Delete the SPECIFIC kernel captured at eviction time, never a re-opened one.
When the manager evicts a session, a concurrent ``open`` of the same id may have
already started a fresh kernel; this deletes only the kernel whose id was captured
and pops the registry entry only when it still points at that same kernel, so the
new kernel survives.
"""
with self._lock:
current = self._kernels.get(session_id)
if current is not None and current.kernel_id == kernel_id:
self._kernels.pop(session_id, None)
self._delete_kernel(kernel_id)
def _delete_kernel(self, kernel_id: str) -> bool:
"""Best-effort DELETE of a gateway kernel, retrying once and never raising."""
last_error = "unknown error"
for _attempt in range(2):
try:
resp = requests.delete(
f"{self._base_url}/api/kernels/{kernel_id}",
headers=self._headers(),
timeout=self._http_timeout,
)
except requests.RequestException as exc: # teardown is best-effort
last_error = str(exc)
continue
# A missing kernel is already in the desired state.
if 200 <= resp.status_code < 300 or resp.status_code == 404:
return True
last_error = f"HTTP {resp.status_code}"
logger.warning("Failed to delete kernel %s after retry: %s", kernel_id, last_error)
return False
def _kernel_alive(self, kernel_id: str) -> bool:
try:
resp = requests.get(
f"{self._base_url}/api/kernels/{kernel_id}",
headers=self._headers(),
timeout=self._http_timeout,
)
return resp.status_code == 200
except requests.RequestException:
return False
def _interrupt(self, kernel_id: str) -> None:
"""Best-effort interrupt so a timed-out/runaway kernel becomes reusable."""
try:
requests.post(
f"{self._base_url}/api/kernels/{kernel_id}/interrupt",
headers=self._headers(),
timeout=self._http_timeout,
)
except requests.RequestException as exc:
logger.warning("Failed to interrupt kernel %s: %s", kernel_id, exc)
def _interrupt_and_drain(self, ws: websocket.WebSocket, msg_id: str, kernel_id: str) -> bool:
"""Interrupt and drain the request, returning whether matching idle was observed."""
self._interrupt(kernel_id)
drain_deadline = time.monotonic() + self._http_timeout
while time.monotonic() < drain_deadline:
try:
ws.settimeout(max(0.05, drain_deadline - time.monotonic()))
raw = ws.recv()
except (websocket.WebSocketTimeoutException, websocket.WebSocketConnectionClosedException):
return False
if not raw:
continue
try:
msg = json.loads(raw)
except json.JSONDecodeError:
continue
if msg.get("parent_header", {}).get("msg_id") != msg_id:
continue
msg_type = msg.get("msg_type") or msg.get("header", {}).get("msg_type")
if msg_type == "status" and msg.get("content", {}).get("execution_state") == "idle":
return True
return False
def _retry_quarantined_kernels(self, session_id: str) -> List[str]:
"""Retry pending exact-id deletes for this session; return unresolved ids."""
with self._lock:
pending = [
(kernel_id, owner)
for kernel_id, owner in self._quarantined_kernels.items()
if owner == session_id
]
unresolved: List[str] = []
for kernel_id, owner in pending:
if self._delete_kernel(kernel_id):
with self._lock:
if (
kernel_id in self._quarantined_kernels
and self._quarantined_kernels[kernel_id] == owner
):
self._quarantined_kernels.pop(kernel_id, None)
else:
unresolved.append(kernel_id)
return unresolved
def _invalidate_kernel(
self, kernel_id: str, session_id: Optional[str] = None
) -> bool:
"""Quarantine and hard-delete one exact kernel without touching a replacement."""
with self._lock:
stale_sessions = [
session_id
for session_id, kernel in self._kernels.items()
if kernel.kernel_id == kernel_id
]
owner = session_id or (stale_sessions[0] if len(stale_sessions) == 1 else None)
# Register before eviction so a failed DELETE never loses the only
# retryable reference to this runaway kernel id.
self._quarantined_kernels[kernel_id] = owner
for session_id in stale_sessions:
self._kernels.pop(session_id, None)
deleted = self._delete_kernel(kernel_id)
if deleted:
with self._lock:
self._quarantined_kernels.pop(kernel_id, None)
return deleted
def _prime(self, kernel: _Kernel) -> None:
"""Create the per-session workspace (mode 0700) and chdir the kernel into it."""
# 0700 on the root and the per-session dir is defense-in-depth only: every
# kernel runs under one shared uid here, so this is not a cross-session
# boundary (that needs distinct uids / per-session VMs -- the Daytona backend).
# A fresh kernel always starts on a clean workspace: rmtree any stale dir a
# prior kernel for the same session id left behind (else artifacts_capture
# would re-read those files every exec). Only genuine new-kernel creation
# reaches here -- open() on a live kernel returns early -- so a warm
# persist=true session is never wiped mid-computation.
setup = (
"import os as _os, shutil as _sh\n"
f"_os.makedirs({_WORKSPACE_ROOT!r}, mode=0o700, exist_ok=True)\n"
f"_os.chmod({_WORKSPACE_ROOT!r}, 0o700)\n"
f"_sh.rmtree({kernel.workspace!r}, ignore_errors=True)\n"
f"_os.makedirs({kernel.workspace!r}, mode=0o700, exist_ok=True)\n"
f"_os.chmod({kernel.workspace!r}, 0o700)\n"
f"_os.chdir({kernel.workspace!r})\n"
)
result = self._run(kernel, setup, self._default_timeout)
if not result.ok:
raise RuntimeError(f"Sandbox workspace setup failed: {result.error_value}")
kernel.initialized = True
# -- Execution -------------------------------------------------------
def exec(self, session_id: str, code: str, timeout: Optional[float] = None) -> ExecResult:
"""Run ``code`` in the session's persistent kernel; state carries across calls."""
kernel = self._get_kernel(session_id)
return self._run(kernel, code, timeout or self._default_timeout)
def _run(
self,
kernel: _Kernel,
code: str,
timeout: float,
max_output_bytes: Optional[int] = None,
) -> ExecResult:
"""Execute one ``execute_request`` over the WS channel and assemble the reply.
``max_output_bytes`` overrides the default output budget for this call only
(file-transfer execs raise it so a multi-MB base64 payload is not truncated).
"""
try:
ws = websocket.create_connection(
self._ws_url(kernel.kernel_id),
timeout=timeout,
header=self._ws_headers(),
)
except Exception as exc: # noqa: BLE001 - connect failure -> error result, never raise
return _error_result(type(exc).__name__, str(exc) or "failed to open kernel channel")
try:
msg_id = uuid.uuid4().hex
ws.send(json.dumps(self._execute_request(msg_id, code)))
return self._collect(
ws,
msg_id,
timeout,
kernel.kernel_id,
max_output_bytes,
session_id=kernel.session_id,
)
finally:
try:
ws.close()
except Exception: # noqa: BLE001 - closing a socket must not mask results
pass
def _ws_headers(self) -> List[str]:
if self._auth_token:
return [f"Authorization: token {self._auth_token}"]
return []
@staticmethod
def _execute_request(msg_id: str, code: str) -> dict:
"""Build a Jupyter ``execute_request`` wire message for the shell channel."""
return {
"header": {
"msg_id": msg_id,
"username": "docsgpt",
"session": uuid.uuid4().hex,
"msg_type": "execute_request",
"version": "5.3",
},
"parent_header": {},
"metadata": {},
"content": {
"code": code,
"silent": False,
"store_history": True,
"user_expressions": {},
"allow_stdin": False,
"stop_on_error": True,
},
"channel": "shell",
}
def _collect(
self,
ws: websocket.WebSocket,
msg_id: str,
timeout: float,
kernel_id: str,
max_output_bytes: Optional[int] = None,
session_id: Optional[str] = None,
) -> ExecResult:
"""Read iopub/shell frames until ``execute_reply``/idle, a wall-clock deadline, or a closed socket."""
effective = max_output_bytes if max_output_bytes is not None else self._max_output_bytes
result = ExecResult()
stdout_parts: List[str] = []
stderr_parts: List[str] = []
buffered = 0
truncated = False
reply_seen = False
idle_seen = False
deadline = time.monotonic() + timeout
while not (reply_seen and idle_seen):
now = time.monotonic()
remaining = deadline - now
if remaining <= 0:
self._fail(result, "TimeoutError", f"execution exceeded {timeout}s")
if not self._interrupt_and_drain(ws, msg_id, kernel_id):
self._invalidate_kernel(kernel_id, session_id)
result.runtime_invalidated = True
break
try:
ws.settimeout(remaining)
raw = ws.recv()
except websocket.WebSocketTimeoutException:
self._fail(result, "TimeoutError", f"execution exceeded {timeout}s")
if not self._interrupt_and_drain(ws, msg_id, kernel_id):
self._invalidate_kernel(kernel_id, session_id)
result.runtime_invalidated = True
break
except websocket.WebSocketConnectionClosedException:
self._fail(result, "KernelDiedError", "kernel channel closed before completion")
break
if not raw:
continue
try:
msg = json.loads(raw)
except json.JSONDecodeError as exc:
self._fail(result, "ProtocolError", f"malformed kernel frame: {exc}")
break
if msg.get("parent_header", {}).get("msg_id") != msg_id:
continue
msg_type = msg.get("msg_type") or msg.get("header", {}).get("msg_type")
content = msg.get("content", {})
if msg_type == "stream":
if not truncated:
text = content.get("text", "")
buffered += len(text.encode("utf-8", "ignore"))
if buffered > effective:
truncated = True
self._interrupt_and_drain(ws, msg_id, kernel_id) # runaway output: stop and drain
break
elif content.get("name") == "stderr":
stderr_parts.append(text)
else:
stdout_parts.append(text)
elif msg_type in ("execute_result", "display_data"):
# Rich outputs (results/display_data/plots) count against the SAME
# byte budget as streams: untrusted code can emit a huge DataFrame /
# IPython.display.HTML / many large images that the backend would
# otherwise buffer unbounded. Drop and truncate once over budget.
if not truncated:
buffered += self._rich_payload_bytes(content)
if buffered > effective:
truncated = True
self._interrupt_and_drain(ws, msg_id, kernel_id)
break
self._capture_rich(result, msg_type, content)
elif msg_type == "error":
result.status = "error"
result.exit_code = 1
result.error_name = content.get("ename")
result.error_value = content.get("evalue")
result.traceback = content.get("traceback", [])
elif msg_type == "execute_reply":
reply_seen = True
result.execution_count = content.get("execution_count")
if content.get("status") == "error":
result.status = "error"
result.exit_code = result.exit_code or 1
result.error_name = result.error_name or content.get("ename")
result.error_value = result.error_value or content.get("evalue")
if not result.traceback:
result.traceback = content.get("traceback", [])
elif msg_type == "status":
if content.get("execution_state") == "idle":
idle_seen = True
if truncated:
stderr_parts.append(f"\n[output truncated at {effective} bytes]")
result.truncated = True # surface the cut without flipping status off "ok"
result.stdout = "".join(stdout_parts)
result.stderr = "".join(stderr_parts)
return result
@staticmethod
def _fail(result: ExecResult, name: str, value: str) -> None:
"""Mark ``result`` as a failed exec with the given error name/value."""
result.status = "error"
result.error_name = name
result.error_value = value
result.exit_code = -1
@staticmethod
def _rich_payload_bytes(content: dict) -> int:
"""Approximate the serialized byte size of a rich-output bundle's data payloads."""
total = 0
for value in (content.get("data") or {}).values():
if isinstance(value, str):
total += len(value.encode("utf-8", "ignore"))
else:
try:
total += len(json.dumps(value, default=str).encode("utf-8", "ignore"))
except Exception:
total += len(str(value).encode("utf-8", "ignore"))
return total
@staticmethod
def _capture_rich(result: ExecResult, msg_type: str, content: dict) -> None:
"""Sort a rich output into results/display_data and pull out any image plots."""
data = content.get("data", {}) or {}
metadata = content.get("metadata", {}) or {}
bundle = DisplayData(data=data, metadata=metadata)
if msg_type == "execute_result":
result.results.append(bundle)
else:
result.display_data.append(bundle)
for mime, payload in data.items():
if mime.startswith("image/"):
result.plots.append(Plot(format=mime.split("/", 1)[1], content_base64=payload))
# -- File transfer ---------------------------------------------------
def _file_transfer_budget(self) -> int:
"""Output budget for a file-transfer exec: max file bytes inflated by base64 plus marker slack."""
return self._max_file_bytes * 4 // 3 + 4096
def put_file(self, session_id: str, dest_path: str, data: bytes) -> None:
"""Decode ``data`` inside the kernel and write it under the session workspace.
The upload is chunked so no single ``execute_request`` exceeds the gateway's
websocket message-size cap: the first chunk creates/truncates the file (``wb``)
and each later chunk appends (``ab``). Every chunk program re-resolves the path
-- kernel globals are not relied on to survive between execs.
"""
kernel = self._get_kernel(session_id)
offset = 0
first = True
# Loop at least once so an empty file is still created (a single wb write of b"").
while first or offset < len(data):
chunk = data[offset:offset + self._PUT_CHUNK_BYTES]
encoded = base64.b64encode(chunk).decode("ascii")
if first:
body = (
f"_p = _resolve({kernel.workspace!r}, {dest_path!r})\n"
"_os.makedirs(_os.path.dirname(_p) or '.', exist_ok=True)\n"
f"_f = open(_p, 'wb'); _f.write(_b64.b64decode({encoded!r})); _f.close()\n"
)
else:
body = (
f"_p = _resolve({kernel.workspace!r}, {dest_path!r})\n"
f"_f = open(_p, 'ab'); _f.write(_b64.b64decode({encoded!r})); _f.close()\n"
)
code = "import base64 as _b64, os as _os\n" + _CONTAINMENT_SNIPPET + body
result = self._run(kernel, code, self._default_timeout)
if not result.ok:
raise IOError(f"put_file failed: {result.error_value}")
offset += self._PUT_CHUNK_BYTES
first = False
def get_file(self, session_id: str, path: str) -> bytes:
"""Read ``path`` inside the kernel and stream its base64 (with a length tag) over stdout."""
kernel = self._get_kernel(session_id)
code = (
"import base64 as _b64, hashlib as _hl, os as _os\n"
+ _CONTAINMENT_SNIPPET
+ f"_p = _resolve({kernel.workspace!r}, {path!r})\n"
f"_sz = _os.path.getsize(_p)\n"
f"if _sz > {self._max_file_bytes}:\n"
f" raise ValueError('file too large: %d > {self._max_file_bytes} bytes' % _sz)\n"
"_d = open(_p, 'rb').read()\n"
"_h = _hl.sha256(_d).hexdigest()\n"
f"print({_FILE_BEGIN!r} + str(len(_d)) + ':' + _h + ':'"
f" + _b64.b64encode(_d).decode('ascii') + {_FILE_END!r})\n"
)
result = self._run(kernel, code, self._default_timeout, max_output_bytes=self._file_transfer_budget())
if not result.ok:
raise IOError(f"get_file failed: {result.error_value}")
out = result.stdout
start = out.find(_FILE_BEGIN)
end = out.find(_FILE_END)
if start == -1 or end == -1:
raise IOError(f"get_file produced no payload for {path!r}")
payload = out[start + len(_FILE_BEGIN):end]
expected_len_s, expected_sha, encoded = payload.split(":", 2)
decoded = base64.b64decode(encoded)
if len(decoded) != int(expected_len_s) or hashlib.sha256(decoded).hexdigest() != expected_sha:
raise IOError(f"get_file integrity check failed for {path!r} (payload truncated)")
return decoded
def list_files(self, session_id: str) -> List[str]:
"""Walk the session workspace inside the kernel and return relative paths."""
kernel = self._get_kernel(session_id)
code = (
"import os as _os, json as _json\n"
f"_root = {kernel.workspace!r}\n"
"_out = []\n"
"for _dp, _dn, _fn in _os.walk(_root):\n"
" for _name in _fn:\n"
" _out.append(_os.path.relpath(_os.path.join(_dp, _name), _root))\n"
f"print({_FILE_BEGIN!r} + _json.dumps(_out) + {_FILE_END!r})\n"
)
result = self._run(kernel, code, self._default_timeout, max_output_bytes=self._file_transfer_budget())
if not result.ok:
raise IOError(f"list_files failed: {result.error_value}")
out = result.stdout
start = out.find(_FILE_BEGIN)
end = out.find(_FILE_END)
if start == -1 or end == -1:
return []
return json.loads(out[start + len(_FILE_BEGIN):end])
def _error_result(name: str, value: str) -> ExecResult:
"""Build a failed ExecResult carrying the given error name/value."""
return ExecResult(status="error", error_name=name, error_value=value, exit_code=-1)
+391
View File
@@ -0,0 +1,391 @@
"""Process-wide registry binding session ids to sandbox handles with idle expiry."""
import logging
import threading
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
from application.sandbox.base import CodeSandbox, ExecResult
logger = logging.getLogger(__name__)
class SandboxCapacityError(RuntimeError):
"""Raised when the concurrent-session cap is reached and no idle session can be freed."""
@dataclass
class _Session:
"""Bookkeeping for one bound sandbox session: its TTL, access timestamps, and backend handle.
``handle`` is the backend handle id returned by ``backend.open``; ``None`` while a
slot is RESERVED (a placeholder occupying a cap slot during a cold backend open that
runs outside the lock). ``ready`` is False for such a placeholder so reuse/reap/evict
skip it until the backend open finalizes it. ``pending_close`` marks a session whose
``close`` arrived while an op held it ``in_use``; the last ``_leave`` then tears it down.
"""
session_id: str
ttl: float
created_at: float
last_access: float = field(default=0.0)
in_use: int = field(default=0)
handle: Optional[str] = field(default=None)
ready: bool = field(default=False)
pending_close: bool = field(default=False)
def is_expired(self, now: float) -> bool:
"""True when the session has been idle longer than its (clamped) TTL."""
return (now - self.last_access) > self.ttl
class SandboxManager:
"""Binds session ids to a shared backend, clamps TTLs, caps concurrency, and reaps idle sessions."""
def __init__(
self,
backend: CodeSandbox,
max_ttl: float,
default_ttl: Optional[float] = None,
max_sessions: Optional[int] = None,
) -> None:
"""Wrap ``backend`` with a registry clamped to ``max_ttl`` and bounded to ``max_sessions``.
The session cap is per-process/worker: each app or Celery process keeps its
own registry, so the effective fleet-wide ceiling is ``max_sessions`` times
the number of live processes.
"""
self._backend = backend
self._max_ttl = max_ttl
self._default_ttl = default_ttl if default_ttl is not None else max_ttl
self._max_sessions = max_sessions if max_sessions and max_sessions > 0 else None
self._sessions: Dict[str, _Session] = {}
self._lock = threading.Lock()
def _clamp_ttl(self, ttl: Optional[float]) -> float:
"""Return an agent-requested TTL bounded to (0, ``max_ttl``]."""
if ttl is None:
ttl = self._default_ttl
if ttl <= 0:
return self._default_ttl
return min(ttl, self._max_ttl)
def open(self, session_id: str, ttl: Optional[float] = None) -> str:
"""Open (or reuse) the sandbox for ``session_id`` with a clamped TTL.
The lock guards ONLY the in-memory registry and refcounts; the (potentially
~60s) cold backend open runs WITHOUT the lock held so it can never serialize
other lock-taking methods. Flow:
1. Under the lock: if the session already exists and is ready, refresh it and
return its cached handle (reuse never does backend I/O). Otherwise reap
idle-expired sessions, evict an LRU-idle victim if at capacity, then RESERVE
a placeholder slot for ``session_id`` so concurrent opens can't overshoot
the cap.
2. Outside the lock: call ``backend.open`` (cold start) and close any reaped /
evicted victims' captured backend resources.
3. Under the lock: finalize the placeholder into a ready session. On failure,
free the reserved slot and re-raise.
"""
with self._lock:
session = self._sessions.get(session_id)
now = time.monotonic()
if session is not None and session.ready:
session.last_access = now
# Reuse must honor an explicit longer keep-alive: a run_code
# persist/ttl on a session first opened by another tool (e.g.
# artifact_generator at the 60s exec timeout) would otherwise be
# dropped, and the kernel + its background state reaped early.
# Extend only (never shrink) so a default reuse can't cut short a
# session another caller kept alive.
if ttl is not None:
session.ttl = max(session.ttl, self._clamp_ttl(ttl))
return session.handle
reaped = self._reap_locked(now)
if session is None:
# Genuinely new key: evict an LRU-idle victim if at capacity (may
# raise SandboxCapacityError), then RESERVE a placeholder slot so
# concurrent opens can't overshoot the cap.
evicted = self._make_room_locked()
self._sessions[session_id] = _Session(
session_id=session_id,
ttl=self._clamp_ttl(ttl),
created_at=now,
last_access=now,
)
else:
# A not-yet-ready placeholder for THIS id already occupies a cap slot
# (another thread is mid cold-open). Overwriting the same key adds no
# slot, so do NOT call _make_room_locked here -- it would wrongly evict
# an innocent LRU-idle session or raise SandboxCapacityError. Refresh
# in place; both threads call the (idempotent) backend.open and the
# finalize step re-checks before binding the handle.
evicted = None
session.last_access = now
session.ttl = self._clamp_ttl(ttl)
# Cold backend open and victim teardown run OUTSIDE the lock.
for sid, handle in reaped:
self._close_backend(sid, handle)
if evicted is not None:
self._close_backend(evicted[0], evicted[1])
try:
handle = self._backend.open(session_id)
except Exception:
# Free the reserved slot so a failed cold open never leaks a cap slot.
with self._lock:
placeholder = self._sessions.get(session_id)
if placeholder is not None and not placeholder.ready:
self._sessions.pop(session_id, None)
self._close_backend(session_id, None) # best-effort: tear down anything created
raise
with self._lock:
session = self._sessions.get(session_id)
if session is None:
# Slot was reclaimed (e.g. closed) while we opened; drop the new runtime.
stale = handle
else:
session.handle = handle
session.ready = True
session.last_access = time.monotonic()
stale = None
if stale is not None:
self._close_backend(session_id, stale)
return handle
def _make_room_locked(self) -> Optional[Tuple[str, Optional[str]]]:
"""Evict the LRU-idle ready session when at capacity; return (id, handle) to close, else None.
Caller holds the lock. The victim is removed from the registry here so its slot
is freed for the reserving caller; its backend ``close`` is deferred to outside
the lock (keyed by the captured handle so a concurrent re-open of the same id
is never torn down). Raises ``SandboxCapacityError`` when the registry is full
and every session is busy or still opening.
"""
if self._max_sessions is None or len(self._sessions) < self._max_sessions:
return None
idle = [s for s in self._sessions.values() if s.ready and s.in_use == 0]
if not idle:
raise SandboxCapacityError(
f"sandbox session cap reached ({self._max_sessions} live, all busy); cannot open another"
)
victim = min(idle, key=lambda s: s.last_access)
self._sessions.pop(victim.session_id, None)
logger.info("SandboxManager evicting LRU-idle session %s to honor cap", victim.session_id)
return (victim.session_id, victim.handle)
def attach(self, session_id: str) -> str:
"""Reattach to an existing session, refreshing its idle clock.
Returns the cached handle without backend I/O for a ready session; falls back
to ``backend.attach`` (outside the lock) only when no handle is cached yet.
"""
with self._lock:
session = self._sessions.get(session_id)
if session is None:
raise KeyError(f"No sandbox session bound for {session_id!r}")
session.last_access = time.monotonic()
cached = session.handle
if cached is not None:
return cached
return self._backend.attach(session_id)
def exec(self, session_id: str, code: str, timeout: Optional[float] = None) -> ExecResult:
"""Execute ``code`` in the bound session, holding it in-use so a reap/evict can't pull it."""
session = self._enter(session_id)
try:
result = self._backend.exec(session_id, code, timeout)
if result.runtime_invalidated:
self._drop_invalidated_session(session_id, session)
return result
finally:
self._leave(session_id, expected=session)
def put_file(self, session_id: str, dest_path: str, data: bytes) -> None:
"""Write ``data`` into the bound session's workspace."""
session = self._enter(session_id)
try:
self._backend.put_file(session_id, dest_path, data)
finally:
self._leave(session_id, expected=session)
def get_file(self, session_id: str, path: str) -> bytes:
"""Read ``path`` from the bound session's workspace."""
session = self._enter(session_id)
try:
return self._backend.get_file(session_id, path)
finally:
self._leave(session_id, expected=session)
def list_files(self, session_id: str) -> List[str]:
"""List files in the bound session's workspace."""
session = self._enter(session_id)
try:
return self._backend.list_files(session_id)
finally:
self._leave(session_id, expected=session)
def remove_path(self, session_id: str, path: str) -> None:
"""Best-effort delete a workspace-relative path; never raises (cleanup must not fail an op)."""
try:
session = self._enter(session_id)
except KeyError:
return
try:
remover = getattr(self._backend, "remove_path", None)
if callable(remover):
remover(session_id, path)
else:
self._remove_via_exec(session_id, path)
except Exception:
logger.exception("SandboxManager: best-effort remove_path failed for %r", path)
finally:
self._leave(session_id, expected=session)
def _remove_via_exec(self, session_id: str, path: str) -> None:
"""Fallback workspace cleanup: run a contained shutil.rmtree of the relative path."""
self._backend.exec(session_id, self._build_remove_program(path), None)
@staticmethod
def _build_remove_program(path: str) -> str:
"""Build the contained shutil.rmtree program for ``path`` with a workspace-root guard."""
# ``path`` is a server-controlled token dir literal; it is passed as a
# repr so its contents are never interpreted as code, and the backend
# already chdirs into the per-session workspace before running. The
# guard rejects absolute, traversal, and empty/'.'/'./' paths so a caller
# can never rmtree the whole workspace root.
return (
"import os, shutil\n"
f"_p = {path!r}\n"
"if _p and not _p.startswith('/') and '..' not in _p.split('/') and os.path.normpath(_p) != '.':\n"
" shutil.rmtree(_p, ignore_errors=True)\n"
)
def close(self, session_id: str) -> None:
"""Tear down the backend runtime and drop the session from the registry.
The backend ``close`` runs OUTSIDE the lock (it may be slow network I/O) and is
keyed by the captured handle so it tears down only the resource this session
owned, never one a concurrent re-open created.
Teardown is DEFERRED when an op holds the session ``in_use`` (a concurrent
exec/put_file/get_file): the session is flagged ``pending_close`` and the last
``_leave`` performs the actual close, so a concurrent close never kills an
in-flight exec (which would lose its captured files). The common path -- the
caller's own exec has already ``_left``, so ``in_use == 0`` -- stays synchronous.
"""
with self._lock:
session = self._sessions.get(session_id)
if session is None:
return
if session.in_use > 0:
session.pending_close = True
return
self._sessions.pop(session_id, None)
handle = session.handle
self._close_backend(session_id, handle)
def has_session(self, session_id: str) -> bool:
"""True when ``session_id`` is currently bound (ready or opening) in the registry."""
with self._lock:
return session_id in self._sessions
def session_count(self) -> int:
"""Return the number of sessions currently bound (ready or reserved) in the registry."""
with self._lock:
return len(self._sessions)
def ttl_for(self, session_id: str) -> Optional[float]:
"""Return the clamped TTL bound to ``session_id``, or None if unbound."""
with self._lock:
session = self._sessions.get(session_id)
return session.ttl if session else None
def _enter(self, session_id: str) -> _Session:
"""Touch the idle clock and mark the session in-use so a concurrent reap/evict skips it."""
with self._lock:
session = self._sessions.get(session_id)
if session is None or not session.ready:
raise KeyError(f"No sandbox session bound for {session_id!r}")
session.last_access = time.monotonic()
session.in_use += 1
return session
def _drop_invalidated_session(self, session_id: str, expected: _Session) -> None:
"""Drop the exact manager entry whose backend runtime was already destroyed."""
with self._lock:
if self._sessions.get(session_id) is expected:
self._sessions.pop(session_id, None)
def _leave(self, session_id: str, expected: Optional[_Session] = None) -> None:
"""Release an in-use hold taken by ``_enter``; run a close deferred by ``close`` on the last release.
Idempotent if the session was already closed. If ``expected`` is supplied,
a newer entry under the same id is never modified (generation/ABA guard).
When the final hold is released and a ``close`` was deferred
(``pending_close``), the session is popped here and its backend torn down
OUTSIDE the lock, keyed by the captured handle.
"""
handle: Optional[str] = None
do_close = False
with self._lock:
session = self._sessions.get(session_id)
if session is not None and (expected is None or session is expected) and session.in_use > 0:
session.in_use -= 1
session.last_access = time.monotonic()
if session.in_use == 0 and session.pending_close:
self._sessions.pop(session_id, None)
handle = session.handle
do_close = True
if do_close:
self._close_backend(session_id, handle)
def _close_backend(self, session_id: str, handle: Optional[str]) -> None:
"""Close the SPECIFIC backend resource captured for this session, best-effort.
When the backend exposes ``close_handle``, the captured handle id is passed so a
concurrent re-open of the same ``session_id`` (which created a new backend handle)
is never the one torn down. Backends without that hook fall back to ``close`` by
id, which is correct for callers (``close``/reap/failed-open) where no concurrent
re-open of a still-registered session can be in flight.
"""
try:
closer = getattr(self._backend, "close_handle", None)
if handle is not None and callable(closer):
closer(session_id, handle)
else:
self._backend.close(session_id)
except Exception:
logger.exception("SandboxManager: backend close failed for %s", session_id)
def reap_expired(self) -> List[str]:
"""Close sessions idle past their TTL and return the reaped session ids.
Artifacts are persisted eagerly by the tools/code node right after each
exec, so a session's workspace is scratch: reaping only closes the kernel
and never loses a user-facing artifact. Busy (in-use) sessions are left
alone so a reap can't pull a workspace out from under a running exec.
"""
now = time.monotonic()
with self._lock:
expired = self._reap_locked(now)
for sid, handle in expired:
self._close_backend(sid, handle)
return [sid for sid, _ in expired]
def _reap_locked(self, now: float) -> List[Tuple[str, Optional[str]]]:
"""Pop idle-expired ready sessions from the registry; return (id, handle) pairs to close.
Caller holds the lock. Placeholders (not yet ready) and busy sessions are left
alone. The backend close of each popped session is deferred to outside the lock.
"""
expired = [
(sid, s.handle)
for sid, s in self._sessions.items()
if s.ready and s.in_use == 0 and s.is_expired(now)
]
for sid, _ in expired:
self._sessions.pop(sid, None)
return expired
+91
View File
@@ -0,0 +1,91 @@
"""Factory + process-wide singleton selecting a sandbox backend from settings."""
from typing import Callable, Dict, Optional
from application.core.settings import settings
from application.sandbox.base import CodeSandbox
from application.sandbox.jupyter_gateway import JupyterKernelGatewaySandbox
from application.sandbox.manager import SandboxManager
def _make_jupyter() -> CodeSandbox:
"""Build the Jupyter Kernel Gateway backend from ``SANDBOX_*`` settings."""
return JupyterKernelGatewaySandbox(
gateway_url=settings.SANDBOX_GATEWAY_URL,
auth_token=settings.SANDBOX_GATEWAY_AUTH_TOKEN,
kernel_name=settings.SANDBOX_KERNEL_NAME,
default_timeout=float(settings.SANDBOX_EXEC_TIMEOUT),
http_timeout=float(settings.SANDBOX_HTTP_TIMEOUT),
max_output_bytes=int(settings.SANDBOX_MAX_OUTPUT_BYTES),
max_file_bytes=int(settings.SANDBOX_MAX_FILE_BYTES),
)
def _make_daytona() -> CodeSandbox:
"""Build the Daytona Cloud backend from ``DAYTONA_*``/``SANDBOX_*`` settings."""
from application.sandbox.daytona import DaytonaSandbox
# Auto-delete is the only backstop against orphaned (paid) sandboxes, so a
# never-expiring value (<= 0) is rejected and clamped to a safe default.
auto_delete_interval = int(settings.DAYTONA_AUTO_DELETE_INTERVAL)
if auto_delete_interval <= 0:
auto_delete_interval = 60
return DaytonaSandbox(
api_key=settings.DAYTONA_API_KEY,
api_url=settings.DAYTONA_API_URL,
target=settings.DAYTONA_TARGET,
snapshot=settings.DAYTONA_SNAPSHOT,
language=settings.DAYTONA_LANGUAGE,
default_timeout=float(settings.SANDBOX_EXEC_TIMEOUT),
create_timeout=float(settings.SANDBOX_HTTP_TIMEOUT) * 6,
auto_stop_interval=int(settings.DAYTONA_AUTO_STOP_INTERVAL),
auto_delete_interval=auto_delete_interval,
# Bounds only what the app propagates from a response: the Daytona SDK still
# buffers the whole HTTP response before we see it, so a full in-sandbox cap
# would need runner-side support. This is a partial but real mitigation.
max_output_bytes=int(settings.SANDBOX_MAX_OUTPUT_BYTES),
max_file_bytes=int(settings.SANDBOX_MAX_FILE_BYTES),
max_sandboxes=int(settings.DAYTONA_MAX_SANDBOXES),
)
class SandboxCreator:
"""Resolves ``SANDBOX_BACKEND`` to a backend and caches a single manager."""
backends: Dict[str, Callable[[], CodeSandbox]] = {
"jupyter": _make_jupyter,
"daytona": _make_daytona,
}
_instance = None
@classmethod
def get_manager(cls) -> SandboxManager:
"""Return the process-wide ``SandboxManager``, building it on first use."""
if cls._instance is None:
backend = cls.create_backend(getattr(settings, "SANDBOX_BACKEND", "jupyter"))
cls._instance = SandboxManager(
backend=backend,
max_ttl=float(settings.SANDBOX_MAX_TTL),
max_sessions=int(settings.SANDBOX_MAX_SESSIONS),
)
return cls._instance
@classmethod
def peek_manager(cls) -> Optional[SandboxManager]:
"""Return the process-wide manager if already built, else None (never constructs one)."""
return cls._instance
@classmethod
def create_backend(cls, type_name: str) -> CodeSandbox:
"""Instantiate the backend registered under ``type_name`` (case-insensitive)."""
factory = cls.backends.get(type_name.lower())
if not factory:
raise ValueError(f"No sandbox backend found for type {type_name}")
return factory()
@classmethod
def reset(cls) -> None:
"""Drop the cached singleton (test/teardown hook)."""
cls._instance = None
+17
View File
@@ -37,6 +37,23 @@ class BaseStorage(ABC):
"""
pass
def generate_presigned_url(self, path: str, expires_in: int = 300) -> str:
"""Return a short-lived presigned download URL; not all backends support it.
Args:
path: Path to the file
expires_in: TTL of the signed URL in seconds
Returns:
str: A presigned URL granting time-limited read access
Raises:
NotImplementedError: If the backend cannot mint presigned URLs.
"""
raise NotImplementedError(
f"{type(self).__name__} does not support presigned URLs"
)
@abstractmethod
def process_file(self, path: str, processor_func: Callable, **kwargs):
"""
+68
View File
@@ -401,6 +401,74 @@ attachments_table = Table(
Column("legacy_mongo_id", Text),
)
# Identity row, one per logical artifact. The stable ``id`` is the handle passed
# around (chat/workflow state/message bodies carry only this reference, never
# bytes). Authz is parent-derived: whoever can reach ``conversation_id`` (chat)
# or ``workflow_run_id`` (run) can reach its artifacts, so a CHECK requires at
# least one parent. ``user_id`` is ownership/quota only; ``team_id`` is a
# nullable forward-compat hook for Teams (no FK, matching ``teams`` itself).
artifacts_table = Table(
"artifacts",
metadata,
Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()),
Column("user_id", Text, nullable=False),
Column("conversation_id", UUID(as_uuid=True)),
Column("workflow_run_id", UUID(as_uuid=True)),
Column("team_id", UUID(as_uuid=True)),
Column("message_id", UUID(as_uuid=True)),
Column("kind", Text, nullable=False),
Column("title", Text),
Column("metadata", JSONB),
Column("current_version", Integer, nullable=False, server_default="1"),
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
CheckConstraint(
"conversation_id IS NOT NULL OR workflow_run_id IS NOT NULL",
name="artifacts_parent_present_check",
),
)
Index(
"artifacts_conversation_idx",
artifacts_table.c.conversation_id,
postgresql_where=artifacts_table.c.conversation_id.isnot(None),
)
Index(
"artifacts_workflow_run_idx",
artifacts_table.c.workflow_run_id,
postgresql_where=artifacts_table.c.workflow_run_id.isnot(None),
)
Index("artifacts_user_idx", artifacts_table.c.user_id)
# Append-only version history; never mutated. Each edit appends a row and bumps
# ``artifacts.current_version``. ``UNIQUE(artifact_id, version)`` keeps versions
# monotonic. ``storage_path`` is the ``BaseStorage`` key (NULL when spec-only);
# bytes live in storage, only metadata + the key live here (pass-by-reference).
artifact_versions_table = Table(
"artifact_versions",
metadata,
Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()),
Column(
"artifact_id",
UUID(as_uuid=True),
ForeignKey("artifacts.id", ondelete="CASCADE"),
nullable=False,
),
Column("version", Integer, nullable=False),
Column("mime_type", Text),
Column("filename", Text),
Column("storage_path", Text),
Column("size", BigInteger),
Column("sha256", Text),
Column("spec", JSONB),
Column("preview_text", Text),
Column("produced_by", JSONB),
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
UniqueConstraint("artifact_id", "version", name="artifact_versions_artifact_version_uidx"),
)
Index("artifact_versions_artifact_idx", artifact_versions_table.c.artifact_id)
memories_table = Table(
"memories",
metadata,
@@ -0,0 +1,627 @@
"""Repository for the ``artifacts`` / ``artifact_versions`` tables.
Append-only artifact store. An ``artifacts`` row is the stable identity; each
edit appends an immutable ``artifact_versions`` row and atomically bumps
``artifacts.current_version`` (mirroring ``workflows.increment_graph_version``).
Authz is parent-derived: reads are not gated on ``user_id`` — callers resolve
access via the artifact's ``conversation_id`` / ``workflow_run_id`` parent.
``user_id`` is carried for ownership / quota accounting only. Pass-by-reference:
only metadata + the ``BaseStorage`` key are stored here, never binary bytes.
"""
from __future__ import annotations
import json
import logging
from typing import Any, Optional
from sqlalchemy import Connection, text
from application.storage.db.base_repository import row_to_dict
logger = logging.getLogger(__name__)
def _artifact_to_dict(row: Any) -> dict:
"""row_to_dict for an ``artifacts`` row (``metadata`` key preserved)."""
return row_to_dict(row)
def _version_to_dict(row: Any) -> dict:
"""row_to_dict for an ``artifact_versions`` row (``spec``/``produced_by`` keys preserved)."""
return row_to_dict(row)
class ArtifactsRepository:
"""CRUD for artifact identities and their append-only versions."""
def __init__(self, conn: Connection) -> None:
"""Bind the repository to an open SQLAlchemy connection."""
self._conn = conn
def create_artifact(
self,
user_id: str,
kind: str,
*,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
team_id: Optional[str] = None,
message_id: Optional[str] = None,
title: Optional[str] = None,
metadata: Any = None,
mime_type: Optional[str] = None,
filename: Optional[str] = None,
storage_path: Optional[str] = None,
size: Optional[int] = None,
sha256: Optional[str] = None,
spec: Any = None,
preview_text: Optional[str] = None,
produced_by: Any = None,
) -> dict:
"""Create the identity row and its version 1 atomically; return the artifact dict.
When a parent is given, a STABLE per-parent ``ref_seq`` (max existing + 1) is
computed BEFORE the insert and stored into ``metadata`` so the short ``A{n}`` ref
the model receives survives later deletions. See ``next_ref_seq`` for the (benign)
concurrent-create dup caveat.
"""
metadata_payload = metadata
if conversation_id is not None or workflow_run_id is not None:
ref_seq = self.next_ref_seq(
conversation_id=conversation_id, workflow_run_id=workflow_run_id
)
merged = dict(metadata) if isinstance(metadata, dict) else {}
merged["ref_seq"] = ref_seq
metadata_payload = merged
artifact = self._conn.execute(
text(
"""
INSERT INTO artifacts (
user_id, conversation_id, workflow_run_id, team_id,
message_id, kind, title, metadata, current_version
)
VALUES (
:user_id,
CAST(:conversation_id AS uuid),
CAST(:workflow_run_id AS uuid),
CAST(:team_id AS uuid),
CAST(:message_id AS uuid),
:kind, :title, CAST(:metadata AS jsonb), 1
)
RETURNING *
"""
),
{
"user_id": user_id,
"conversation_id": conversation_id,
"workflow_run_id": workflow_run_id,
"team_id": team_id,
"message_id": message_id,
"kind": kind,
"title": title,
"metadata": json.dumps(metadata_payload) if metadata_payload is not None else None,
},
).fetchone()
artifact_dict = _artifact_to_dict(artifact)
self._insert_version(
artifact_id=artifact_dict["id"],
version=1,
mime_type=mime_type,
filename=filename,
storage_path=storage_path,
size=size,
sha256=sha256,
spec=spec,
preview_text=preview_text,
produced_by=produced_by,
)
return artifact_dict
def get_artifact(self, artifact_id: str) -> Optional[dict]:
"""Fetch one artifact by id, unscoped (prefer ``get_artifact_in_parent`` at request boundaries)."""
result = self._conn.execute(
text("SELECT * FROM artifacts WHERE id = CAST(:id AS uuid)"),
{"id": artifact_id},
)
row = result.fetchone()
return _artifact_to_dict(row) if row is not None else None
def get_artifact_in_parent(
self,
artifact_id: str,
*,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
) -> Optional[dict]:
"""Fetch an artifact only if it belongs to the given parent; the safe gate for request handlers."""
if conversation_id is None and workflow_run_id is None:
raise ValueError("get_artifact_in_parent requires conversation_id or workflow_run_id")
clauses = ["id = CAST(:id AS uuid)"]
params: dict[str, Any] = {"id": artifact_id}
if conversation_id is not None:
clauses.append("conversation_id = CAST(:conversation_id AS uuid)")
params["conversation_id"] = conversation_id
if workflow_run_id is not None:
clauses.append("workflow_run_id = CAST(:workflow_run_id AS uuid)")
params["workflow_run_id"] = workflow_run_id
result = self._conn.execute(
text(f"SELECT * FROM artifacts WHERE {' AND '.join(clauses)}"),
params,
)
row = result.fetchone()
return _artifact_to_dict(row) if row is not None else None
def find_bridged_attachment(
self,
attachment_id: str,
*,
conversation_id: str,
) -> Optional[dict]:
"""Return the conversation artifact already bridged from ``attachment_id``, or None (idempotency gate)."""
result = self._conn.execute(
text(
"SELECT a.* FROM artifacts a "
"JOIN artifact_versions v "
" ON v.artifact_id = a.id AND v.version = a.current_version "
"WHERE a.conversation_id = CAST(:conversation_id AS uuid) "
" AND v.produced_by ->> 'attachment_id' = :attachment_id "
"ORDER BY a.created_at ASC, a.id ASC LIMIT 1"
),
{"conversation_id": conversation_id, "attachment_id": str(attachment_id)},
)
row = result.fetchone()
return _artifact_to_dict(row) if row is not None else None
def list_artifacts(
self,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
user_id: Optional[str] = None,
) -> list[dict]:
"""List artifacts filtered by parent / owner (newest first); at least one filter is required."""
if conversation_id is None and workflow_run_id is None and user_id is None:
raise ValueError("list_artifacts requires at least one of conversation_id, workflow_run_id, user_id")
clauses: list[str] = []
params: dict[str, Any] = {}
if conversation_id is not None:
clauses.append("conversation_id = CAST(:conversation_id AS uuid)")
params["conversation_id"] = conversation_id
if workflow_run_id is not None:
clauses.append("workflow_run_id = CAST(:workflow_run_id AS uuid)")
params["workflow_run_id"] = workflow_run_id
if user_id is not None:
clauses.append("user_id = :user_id")
params["user_id"] = user_id
where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
result = self._conn.execute(
text(f"SELECT * FROM artifacts {where} ORDER BY created_at DESC"),
params,
)
return [_artifact_to_dict(r) for r in result.fetchall()]
def list_artifacts_for_agent(
self,
agent_id: str,
user_id: str,
conversation_id: Optional[str] = None,
) -> list[dict]:
"""List an agent's artifacts (owner-scoped), optionally narrowed to one conversation.
Scopes a per-agent api-key's artifact visibility to the conversations that
agent produced, so the key cannot enumerate the owner's whole corpus. When
``conversation_id`` is given, the SQL narrows to that single conversation
(the per-visitor bearer capability for a public widget key).
"""
clauses = ["c.agent_id = CAST(:agent_id AS uuid)", "a.user_id = :user_id"]
params: dict[str, Any] = {"agent_id": str(agent_id), "user_id": user_id}
if conversation_id is not None:
clauses.append("a.conversation_id = CAST(:conversation_id AS uuid)")
params["conversation_id"] = conversation_id
result = self._conn.execute(
text(
"SELECT a.* FROM artifacts a "
"JOIN conversations c ON a.conversation_id = c.id "
f"WHERE {' AND '.join(clauses)} "
"ORDER BY a.created_at DESC, a.id DESC"
),
params,
)
return [_artifact_to_dict(r) for r in result.fetchall()]
def artifact_in_agent_scope(self, artifact_id: str, agent_id: str) -> bool:
"""True if ``artifact_id``'s parent conversation belongs to ``agent_id``."""
clauses = [
"a.id = CAST(:id AS uuid)",
"c.agent_id = CAST(:agent_id AS uuid)",
]
result = self._conn.execute(
text(
"SELECT 1 FROM artifacts a "
"JOIN conversations c ON a.conversation_id = c.id "
f"WHERE {' AND '.join(clauses)} "
"LIMIT 1"
),
{"id": artifact_id, "agent_id": str(agent_id)},
)
return result.fetchone() is not None
def position_in_parent(
self,
artifact_id: str,
*,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
) -> int:
"""Return the 1-based position of an artifact within its parent (by created_at, id tie-break); 0 if absent."""
if conversation_id is None and workflow_run_id is None:
raise ValueError("position_in_parent requires conversation_id or workflow_run_id")
outer, params = self._parent_clauses(conversation_id, workflow_run_id, alias="a")
inner, _ = self._parent_clauses(conversation_id, workflow_run_id, alias="t")
params["id"] = artifact_id
# The inner SELECT applies the same parent scope, so an artifact in a
# different parent yields no anchor row -> count() is 0 (not in this parent).
row = self._conn.execute(
text(
f"SELECT count(*) FROM artifacts a "
f"WHERE {' AND '.join(outer)} AND (a.created_at, a.id) <= ("
f" SELECT t.created_at, t.id FROM artifacts t "
f" WHERE {' AND '.join(inner)} AND t.id = CAST(:id AS uuid)"
f")"
),
params,
).fetchone()
return int(row[0]) if row is not None else 0
def artifact_id_at_position(
self,
n: int,
*,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
) -> Optional[str]:
"""Return the id of the n-th artifact (1-based, created_at asc, id tie-break) in a parent, or None."""
if conversation_id is None and workflow_run_id is None:
raise ValueError("artifact_id_at_position requires conversation_id or workflow_run_id")
if not isinstance(n, int) or n < 1:
return None
clauses, params = self._parent_clauses(conversation_id, workflow_run_id)
params["offset"] = n - 1
row = self._conn.execute(
text(
f"SELECT id FROM artifacts WHERE {' AND '.join(clauses)} "
f"ORDER BY created_at ASC, id ASC OFFSET :offset LIMIT 1"
),
params,
).fetchone()
return str(row[0]) if row is not None else None
def next_ref_seq(
self,
*,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
) -> int:
"""Return the next stable ``ref_seq`` for a parent (max stored + 1; 1 when empty).
Caveat: two concurrent creates in one parent can read the same MAX and get a
duplicate ``ref_seq`` (creates within a parent are normally serial from one agent
turn); a dup resolves to the earliest — acceptable, and strictly better than the
re-point bug. A DB unique index is a possible future hardening (no migration here).
"""
if conversation_id is None and workflow_run_id is None:
raise ValueError("next_ref_seq requires conversation_id or workflow_run_id")
clauses, params = self._parent_clauses(conversation_id, workflow_run_id)
# Guard the cast: legacy rows have no ``ref_seq`` (NULL, ignored), and the regex
# match keeps any non-numeric value from erroring the cast.
row = self._conn.execute(
text(
"SELECT COALESCE(MAX(CASE WHEN metadata ->> 'ref_seq' ~ '^[0-9]+$' "
"THEN (metadata ->> 'ref_seq')::int END), 0) + 1 "
f"FROM artifacts WHERE {' AND '.join(clauses)}"
),
params,
).fetchone()
return int(row[0]) if row is not None else 1
def resolve_id_by_ref_seq(
self,
seq: int,
*,
conversation_id: Optional[str] = None,
workflow_run_id: Optional[str] = None,
) -> Optional[str]:
"""Return the id of the artifact whose stored ``ref_seq`` == ``seq`` in a parent, or None.
Earliest ``created_at`` wins on the rare concurrent-create tie (see ``next_ref_seq``).
"""
if conversation_id is None and workflow_run_id is None:
raise ValueError("resolve_id_by_ref_seq requires conversation_id or workflow_run_id")
if not isinstance(seq, int) or seq < 1:
return None
clauses, params = self._parent_clauses(conversation_id, workflow_run_id)
params["seq"] = str(seq)
row = self._conn.execute(
text(
f"SELECT id FROM artifacts WHERE {' AND '.join(clauses)} "
"AND metadata ->> 'ref_seq' = :seq "
"ORDER BY created_at ASC, id ASC LIMIT 1"
),
params,
).fetchone()
return str(row[0]) if row is not None else None
@staticmethod
def _parent_clauses(
conversation_id: Optional[str], workflow_run_id: Optional[str], alias: str = ""
) -> tuple[list[str], dict[str, Any]]:
"""Build the parent-scope WHERE clauses + params (optionally column-aliased) for the position helpers."""
prefix = f"{alias}." if alias else ""
clauses: list[str] = []
params: dict[str, Any] = {}
if conversation_id is not None:
clauses.append(f"{prefix}conversation_id = CAST(:conversation_id AS uuid)")
params["conversation_id"] = conversation_id
if workflow_run_id is not None:
clauses.append(f"{prefix}workflow_run_id = CAST(:workflow_run_id AS uuid)")
params["workflow_run_id"] = workflow_run_id
return clauses, params
def count_for_user(self, user_id: str) -> int:
"""Return how many artifacts ``user_id`` currently owns (quota accounting)."""
row = self._conn.execute(
text("SELECT count(*) FROM artifacts WHERE user_id = :user_id"),
{"user_id": user_id},
).fetchone()
return int(row[0]) if row is not None else 0
def total_bytes_for_user(self, user_id: str) -> int:
"""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(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()
return int(row[0]) if row is not None else 0
def delete_artifact(self, artifact_id: str) -> list[str]:
"""Delete an artifact (+ its versions, via FK cascade); return its stored paths.
The caller deletes the returned storage objects best-effort AFTER commit;
the DB delete also frees the owner's count/byte quota.
"""
paths = [
v["storage_path"]
for v in self.list_versions(artifact_id)
if v.get("storage_path")
]
self._conn.execute(
text("DELETE FROM artifacts WHERE id = CAST(:id AS uuid)"),
{"id": artifact_id},
)
return paths
def storage_paths_for_conversation(self, conversation_id: str) -> list[str]:
"""Return every stored version path for a conversation's artifacts (for byte cleanup)."""
result = self._conn.execute(
text(
"SELECT v.storage_path FROM artifact_versions v "
"JOIN artifacts a ON a.id = v.artifact_id "
"WHERE a.conversation_id = CAST(:cid AS uuid) "
"AND v.storage_path IS NOT NULL"
),
{"cid": conversation_id},
)
return [r[0] for r in result.fetchall()]
def storage_paths_for_user_conversations(self, user_id: str) -> list[str]:
"""Return every stored version path for artifacts under a user's conversations."""
result = self._conn.execute(
text(
"SELECT v.storage_path FROM artifact_versions v "
"JOIN artifacts a ON a.id = v.artifact_id "
"JOIN conversations c ON c.id = a.conversation_id "
"WHERE c.user_id = :user_id AND v.storage_path IS NOT NULL"
),
{"user_id": user_id},
)
return [r[0] for r in result.fetchall()]
def delete_for_conversation(self, conversation_id: str) -> int:
"""Delete all artifacts (+ versions, via cascade) parented to a conversation.
``conversation_id`` is a bare uuid column (no FK), so the parent delete
does not cascade — callers invoke this to reclaim the rows + their quota.
"""
result = self._conn.execute(
text("DELETE FROM artifacts WHERE conversation_id = CAST(:cid AS uuid)"),
{"cid": conversation_id},
)
return result.rowcount
def delete_for_user_conversations(self, user_id: str) -> int:
"""Delete every artifact parented to one of a user's conversations (+ versions)."""
result = self._conn.execute(
text(
"DELETE FROM artifacts WHERE conversation_id IN "
"(SELECT id FROM conversations WHERE user_id = :user_id)"
),
{"user_id": user_id},
)
return result.rowcount
def storage_paths_for_workflow(self, workflow_id: str) -> list[str]:
"""Return every stored version path for artifacts under a workflow's runs (for byte cleanup)."""
result = self._conn.execute(
text(
"SELECT v.storage_path FROM artifact_versions v "
"JOIN artifacts a ON a.id = v.artifact_id "
"JOIN workflow_runs r ON r.id = a.workflow_run_id "
"WHERE r.workflow_id = CAST(:wid AS uuid) "
"AND v.storage_path IS NOT NULL"
),
{"wid": workflow_id},
)
return [r[0] for r in result.fetchall()]
def delete_for_workflow(self, workflow_id: str) -> int:
"""Delete all artifacts (+ versions, via cascade) parented to a workflow's runs.
``artifacts.workflow_run_id`` is a bare uuid column (no FK), so deleting the
workflow (which cascade-deletes its ``workflow_runs``) does NOT reap these
rows -- callers invoke this to reclaim the rows + their quota. Must run
BEFORE the workflow delete, while the run rows still resolve the subquery.
"""
result = self._conn.execute(
text(
"DELETE FROM artifacts WHERE workflow_run_id IN "
"(SELECT id FROM workflow_runs WHERE workflow_id = CAST(:wid AS uuid))"
),
{"wid": workflow_id},
)
return result.rowcount
@staticmethod
def reap_storage(paths: list) -> None:
"""Best-effort delete a batch of artifact byte objects; never raises."""
if not paths:
return
try:
from application.storage.storage_creator import StorageCreator
storage = StorageCreator.get_storage()
except Exception:
logger.warning("artifact cleanup: storage unavailable", exc_info=True)
return
for path in paths:
try:
storage.delete_file(path)
except Exception:
logger.warning("artifact cleanup: failed to delete bytes %s", path, exc_info=True)
def get_version(self, artifact_id: str, version: int) -> Optional[dict]:
"""Fetch a single version row by artifact id and version number."""
result = self._conn.execute(
text(
"SELECT * FROM artifact_versions "
"WHERE artifact_id = CAST(:artifact_id AS uuid) AND version = :version"
),
{"artifact_id": artifact_id, "version": version},
)
row = result.fetchone()
return _version_to_dict(row) if row is not None else None
def list_versions(self, artifact_id: str) -> list[dict]:
"""List every version of an artifact, oldest first."""
result = self._conn.execute(
text(
"SELECT * FROM artifact_versions "
"WHERE artifact_id = CAST(:artifact_id AS uuid) ORDER BY version ASC"
),
{"artifact_id": artifact_id},
)
return [_version_to_dict(r) for r in result.fetchall()]
def append_version(
self,
artifact_id: str,
*,
mime_type: Optional[str] = None,
filename: Optional[str] = None,
storage_path: Optional[str] = None,
size: Optional[int] = None,
sha256: Optional[str] = None,
spec: Any = None,
preview_text: Optional[str] = None,
produced_by: Any = None,
) -> dict:
"""Append a new version and atomically bump ``current_version``; return the version dict."""
new_version = self._conn.execute(
text(
"UPDATE artifacts "
"SET current_version = current_version + 1, updated_at = now() "
"WHERE id = CAST(:id AS uuid) "
"RETURNING current_version"
),
{"id": artifact_id},
).fetchone()
if new_version is None:
raise ValueError(f"artifact {artifact_id} not found")
return self._insert_version(
artifact_id=artifact_id,
version=new_version[0],
mime_type=mime_type,
filename=filename,
storage_path=storage_path,
size=size,
sha256=sha256,
spec=spec,
preview_text=preview_text,
produced_by=produced_by,
)
def _insert_version(
self,
*,
artifact_id: str,
version: int,
mime_type: Optional[str],
filename: Optional[str],
storage_path: Optional[str],
size: Optional[int],
sha256: Optional[str],
spec: Any,
preview_text: Optional[str],
produced_by: Any,
) -> dict:
"""Insert one append-only version row; UNIQUE(artifact_id, version) guards duplicates."""
result = self._conn.execute(
text(
"""
INSERT INTO artifact_versions (
artifact_id, version, mime_type, filename, storage_path,
size, sha256, spec, preview_text, produced_by
)
VALUES (
CAST(:artifact_id AS uuid), :version, :mime_type, :filename,
:storage_path, :size, :sha256, CAST(:spec AS jsonb),
:preview_text, CAST(:produced_by AS jsonb)
)
RETURNING *
"""
),
{
"artifact_id": artifact_id,
"version": version,
"mime_type": mime_type,
"filename": filename,
"storage_path": storage_path,
"size": size,
"sha256": sha256,
"spec": json.dumps(spec) if spec is not None else None,
"preview_text": preview_text,
"produced_by": json.dumps(produced_by) if produced_by is not None else None,
},
)
return _version_to_dict(result.fetchone())
@@ -15,6 +15,7 @@ Covers every operation the legacy Mongo code performs on
from __future__ import annotations
import json
import logging
from enum import Enum
from typing import Optional
@@ -25,6 +26,8 @@ from application.storage.db.base_repository import looks_like_uuid, row_to_dict
from application.storage.db.models import conversations_table, conversation_messages_table
from application.storage.db.serialization import PGNativeJSONEncoder
logger = logging.getLogger(__name__)
class MessageUpdateOutcome(str, Enum):
"""Discriminated result of ``update_message_by_id``.
@@ -535,6 +538,14 @@ class ConversationsRepository:
# a non-UUID id reaches this code path.
if not looks_like_uuid(conversation_id):
return False
# Artifacts are parent-derived: deleting a conversation must take its
# artifacts with it (rows + bytes + the quota they consume). conversation_id
# is a bare uuid (no FK cascade), so delete the rows explicitly and reap the
# stored bytes best-effort.
from application.storage.db.repositories.artifacts import ArtifactsRepository
artifacts = ArtifactsRepository(self._conn)
paths = artifacts.storage_paths_for_conversation(conversation_id)
result = self._conn.execute(
text(
"DELETE FROM conversations "
@@ -542,15 +553,45 @@ class ConversationsRepository:
),
{"id": conversation_id, "user_id": user_id},
)
return result.rowcount > 0
deleted = result.rowcount > 0
if deleted:
artifacts.delete_for_conversation(conversation_id)
self._reap_artifact_storage(paths)
return deleted
def delete_all_for_user(self, user_id: str) -> int:
from application.storage.db.repositories.artifacts import ArtifactsRepository
# Reap artifacts BEFORE the conversations: delete_for_user_conversations
# resolves them via a subquery over the still-present conversation rows.
artifacts = ArtifactsRepository(self._conn)
paths = artifacts.storage_paths_for_user_conversations(user_id)
artifacts.delete_for_user_conversations(user_id)
result = self._conn.execute(
text("DELETE FROM conversations WHERE user_id = :user_id"),
{"user_id": user_id},
)
self._reap_artifact_storage(paths)
return result.rowcount
@staticmethod
def _reap_artifact_storage(paths: list) -> None:
"""Best-effort delete artifact bytes whose rows were cascade-removed; never raises."""
if not paths:
return
try:
from application.storage.storage_creator import StorageCreator
storage = StorageCreator.get_storage()
except Exception:
logger.warning("conversation delete: storage unavailable for artifact cleanup", exc_info=True)
return
for path in paths:
try:
storage.delete_file(path)
except Exception:
logger.warning("conversation delete: failed to delete artifact bytes %s", path, exc_info=True)
# ------------------------------------------------------------------
# Messages
# ------------------------------------------------------------------
@@ -566,6 +607,44 @@ class ConversationsRepository:
)
return [_message_row_to_dict(r) for r in result.fetchall()]
def first_n_message_ids(self, conversation_id: str, first_n: int) -> set[str]:
"""Return the ids of the first ``first_n`` messages (by position) of the conversation.
Same ordering a share snapshot uses (``get_messages``' ``position ASC``
then ``[:first_n]``); empty set for ``first_n <= 0`` or a non-UUID id.
"""
if first_n <= 0 or not looks_like_uuid(conversation_id):
return set()
result = self._conn.execute(
text(
"SELECT id FROM conversation_messages "
"WHERE conversation_id = CAST(:cid AS uuid) "
"ORDER BY position ASC LIMIT :n"
),
{"cid": conversation_id, "n": int(first_n)},
)
return {str(r[0]) for r in result.fetchall()}
def message_in_first_n(self, conversation_id: str, message_id: str, first_n: int) -> bool:
"""True if message_id is among the first ``first_n`` messages (by position) of the conversation."""
if not message_id or first_n <= 0:
return False
# Shape-gate both ids: a non-UUID reaching CAST(... AS uuid) would raise
# and poison the enclosing transaction (see ``rename``).
if not looks_like_uuid(conversation_id) or not looks_like_uuid(message_id):
return False
row = self._conn.execute(
text(
"SELECT 1 FROM ( "
" SELECT id FROM conversation_messages "
" WHERE conversation_id = CAST(:cid AS uuid) "
" ORDER BY position ASC LIMIT :n "
") t WHERE t.id = CAST(:mid AS uuid)"
),
{"cid": conversation_id, "mid": message_id, "n": int(first_n)},
).fetchone()
return row is not None
def get_message_at(self, conversation_id: str, position: int) -> Optional[dict]:
# Shape-gate: see ``rename``. Callers today always pass a resolved
# UUID (via ``get_any`` first), but the guard costs nothing and
@@ -6,9 +6,10 @@ written once after workflow execution completes and never updated.
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
@@ -25,6 +26,7 @@ class WorkflowRunsRepository:
user_id: str,
status: str,
*,
run_id: str | None = None,
inputs: dict | None = None,
result: dict | None = None,
steps: list | None = None,
@@ -37,6 +39,10 @@ class WorkflowRunsRepository:
"user_id": user_id,
"status": status,
}
# An explicit id lets the engine bind run-scoped artifacts to this row
# before the run is persisted, so artifact authz can resolve the parent.
if run_id is not None:
values["id"] = run_id
if inputs is not None:
values["inputs"] = inputs
if result is not None:
@@ -54,6 +60,60 @@ class WorkflowRunsRepository:
res = self._conn.execute(stmt)
return row_to_dict(res.fetchone())
def finalize(
self,
run_id: str,
user_id: str,
status: str,
*,
result: dict | None = None,
steps: list | None = None,
ended_at: Optional[datetime] = None,
) -> bool:
"""Update a pre-created run row with its terminal status/result; owner-scoped."""
values: dict = {"status": status}
if result is not None:
values["result"] = result
if steps is not None:
values["steps"] = steps
if ended_at is not None:
values["ended_at"] = ended_at
stmt = (
workflow_runs_table.update()
.where(workflow_runs_table.c.id == run_id)
.where(workflow_runs_table.c.user_id == user_id)
.values(**values)
)
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(),
# failed_reason marks the reap so readers know ended_at is the
# reap time, not when the run actually died (duration is bogus).
result={
"error": "Run did not complete (timed out or the client disconnected).",
"failed_reason": "stale_reaper",
},
)
)
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)"),
@@ -13,7 +13,7 @@ from typing import Optional
from sqlalchemy import Connection, text
from sqlalchemy.dialects.postgresql import insert as pg_insert
from application.storage.db.base_repository import row_to_dict
from application.storage.db.base_repository import looks_like_uuid, row_to_dict
from application.storage.db.models import workflows_table
@@ -116,6 +116,31 @@ class WorkflowsRepository:
return row[0] if row else None
def delete(self, workflow_id: str, user_id: str) -> bool:
# Run artifacts parent to ``workflow_runs`` (which cascade-delete with the
# workflow), but ``artifacts.workflow_run_id`` is a bare uuid with no FK, so
# the rows + bytes + the quota they consume would leak. Reclaim them BEFORE
# deleting the workflow, while the run rows still resolve the subquery.
# Ownership is confirmed first so this never reaps another user's artifacts.
from application.storage.db.repositories.artifacts import ArtifactsRepository
if not looks_like_uuid(workflow_id):
return False
owned = (
self._conn.execute(
text(
"SELECT 1 FROM workflows "
"WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
),
{"id": workflow_id, "user_id": user_id},
).fetchone()
is not None
)
if not owned:
return False
artifacts = ArtifactsRepository(self._conn)
paths = artifacts.storage_paths_for_workflow(workflow_id)
artifacts.delete_for_workflow(workflow_id)
result = self._conn.execute(
text(
"DELETE FROM workflows "
@@ -123,7 +148,10 @@ class WorkflowsRepository:
),
{"id": workflow_id, "user_id": user_id},
)
return result.rowcount > 0
deleted = result.rowcount > 0
if deleted:
artifacts.reap_storage(paths)
return deleted
def delete_by_legacy_id(self, legacy_mongo_id: str, user_id: str) -> bool:
"""Delete a workflow addressed by the Mongo ObjectId string.
+10
View File
@@ -121,6 +121,16 @@ class S3Storage(BaseStorage):
file_obj.seek(0)
return file_obj
def generate_presigned_url(self, path: str, expires_in: int = 300) -> str:
"""Return a short-lived presigned GET URL for a private object (TTL <= 1h)."""
path = self._validate_path(path)
expires_in = min(expires_in, 3600)
return self.s3.generate_presigned_url(
"get_object",
Params={"Bucket": self.bucket_name, "Key": path},
ExpiresIn=expires_in,
)
def delete_file(self, path: str) -> bool:
"""Delete a file from S3 storage."""
path = self._validate_path(path)
+159 -8
View File
@@ -1,4 +1,5 @@
import logging
import re
import uuid
from abc import ABC, abstractmethod
from datetime import datetime, timezone
@@ -125,32 +126,180 @@ class ToolsNamespace(NamespaceBuilder):
return "tools"
def build(
self, tools_data: Optional[Dict[str, Any]] = None, **kwargs
self,
tools_data: Optional[Dict[str, Any]] = None,
enabled_tools: Optional[Any] = None,
**kwargs,
) -> Dict[str, Any]:
"""
Build tools context with pre-executed tool results.
Build tools context: pre-executed tool results plus the enabled-tool set.
Args:
tools_data: Dictionary of pre-fetched tool results organized by tool name
e.g., {"memory": {"notes": "content", "tasks": "list"}}
enabled_tools: Names of the tools enabled for this turn. When provided,
exposed as the reserved ``tools.enabled`` list for gating
tool-specific prompt sections. Left absent (not just empty)
when the caller did not compute it, so a gate can fail open.
Returns:
Dictionary with tool results organized by tool name
Dictionary with tool results by tool name, plus ``enabled`` when known
"""
if not tools_data:
return {}
safe_data = {}
for tool_name, tool_result in tools_data.items():
safe_data: Dict[str, Any] = {}
for tool_name, tool_result in (tools_data or {}).items():
if isinstance(tool_result, (str, dict, list, int, float, bool, type(None))):
safe_data[tool_name] = tool_result
else:
logger.warning(
f"Skipping non-serializable tool result for '{tool_name}': {type(tool_result)}"
)
if enabled_tools is not None:
safe_data["enabled"] = sorted({str(t) for t in enabled_tools if t})
return safe_data
# Artifact-reference metadata is the only thing that may enter template context;
# bytes (and anything non-primitive) are never exposed through this namespace.
_ARTIFACT_METADATA_KEYS = (
"artifact_id",
"version",
"mime_type",
"filename",
"size",
"kind",
"title",
)
def _artifact_view(ref: Any) -> Optional[Dict[str, Any]]:
"""Project an artifact reference to a serializable metadata view, or None if it isn't one."""
if not isinstance(ref, dict) or not ref.get("artifact_id"):
return None
view: Dict[str, Any] = {}
for key in _ARTIFACT_METADATA_KEYS:
value = ref.get(key)
if isinstance(value, (str, int, float, bool, type(None))):
view[key] = value
# ``{{ artifacts.<name>.id }}`` is the documented accessor; mirror artifact_id to id.
view["id"] = view.get("artifact_id")
return view
class ArtifactsNamespace(NamespaceBuilder):
"""Artifact references namespace: {{ artifacts.<name>.id }} / .mime_type / .filename + artifact(id)."""
@property
def namespace_name(self) -> str:
return "artifacts"
def build(
self,
artifacts_data: Optional[Dict[str, Any]] = None,
artifact_parent: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Dict[str, Any]:
"""Build the artifacts namespace from named references plus a parent-scoped ``artifact(id)`` lookup.
Args:
artifacts_data: Map of output-variable name -> artifact reference dict.
artifact_parent: Parent scope ({"conversation_id"|"workflow_run_id"}) for ``artifact(id)``.
Returns:
A dict of named metadata views plus an ``artifact`` callable; never any bytes.
"""
context: Dict[str, Any] = {}
for name, ref in (artifacts_data or {}).items():
view = _artifact_view(ref)
if view is not None:
context[str(name)] = view
context["artifact"] = self._make_lookup(artifact_parent or {})
return context
def _make_lookup(self, parent: Dict[str, Any]):
"""Return an ``artifact(id)`` helper that resolves metadata scoped to the run/conversation parent."""
conversation_id = parent.get("conversation_id")
workflow_run_id = parent.get("workflow_run_id")
def artifact(artifact_id: Any) -> Dict[str, Any]:
"""Resolve one artifact's metadata by id, scoped to this parent (never cross-tenant)."""
if not artifact_id or (conversation_id is None and workflow_run_id is None):
return {}
try:
from application.storage.db.repositories.artifacts import (
ArtifactsRepository,
)
from application.storage.db.session import db_readonly
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
row = repo.get_artifact_in_parent(
str(artifact_id),
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
)
if row is None:
return {}
version = repo.get_version(str(artifact_id), row.get("current_version", 1))
except Exception:
logger.exception("Failed to resolve artifact %s in namespace", artifact_id)
return {}
view = {
"artifact_id": str(row["id"]),
"id": str(row["id"]),
"version": row.get("current_version"),
"kind": row.get("kind"),
"title": row.get("title"),
}
if version is not None:
for key in ("mime_type", "filename", "size"):
value = version.get(key)
if isinstance(value, (str, int, float, bool, type(None))):
view[key] = value
return view
return artifact
class AttachmentsNamespace(NamespaceBuilder):
"""Attached-file metadata namespace: {{ attachments.files }} (name / mime_type / size)."""
@property
def namespace_name(self) -> str:
return "attachments"
def build(self, attachments: Optional[list] = None, **kwargs) -> Dict[str, Any]:
"""Expose metadata for files attached to this message.
Only the filename, MIME type, and size enter the prompt; file bytes and
any extracted ``content`` are never surfaced through this namespace. Field
names mirror the ``artifacts`` namespace (``filename`` / ``mime_type`` / ``size``).
"""
if not attachments:
return {}
files = []
for att in attachments:
if not isinstance(att, dict):
continue
filename = att.get("filename") or att.get("name")
if not filename:
continue
# Filenames are user-controlled and rendered unescaped into the prompt;
# strip control chars/newlines and cap length so a crafted name cannot
# inject fake markdown sections regardless of the upstream upload path.
clean = re.sub(r"[\x00-\x1f\x7f]", " ", str(filename))[:255].strip()
if not clean:
continue
entry: Dict[str, Any] = {
"filename": clean,
"mime_type": att.get("mime_type") or "application/octet-stream",
}
size = att.get("size")
if isinstance(size, int) and size > 0:
entry["size"] = size
files.append(entry)
return {"files": files} if files else {}
class NamespaceManager:
"""Manages all namespace builders and context assembly"""
@@ -160,6 +309,8 @@ class NamespaceManager:
"passthrough": PassthroughNamespace(),
"source": SourceNamespace(),
"tools": ToolsNamespace(),
"artifacts": ArtifactsNamespace(),
"attachments": AttachmentsNamespace(),
}
def build_context(self, **kwargs) -> Dict[str, Any]:
+116
View File
@@ -34,6 +34,7 @@ from application.parser.schema.base import Document
from application.storage.db.base_repository import looks_like_uuid
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.repositories.attachments import AttachmentsRepository
from application.storage.db.repositories.ingest_chunk_progress import (
IngestChunkProgressRepository,
@@ -1609,6 +1610,121 @@ def attachment_worker(self, file_info, user):
raise
def parse_document_worker(self, artifact_id, parent, user_id, options):
"""Thin Celery-task wrapper; delegates to the process-agnostic ``run_parse_document``."""
return run_parse_document(artifact_id, parent, user_id, options)
def run_parse_document(artifact_id, parent, user_id, options):
"""Parse an input artifact's bytes to a shaped result; runnable inline OR on the parsing queue.
Security: the artifact is re-resolved through the run-scoped gate here (never trusting a
raw storage path) so authz is enforced independently, in addition to the pre-enqueue check
in the tool. ``read_document`` calls this directly (in-process) when it already runs inside a
Celery worker; the web process dispatches ``parse_document`` to the parsing queue, which lands
here via ``parse_document_worker``.
"""
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.parser.document_reader import bound_parse_payload, parse_document_bytes
options = options or {}
parent = parent or {}
conversation_id = parent.get("conversation_id")
workflow_run_id = parent.get("workflow_run_id")
if conversation_id is None and workflow_run_id is None:
return {"status": "error", "error": "parse_document requires a conversation_id or workflow_run_id."}
# Re-resolve through the parent-scoped gate so a forged/cross-run id is rejected
# here too; resolve a short ref to an id within this parent only.
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
resolved_id = resolve_artifact_id(
repo, artifact_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
)
artifact = (
repo.get_artifact_in_parent(
resolved_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
)
if resolved_id is not None
else None
)
if artifact is None:
return {"status": "error", "error": f"input artifact {artifact_id} not found in this conversation/run."}
version = repo.get_version(resolved_id, artifact["current_version"])
except Exception:
logging.error("run_parse_document: failed to resolve input artifact", exc_info=True)
return {"status": "error", "error": f"failed to load input artifact {artifact_id}."}
if not version or not version.get("storage_path"):
return {"status": "error", "error": f"input artifact {artifact_id} has no stored content."}
display_name = version.get("filename") or artifact.get("title") or str(resolved_id)
filename = safe_filename(display_name)
try:
data = StorageCreator.get_storage().get_file(version["storage_path"]).read()
except Exception:
logging.error("run_parse_document: failed to read input artifact bytes", exc_info=True)
return {"status": "error", "error": f"failed to read input artifact {artifact_id}."}
# Parse returns the FULL content (no max_chars/window here) so the persisted artifact is the
# complete parse; all view-bounding is applied by ``bound_parse_payload`` below.
result = parse_document_bytes(
data,
filename,
output=options.get("output", "markdown"),
ocr=options.get("ocr", "auto"),
pages=options.get("pages"),
engine=options.get("engine", "auto"),
include_tables=bool(options.get("include_tables", True)),
)
if result.get("error"):
return {"status": "error", "error": result["error"]}
payload = {"status": "ok", **result}
if options.get("persist"):
# Persist the FULL shaped result by reference (bytes live in the artifact); only the
# bounded view computed below rides back through the Redis result backend.
artifact_ref = _persist_parse_result(result, display_name, user_id, parent, options)
if isinstance(artifact_ref, dict) and artifact_ref.get("error"):
payload["artifact_error"] = artifact_ref["error"]
elif artifact_ref is not None:
payload["artifact"] = artifact_ref
# Bound the Redis-backed VIEW across all shapes: content is capped by max_chars (else the
# default head+tail window), chunks are count/length-capped, and structured (needed for
# json_schema validation) rides back as-is. The FULL result already lives in the artifact.
payload = bound_parse_payload(payload, max_chars=options.get("max_chars"))
return payload
def _persist_parse_result(result, title, user_id, parent, options):
"""Persist the full shaped parse result as an owner/parent-scoped ``data`` artifact; return its ref."""
from application.sandbox.artifacts_capture import QuotaExceeded, persist_new_artifact
try:
data = json.dumps(result).encode("utf-8")
except (TypeError, ValueError):
logging.error("parse_document_worker: parse result is not JSON-serializable", exc_info=True)
return {"error": "parse result is not JSON-serializable."}
base = safe_filename(title) or "document"
filename = f"{base}.parsed.json"
try:
return persist_new_artifact(
user_id=user_id,
kind="data",
data=data,
filename=filename,
mime_type="application/json",
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:
return {"error": str(exc)}
def agent_webhook_worker(self, agent_id, payload):
"""Process the webhook payload for an agent.
+9 -1
View File
@@ -32,7 +32,8 @@ services:
worker:
build: ../application
command: celery -A application.app.celery worker -l INFO
# `parsing` queue carries read_document/parse_document; required for its await to resolve.
command: celery -A application.app.celery worker -l INFO -Q docsgpt,parsing
env_file:
- ../.env
environment:
@@ -47,6 +48,13 @@ services:
postgres:
condition: service_healthy
# 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
ports:
+9 -1
View File
@@ -39,7 +39,8 @@ services:
worker:
user: root
image: arc53/docsgpt:develop
command: celery -A application.app.celery worker -l INFO -B
# `parsing` queue carries read_document/parse_document; required for its await to resolve.
command: celery -A application.app.celery worker -l INFO -B -Q docsgpt,parsing
env_file:
- ../.env
environment:
@@ -58,6 +59,13 @@ services:
postgres:
condition: service_healthy
# 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
ports:
+4 -2
View File
@@ -39,7 +39,10 @@ services:
worker:
user: root
build: ../application
command: celery -A application.app.celery worker -l INFO -B
# 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`.
command: celery -A application.app.celery worker -l INFO -B -Q docsgpt,parsing
env_file:
- ../.env
environment:
@@ -82,4 +85,3 @@ services:
volumes:
postgres_data:
@@ -85,7 +85,9 @@ spec:
containers:
- name: docsgpt-worker
image: arc53/docsgpt
command: ["celery", "-A", "application.app.celery", "worker", "-l", "INFO", "-n", "worker.%h"]
# `parsing` queue carries read_document/parse_document; required for its await to resolve.
# For heavy/OCR parsing, run a separate deployment with `-Q parsing` (and GPU env).
command: ["celery", "-A", "application.app.celery", "worker", "-l", "INFO", "-n", "worker.%h", "-Q", "docsgpt,parsing"]
resources:
limits:
memory: "4Gi"
@@ -0,0 +1,112 @@
# 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).
#
# SINGLE TRUST DOMAIN: all sessions in this pod share one container/uid and are
# isolated by working directory only (per-session cwd) — not by a kernel/OS
# boundary. This image ships an env-scrubbing kernelspec under the distinct name
# "docsgpt-python"; the app (docsgpt-api / docsgpt-worker) MUST select it via
# SANDBOX_KERNEL_NAME=docsgpt-python (already set in docsgpt-deploy.yaml) so the
# scrubber runs. The stock "python3" spec stays untouched, so it can never
# shadow the scrubbing kernel. Sibling workspaces are still readable under the
# shared uid and kernels share one address space. Do NOT inject app secrets
# (DB/LLM creds) into this pod's env; the runner needs none. For cross-tenant /
# untrusted multi-tenant workloads use a per-session VM via
# SANDBOX_BACKEND=daytona instead.
apiVersion: apps/v1
kind: Deployment
metadata:
name: docsgpt-sandbox
spec:
replicas: 1
selector:
matchLabels:
app: docsgpt-sandbox
template:
metadata:
labels:
app: docsgpt-sandbox
spec:
# runtimeClassName: gvisor
securityContext:
runAsNonRoot: true
# Numeric UID the kubelet can verify as non-root; MUST match the sandbox
# image's `USER` (deployment/sandbox/Dockerfile). Without it the pod
# fails to start with CreateContainerConfigError.
runAsUser: 10001
seccompProfile:
type: RuntimeDefault
containers:
- name: docsgpt-sandbox
image: arc53/docsgpt-sandbox
ports:
- containerPort: 8888
securityContext:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: true
capabilities:
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
value: /tmp/jupyter-data
resources:
limits:
memory: "1Gi"
cpu: "1"
requests:
memory: "256Mi"
cpu: "250m"
volumeMounts:
# Per-session workspaces + Jupyter runtime files on an emptyDir;
# the root FS stays read-only everywhere else.
- name: scratch
mountPath: /tmp
volumes:
- name: scratch
emptyDir: {}
---
apiVersion: v1
kind: Service
metadata:
name: docsgpt-sandbox
spec:
selector:
app: docsgpt-sandbox
ports:
- protocol: TCP
port: 8888
targetPort: 8888
# Cluster-internal only: the runner must never be exposed outside the cluster.
type: ClusterIP
@@ -0,0 +1,102 @@
# Egress / SSRF NetworkPolicy for the docsgpt-sandbox runner.
#
# The sandbox executes arbitrary LLM-authored code, so app-level URL checks
# cannot contain it -- the code opens its own sockets. The hardened container
# runs without NET_ADMIN and cannot self-apply iptables, so egress is enforced
# here at the network layer: broad public-internet egress is allowed, but all
# private / link-local / metadata ranges are DENIED. Ingress is restricted to
# the backend and worker pods only.
#
# Requires a NetworkPolicy-enforcing CNI (Calico, Cilium, etc.). The default
# kube-proxy / flannel setups do NOT enforce these rules. Apply with:
# kubectl apply -f deployment/k8s/network-policies/sandbox-egress-policy.yaml
#
# Pods this selects must carry the label `app: docsgpt-sandbox`.
apiVersion: networking.k8s.io/v1
kind: NetworkPolicy
metadata:
name: docsgpt-sandbox-egress
labels:
app: docsgpt-sandbox
spec:
podSelector:
matchLabels:
app: docsgpt-sandbox
policyTypes:
- Ingress
- Egress
ingress:
# Only the backend/worker may reach the sandbox gateway (TCP 8888).
- from:
- podSelector:
matchLabels:
app: docsgpt-api
- podSelector:
matchLabels:
app: docsgpt-worker
ports:
- protocol: TCP
port: 8888
egress:
# DNS to cluster resolvers (kube-dns). Restricted to UDP/TCP 53 so the
# broad-egress rule below does not need to whitelist the resolver IP.
# Kept for clusters WITHOUT NodeLocal DNSCache, where the resolver is a pod.
- to:
- namespaceSelector: {}
ports:
- protocol: UDP
port: 53
- protocol: TCP
port: 53
# NodeLocal DNSCache: on clusters running it, the stub resolver is a
# link-local IP (commonly 169.254.20.10), NOT a pod, and it sits inside the
# 169.254.0.0/16 range carved out of the broad-egress rule below -- so
# without this explicit allow ALL DNS is dropped and every hostname lookup
# (pip install, public APIs) times out. NetworkPolicy rules are
# additive-allow, so this narrowly re-permits port 53 to that one IP.
# Operators whose NodeLocal DNSCache uses a different link-local IP must
# adjust this cidr.
- to:
- ipBlock:
cidr: 169.254.20.10/32 # NodeLocal DNSCache link-local resolver
ports:
- protocol: UDP
port: 53
- protocol: TCP
port: 53
# Broad public-internet egress, with every private / link-local / ULA /
# carrier-grade-NAT / documentation range carved out via `except`. The
# cloud metadata IP 169.254.169.254 is inside the 169.254.0.0/16 hole, so
# it is unreachable. SSRF to internal services is therefore blocked while
# legitimate outbound calls (package installs, public APIs) still work.
- to:
- ipBlock:
cidr: 0.0.0.0/0
except:
- 0.0.0.0/8 # "this network" (RFC1122)
- 10.0.0.0/8 # RFC1918 private
- 172.16.0.0/12 # RFC1918 private
- 192.168.0.0/16 # RFC1918 private
- 169.254.0.0/16 # link-local incl. 169.254.169.254 metadata
- 100.64.0.0/10 # RFC6598 carrier-grade NAT
- 127.0.0.0/8 # loopback
- 192.0.0.0/24 # IETF protocol assignments
- 192.0.2.0/24 # TEST-NET-1 (documentation)
- 198.18.0.0/15 # benchmarking
- 198.51.100.0/24 # TEST-NET-2 (documentation)
- 203.0.113.0/24 # TEST-NET-3 (documentation)
- 224.0.0.0/4 # multicast
- 240.0.0.0/4 # reserved
- 255.255.255.255/32 # limited broadcast (RFC1122)
# IPv6 public egress with private / ULA / link-local ranges carved out.
# The IPv6 metadata endpoint (fd00:ec2::254) sits inside the fc00::/7 ULA
# hole and is unreachable. Drop this block on IPv4-only clusters.
- to:
- ipBlock:
cidr: "::/0"
except:
- "::1/128" # loopback
- "fc00::/7" # unique local (ULA) incl. fd00:ec2::254
- "fe80::/10" # link-local
- "ff00::/8" # multicast
- "2001:db8::/32" # documentation
@@ -0,0 +1,121 @@
# Optional egress-firewall overlay for the docsgpt-sandbox runner (compose).
#
# Docker Compose cannot express L3 egress filtering the way a Kubernetes
# NetworkPolicy can, so SSRF containment in compose deployments is delivered by
# taking away the runner's direct internet route and forcing its outbound
# traffic through an egress-gateway sidecar that DENIES private / link-local /
# metadata ranges and ALLOWS the public internet.
#
# HOW THIS OVERLAY WORKS
# The 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
# (on its own `sandbox-egress-out` network) and its ACL denies private
# destinations. We flip a network's scalar property by KEY (a well-defined
# Compose merge); we do NOT try to remove an item from the service's `networks`
# LIST, which Compose cannot express (it unions lists) -- the reason a naive
# `networks: [sandbox-net]` override silently leaves the runner on its old net.
#
# WHAT THIS OVERLAY DOES NOT CONTAIN (read before enabling the sandbox):
# The runner STILL shares `sandbox-net` with `backend` and `worker` -- that is
# its control path and it cannot be removed without breaking code execution. A
# shared Docker network is bidirectional and Compose cannot sever it
# one-directionally, so arbitrary sandbox code can still open sockets to
# `backend:7091` and reach the worker. The internal flip contains raw sockets to
# the internet / host / RFC1918 / metadata, NOT to backend/worker on sandbox-net.
# (The Kubernetes NetworkPolicy DOES block this, via its RFC1918 egress
# carve-out; compose has no equivalent.)
#
# MITIGATION -- required when enabling the sandbox (do the first, ideally both):
# - Run the backend with real authentication (`AUTH_TYPE` != none / a real auth
# provider) so a reachable API rejects unauthenticated requests. Without it,
# runner->backend reach is a free control-plane bypass.
# - Add a host-firewall DROP for runner->backend/worker on sandbox-net
# (approach (1) below).
#
# Two layers below: (2) is what this overlay wires up (internet egress via the
# deny-private proxy); (1) is the host-firewall DROP that closes the
# runner->backend/worker gap. Apply both for the strongest posture.
#
# (1) Host-firewall DROP for runner->backend/worker (closes the gap above). Run
# as root on the Docker host AFTER `... up`. The project name is pinned to
# `docsgpt-oss` (see `name:` in docker-compose.yaml), so the shared control
# net is `docsgpt-oss_sandbox-net`. `docker compose ... ps -q` resolves each
# container id regardless of its generated name (there is no fixed
# `container_name`), and the `index` function is REQUIRED because a
# Go-template dotted key cannot contain hyphens -- the old
# `{{.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.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)")
# WK=$( docker inspect -f "{{(index .NetworkSettings.Networks \"$NET\").IPAddress}}" "$(docker compose $CF ps -q worker)")
# iptables -I DOCKER-USER -s "$SBX" -d "$BE" -j DROP # runner -> backend (incl. :7091)
# iptables -I DOCKER-USER -s "$SBX" -d "$WK" -j DROP # runner -> worker
# (These container-to-container DROPs require the host's
# `bridge-nf-call-iptables=1`, Docker's default, so bridged traffic
# traverses DOCKER-USER.) If a lookup is empty (renamed project, or the
# container is not up yet), resolve by hand instead of running a broken
# command: `docker compose $CF ps` to list ids, then `docker inspect
# <runner-id>` / `<backend-id>`, read the sandbox-net
# `.NetworkSettings.Networks.*.IPAddress`, and add the same DROP rules.
# Re-run after any recreate -- container IPs are not stable across `up`.
# (Internet / RFC1918 / metadata egress is already gone via the internal
# flip, so no separate SSRF drop is needed here.)
#
# (2) Egress-gateway sidecar (this overlay): the runner has NO direct internet
# route (see above); its only path out is the deny-private forward proxy.
# Point the runner's HTTP(S) client at the proxy via the env vars below.
# This does NOT stop runner->backend/worker on sandbox-net -- pair it with
# backend auth and/or approach (1).
#
# 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
#
# 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
# code does (there is simply no route). HTTP(S) that honors the proxy env is
# filtered by the proxy ACL. This does NOT block raw egress to `backend:7091` /
# the worker on the shared `sandbox-net` (see WHAT THIS OVERLAY DOES NOT CONTAIN
# / MITIGATION above): enable backend auth and/or apply approach (1). On a
# multi-tenant or untrusted deployment, prefer the Kubernetes NetworkPolicy,
# which blocks the internal path too.
services:
# Forward proxy that denies private/link-local/metadata destinations and
# allows the public internet. Replace the image/command with your proxy of
# choice (tinyproxy, squid, mitmproxy with a deny rule, etc.); the ACL must
# DENY 10/8, 172.16/12, 192.168/16, 169.254/16, 127/8 and ALLOW the rest.
sandbox-egress-proxy:
image: ghcr.io/example/egress-deny-private:latest # replace with your proxy image
restart: unless-stopped
networks:
- sandbox-net # reachable by the runner (control-plane net)
- sandbox-egress-out # the ONLY container with an internet route
docsgpt-sandbox:
# 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
- NO_PROXY=localhost,127.0.0.1
networks:
# Flip the runner's egress net to internal: no direct route to the internet,
# host, RFC1918, or cloud metadata. (Merged by key over the base definition.)
sandbox-egress:
internal: true
# Internet-facing route for the proxy ONLY.
sandbox-egress-out:
driver: bridge
@@ -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: {}
+75
View File
@@ -0,0 +1,75 @@
# docsgpt-sandbox: the always-on code-execution runner.
#
# One container runs a Jupyter Kernel Gateway; each sandbox session is an
# in-process kernel (a child process), NOT a child container. This image does
# NOT mount the Docker socket and never spawns containers — that avoids the
# host-root risk of Docker-in-Docker and is the core security win.
#
# The app (backend/worker) is the CLIENT and reaches this service over
# HTTP + WebSocket via SANDBOX_GATEWAY_URL.
#
# Pre-baked doc libs are all permissive (MIT/BSD/Apache) for speed/reliability;
# runtime `pip install` still works because egress is open. PyMuPDF and any
# other AGPL lib are intentionally excluded; Docling is deferred to its own slice.
#
# HARDENING (separate slice — NOT done here): run under the gVisor `runsc`
# runtime, add network-layer SSRF blocks (drop RFC1918 / link-local /
# 169.254.169.254), seccomp profile, read-only root FS + quota'd scratch dir,
# and cgroup CPU/mem/PID caps wired from SANDBOX_MEMORY / SANDBOX_CPUS.
FROM python:3.12-slim
# Non-root user for the runner (untrusted code runs as this UID).
RUN useradd --create-home --uid 10001 sandbox
# Exact pins for reproducible builds and no license drift (all MIT/BSD/Apache).
RUN pip install --no-cache-dir \
jupyter-kernel-gateway==3.0.1 \
ipykernel==6.29.5 \
python-pptx==1.0.2 \
python-docx==1.1.2 \
openpyxl==3.1.5 \
reportlab==4.2.5 \
pandas==2.2.3 \
matplotlib==3.9.2
# Docling (MIT) — OFF by default because it pulls torch + models and makes the
# image multi-GB. NOTE: document parsing now runs on the Celery `parsing` worker
# (see read_document / parse_document_worker), NOT in this sandbox, so this
# INSTALL_DOCLING build-arg path is effectively unused for the read_document flow;
# it remains only for sandbox code that opts into Docling explicitly. The build
# mechanics are kept for that case. Docling is MIT and uses its own PDF backend;
# PyMuPDF (AGPL) is NOT installed here. The base image stays docling-free.
ARG INSTALL_DOCLING=false
RUN if [ "$INSTALL_DOCLING" = "true" ]; then \
pip install --no-cache-dir docling==2.8.3; \
fi
# Env-scrubbing kernel launcher + custom kernelspec. The launcher re-execs
# ipykernel under a minimal allowlisted env (env -i) so NO secret in the
# gateway's environment (*_API_KEY, *_TOKEN, POSTGRES_URI, the gateway auth
# token, ...) ever reaches kernel code. The kernelspec ships under a DISTINCT
# name ("docsgpt-python"), so the app selects it with SANDBOX_KERNEL_NAME and it
# is never shadowed by the stock ipykernel "python3" spec regardless of the
# python prefix. The stock "python3" spec is left untouched (no overwrite, no
# kernelspec-name precedence to rely on). SECURITY: never give this image
# `env_file: ../.env` -- the scrubber blocks exfil from the kernel, but the
# 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
# the user is non-root without resolving /etc/passwd. This uid MUST match
# `runAsUser` in deployment/k8s/deployments/sandbox-deploy.yaml.
USER 10001
WORKDIR /home/sandbox
EXPOSE 8888
# 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"]
+284
View File
@@ -0,0 +1,284 @@
# docsgpt-sandbox runner
Opt-in Jupyter Kernel Gateway that executes sandboxed LLM code. The DocsGPT
backend/worker is the **client** and connects over HTTP + WebSocket via
`SANDBOX_GATEWAY_URL`. Each session is an **in-process kernel** (child process),
never a child container; the Docker socket is **not** mounted.
## Enabling code execution (opt-in)
The runner is **opt-in**. 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
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
```
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.yaml \
-f deployment/optional/docker-compose.optional.sandbox-egress.yaml up
```
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
Read this before pointing untrusted or multi-tenant workloads at the runner.
A single Jupyter runner is **one trust domain**. Every session is an in-process
kernel under **one shared uid (10001)** in **one container**; sessions are
isolated by **working directory only** — each session's code runs with its cwd
set to its own `/tmp/docsgpt-sandbox/<session_id>` directory. That is a
convenience boundary, not a security boundary between sessions.
What this slice does close:
- **Env-secret exfil is closed.** The custom kernelspec
(`kernels/docsgpt-python/kernel.json` → `/opt/docsgpt/kernel-launch.sh`)
re-execs ipykernel under a minimal allowlisted env (`env -i` keeping only
`PATH`, `HOME`, `LANG`, `JUPYTER_RUNTIME_DIR`, `JUPYTER_DATA_DIR`). The image
installs this spec under the **distinct name `docsgpt-python`** and the app
selects it via `SANDBOX_KERNEL_NAME=docsgpt-python`; because the name is
distinct, it is **never shadowed** by the stock ipykernel `python3` spec
(kernelspec name resolution prefers `sys.prefix/share` over
`/usr/local/share`, so reusing `python3` would silently fall back to the
unscrubbed stock spec on a different python prefix). The stock `python3` spec
is left untouched. So even though the gateway process inherits the operator's
full environment, **no `*_API_KEY` / `*_TOKEN` / `POSTGRES_URI` / gateway auth
token reaches kernel code** via `os.environ`, regardless of how the gateway is
launched. Loopback ZMQ reachability is preserved because `{connection_file}`
is forwarded untouched.
- **Per-session workspace perms.** The workspace root and each session dir are
created `0700` (defense-in-depth). Under one shared uid this does **not** stop
a sibling session from reading another's files — it only narrows exposure to
other uids on the box.
Residual gaps (treat all sessions in one runner as mutually trusting):
- **Sibling-workspace reads.** All kernels run as the same uid, so one session's
code can read another session's files (and `/tmp`) despite `0700`. Distinct
uids / per-session VMs are required to close this.
- **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
egress / SSRF* below), never by the runner itself.
For real per-tenant isolation (cross-tenant or untrusted code), use the
**Daytona backend** (`SANDBOX_BACKEND=daytona`), which gives each session its
own VM. To harden the self-hosted Jupyter runner as a whole (host protection +
egress), layer the **gVisor `runsc` runtime**, the **NetworkPolicy**, and a
**host firewall** as documented below — those protect the host and constrain
egress; they do **not** create a boundary between sessions inside one runner.
## Run standalone for dev
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 -e SANDBOX_GATEWAY_AUTH_TOKEN=devtoken docsgpt-sandbox
# in the app's .env: SANDBOX_GATEWAY_URL=http://localhost:8888
# SANDBOX_GATEWAY_AUTH_TOKEN=devtoken
```
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
```
`--ZMQChannelsWebsocketConnection.limit_rate=False` raises the iopub data-rate
limit so large `get_file` base64 payloads aren't silently truncated. (On older
gateways the trait may live elsewhere; the client's `get_file` integrity check
catches any truncation regardless.)
A bare-venv gateway uses the **stock** `python3` kernelspec, which inherits the
gateway's full env (no secret scrubbing). The default `SANDBOX_KERNEL_NAME` is
`python3`, so plain venv dev gets no scrubbing — acceptable for single-trust
dev. The Docker image instead ships the env-scrubbing spec under the distinct
name `docsgpt-python` (see *Isolation model*) and the runner stack sets
`SANDBOX_KERNEL_NAME=docsgpt-python`. To get the scrubbing behavior in a venv,
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.
## Gateway authentication
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 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)
The `artifact` tool renders `presentation` / `document` / `spreadsheet` / `pdf`
specs by running a fixed renderer **inside the sandbox** that imports
`python-pptx`, `python-docx`, `openpyxl`, and `reportlab` (HTML/markdown need no
library). The self-hosted Jupyter runner inherits these from the backend venv,
but Daytona's default snapshot is a plain Python image — so under
`SANDBOX_BACKEND=daytona` those renders fail with `render failed: ExecutionError`
(a `ModuleNotFoundError` raised inside the sandbox). HTML/markdown still work.
Bake the libraries into a Daytona snapshot once, then point `DAYTONA_SNAPSHOT`
at it:
```bash
# Reads DAYTONA_API_KEY / DAYTONA_API_URL / DAYTONA_TARGET from .env:
python scripts/build_daytona_snapshot.py # builds "docsgpt-artifacts-py312"
# then in .env:
# DAYTONA_SNAPSHOT=docsgpt-artifacts-py312
```
The snapshot lives in **your** Daytona account, so each deployment builds its
own — the script is idempotent and skips if the name already exists. Keep the
pins in `scripts/build_daytona_snapshot.py` in sync with the backend venv so the
Daytona render output matches the Jupyter-backend output.
## Document reading (parsing worker — not the sandbox)
Document reading no longer runs in this sandbox. The `read_document` tool and the
workflow native-file extract branch enqueue a `parse_document` Celery task that
parses the document **in the backend** (Docling, already in
`application/requirements.txt`) and awaits the result. The task is routed to a
dedicated **`parsing` queue** (`settings.DOCUMENT_PARSE_QUEUE`, default
`"parsing"`) so a parse enqueued from inside a Celery worker (headless/scheduled
agent) is served by a separate worker and never self-deadlocks the awaiting one.
Run a dedicated parsing worker that consumes the `parsing` queue:
```bash
celery -A application.app.celery worker -Q parsing -l INFO
```
It can be GPU-enabled with its own env (`DOCLING_OCR_ENABLED=true` plus GPU
libraries) so OCR-heavy parsing runs on a separate, optionally larger pool.
**Dev / single-worker setups:** without a dedicated parsing worker the default
worker must also consume `parsing`, or the tool's await never resolves:
```bash
celery -A application.app.celery worker -Q docsgpt,parsing -l INFO
```
Tuning settings: `DOCUMENT_PARSE_TIMEOUT` (seconds the tool awaits before
degrading to an error), `DOCUMENT_PARSE_MAX_BYTES` (per-document byte cap; 0
reuses `SANDBOX_MAX_INPUT_BYTES`).
## Network egress / SSRF
The runner allows **broad outbound egress** (so sandboxed code can `pip install`
and call public APIs) but private, link-local, and cloud-metadata ranges **MUST
be blocked at the network layer**. This is not optional: the sandbox executes
arbitrary LLM-authored code, which opens its own sockets — app-level URL checks
(the `mcp_tool.py` approach) cannot contain it. Without a network-layer block,
sandbox code can reach `169.254.169.254` (cloud instance metadata / credentials)
and internal services on the private network.
The hardened container runs **without `NET_ADMIN`**, so it cannot self-apply
`iptables`. Enforcement therefore lives in deployment config:
- **Kubernetes** — apply
[`deployment/k8s/network-policies/sandbox-egress-policy.yaml`](../k8s/network-policies/sandbox-egress-policy.yaml).
It allows `0.0.0.0/0` egress with `except` carve-outs for RFC1918
(`10/8`, `172.16/12`, `192.168/16`), link-local (`169.254/16`, which contains
`169.254.169.254`), loopback, CGNAT, documentation/test ranges, and the IPv6
ULA/link-local equivalents — and restricts ingress to the API/worker pods on
TCP 8888. It requires a policy-enforcing CNI (Calico, Cilium, …); plain
flannel/kube-proxy will silently not enforce it. The matching sandbox pod is
[`deployment/k8s/deployments/sandbox-deploy.yaml`](../k8s/deployments/sandbox-deploy.yaml)
(label `app: docsgpt-sandbox`).
```bash
kubectl apply -f deployment/k8s/deployments/sandbox-deploy.yaml
kubectl apply -f deployment/k8s/network-policies/sandbox-egress-policy.yaml
```
- **docker-compose** — compose cannot express L3 egress filtering natively. The
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
[`deployment/optional/docker-compose.optional.sandbox-egress.yaml`](../optional/docker-compose.optional.sandbox-egress.yaml),
which flips `sandbox-egress` to `internal: true` (removing the runner's direct
internet/RFC1918/metadata route entirely) and forces egress through a
deny-private **egress-gateway proxy** sidecar. That `internal` flip is what
contains **raw sockets to the internet/host/RFC1918/metadata** (a forward proxy
only filters code that honors `HTTP(S)_PROXY`).
**What compose canNOT contain:** the runner stays on `sandbox-net` with the
backend and worker — that is its control path, and a shared Docker network is
bidirectional, so Compose cannot sever it one-directionally. Sandbox code can
therefore still open sockets to `backend:7091` and the worker. This is a real
gap the Kubernetes NetworkPolicy closes (via its RFC1918 egress carve-out) but
compose cannot. **Mitigate it** when enabling the sandbox: run the backend with
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). 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/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
above.
+38
View File
@@ -0,0 +1,38 @@
#!/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-launch.sh re-execs ipykernel under a
# scrubbed `env -i` allowlist that excludes the token, so it is NOT in kernel code's
# os.environ. NOTE: this is not a hard secret boundary -- the gateway and kernels
# share one container and uid, so kernel code can still read the gateway's env via
# /proc/<gateway_pid>/environ. The token limits the blast radius (an outside caller
# cannot drive the gateway), but the Jupyter runner is a single trust domain: for
# real isolation against untrusted code use the Daytona backend or run under gVisor
# (see deployment/sandbox/README.md).
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
+20
View File
@@ -0,0 +1,20 @@
#!/bin/sh
# Env-scrubbing launcher for the docsgpt-sandbox ipykernel.
#
# The gateway process inherits the operator's full environment, which can carry
# secrets (*_API_KEY, *_TOKEN, POSTGRES_URI, the gateway auth token, ...). Stock
# kernels inherit that env verbatim, so LLM-authored code could read it via
# os.environ. This wrapper re-execs ipykernel under a MINIMAL allowlisted env so
# NO secret reaches kernel code, regardless of how the gateway was launched.
#
# Only what ipykernel needs is kept: PATH (find python), HOME (~/.ipython etc),
# LANG (encoding), and the Jupyter runtime/data dirs (writable tmpfs paths). The
# {connection_file} the gateway passes is forwarded via "$@" so loopback ZMQ
# reachability is preserved -- do NOT drop or rewrite those args.
exec env -i \
PATH="${PATH}" \
HOME="${HOME}" \
LANG="${LANG}" \
JUPYTER_RUNTIME_DIR="${JUPYTER_RUNTIME_DIR}" \
JUPYTER_DATA_DIR="${JUPYTER_DATA_DIR}" \
python -m ipykernel_launcher "$@"
@@ -0,0 +1,12 @@
{
"argv": [
"/opt/docsgpt/kernel-launch.sh",
"-f",
"{connection_file}"
],
"display_name": "Python 3 (docsgpt-sandbox, scrubbed env)",
"language": "python",
"metadata": {
"debugger": true
}
}
+22
View File
@@ -35,6 +35,16 @@ When the agent completes its task, it stores the result in the shared state.
- *Custom*: You can set this to something meaningful, like `summary` or `translated_text`.
- **Streaming**: If "Stream to user" is enabled, the output is sent to the user in real-time as it is generated, in addition to being saved to the state.
### Documents
An agent node can receive documents as inputs. Choose which documents the node sees:
- **All**: every document attached to the run.
- **None**: no documents.
- **Choose**: a specific set that you select.
For how a chosen document reaches the model, the node can pass it natively (send the file to a model that accepts files) or extract it to text first. The default picks automatically based on the model and the file type.
---
## Set State Node
@@ -63,3 +73,15 @@ You can define multiple operations in a single node. Each operation targets a sp
- **Loop Counters**: Use a *Set State* node to initialize a counter (`i = 0`) before a loop, and another to increment it inside the loop.
- **Accumulators**: Use *Append* to collect results from multiple parallel branches into a single list.
- **Renaming**: Copy the output of a previous node to a more generic name (e.g., set `context` to `{{search_results}}`) so subsequent nodes can use a standard variable name.
---
## Code Node
The **Code Node** runs a script in a sandboxed session bound to the workflow run. Use it to transform data, parse files, cross-check documents, or build a report that later nodes consume.
- **Code**: the script to run. The workflow state is available to the script as data, so it can read variables, compute, and write results back.
- **Inputs**: documents or artifacts to place in the workspace before the script runs. Each input accepts a produced artifact reference, a full artifact id, or an attached file.
- **Outputs**: any files the script writes are captured as artifacts and surfaced in the run view. Write a small JSON value back to the state to pass a decision or summary to later nodes.
Runs are sandboxed and have a fixed time limit, so keep each step focused. See [Artifacts and Code Execution](/Tools/artifacts-and-code-execution) for the sandbox backends and configuration.
+4
View File
@@ -11,6 +11,10 @@ export default {
"title": "🖥️ Remote Device",
"href": "/Tools/remote-device"
},
"artifacts-and-code-execution": {
"title": "📦 Artifacts and Code Execution",
"href": "/Tools/artifacts-and-code-execution"
},
"creating-a-tool": {
"title": "🛠️ Creating a Custom Tool",
"href": "/Tools/creating-a-tool"
@@ -0,0 +1,101 @@
---
title: 📦 Artifacts and Code Execution
description: Generate editable documents, run sandboxed code, and read files with the DocsGPT artifact, code executor, and read document tools.
---
import { Callout } from 'nextra/components';
# Artifacts and Code Execution
DocsGPT can generate documents, run code, and read files for you during a chat or inside an agent workflow. Three built-in tools work together for this:
- **Artifact** builds and edits documents such as slide decks, Word documents, spreadsheets, PDFs, and HTML.
- **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.
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
An **artifact** is a file the assistant produces that you can open, download, and edit later. Artifacts are versioned: every edit appends a new version at the same artifact, so you can review or restore an earlier one. Bytes are stored on the server and are never passed through the model, so large files stay fast and cheap to work with.
In a normal chat, a produced artifact appears as a chip under the assistant's reply (for example, "Artifact"). Open it to:
- preview the content in a side panel (HTML and similar kinds render inline),
- download the file,
- ask for changes, which create a new version.
### The Artifact tool
The Artifact tool keeps a structured spec as the source of truth and renders the file from it. Editing means changing the spec and re-rendering, which keeps a clean and diffable version history. It exposes three actions:
- **create** a new artifact from a kind and a spec,
- **edit** an existing artifact with a targeted change (preferred for small edits),
- **rewrite** an artifact with a full replacement spec.
Supported kinds are presentation (.pptx), document (.docx), spreadsheet (.xlsx), pdf, and html.
## Code Executor
The Code Executor runs Python in a sandboxed, stateful session bound to your conversation (or to a workflow run). Files the code writes into the workspace are captured as artifacts automatically, so a script that produces `report.csv` gives you a downloadable artifact with no extra steps.
Key points:
- The session is stateful, so variables, imports, and files persist between calls in the same conversation while it is kept alive.
- Only a compact summary is returned to the model (an output tail plus artifact references), never raw file bytes.
- Each run has a fixed wall-clock limit of about 60 seconds. For longer work, start it in the background and check back with more calls.
- Install packages from inside the code itself when you need them.
<Callout type="warning">
By default the Code Executor runs model-authored code without an approval
prompt. Turn on per-action approval on the agent if you want a human to
confirm each run.
</Callout>
### Run inputs
You can pass files into a run. Each input accepts a short reference returned by a previous artifact action, a full artifact id, or the name or id of a file you attached to the conversation. The referenced files are placed in the workspace before the code runs.
## Read Document
Read Document parses a file into text, markdown, structured JSON (with tables), or chunks. It is meant for workflows where a node needs the contents of an uploaded or produced document, for example a compliance flow that extracts fields and validates them against a JSON schema.
Parsing runs in the DocsGPT parsing worker rather than the sandbox, so it uses the same document engine as the rest of the product and works with every sandbox backend.
## Using these tools in workflows
Inside an agent workflow the same building blocks are available:
- A **Code node** runs a script bound to the run, with the workflow state available to it as data.
- Agent and code nodes can receive selected documents as inputs and hand files to each other by reference.
- Produced artifacts are visible from the workflow run view.
See [Workflow Nodes](/Agents/nodes) for the node-level details.
## Configuration
Code execution and rendering run inside a sandbox. DocsGPT ships two backends, selected with `SANDBOX_BACKEND`:
- **`jupyter`** (default): a self-hosted runner that you operate. See the runner setup notes under `deployment/sandbox`.
- **`daytona`**: Daytona Cloud, a managed per-session sandbox. Set `DAYTONA_API_KEY` and choose a region with `DAYTONA_TARGET`.
### Render libraries on Daytona
Artifact rendering imports `python-pptx`, `python-docx`, `openpyxl`, and `reportlab` inside the sandbox. The self-hosted runner already has them. Daytona's default image does not, so presentation, document, spreadsheet, and pdf rendering fail there until you point Daytona at an image that includes them. HTML and markdown artifacts need no extra libraries and work on any image.
Build a snapshot with the libraries once, then set `DAYTONA_SNAPSHOT` to it:
```bash
# Reads DAYTONA_API_KEY / DAYTONA_API_URL / DAYTONA_TARGET from your environment:
python scripts/build_daytona_snapshot.py
# then set in your environment:
# DAYTONA_SNAPSHOT=docsgpt-artifacts-py312
```
### Useful settings
- `SANDBOX_BACKEND`: `jupyter` or `daytona`.
- `SANDBOX_EXEC_TIMEOUT`: per-run wall-clock cap in seconds.
- `SANDBOX_MAX_TTL`: upper bound on how long a kept-alive session lives.
- `DAYTONA_API_KEY`, `DAYTONA_TARGET`, `DAYTONA_SNAPSHOT`: Daytona Cloud settings.
- Per-user artifact quotas bound how many artifacts and how much storage each user can hold.
+11 -3
View File
@@ -51,7 +51,10 @@ import {
import PromptsModal from '../preferences/PromptsModal';
import Prompts from '../settings/Prompts';
import { UserToolType } from '../settings/types';
import { getToolDisplayName } from '../utils/toolUtils';
import {
getToolDisplayName,
isClassicAgentToolVisible,
} from '../utils/toolUtils';
import AgentPageHeader from './AgentPageHeader';
import AgentPreview from './AgentPreview';
import { Agent, ToolSummary } from './types';
@@ -480,6 +483,11 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) {
]);
if (!toolsResponse.ok) throw new Error('Failed to fetch tools');
const data = await toolsResponse.json();
// Hide workflow-only builtins (e.g. read_document) from the classic
// agent picker; they belong to the workflow-node picker only.
const visibleTools = (data.tools as UserToolType[]).filter(
isClassicAgentToolVisible,
);
const devicesById = new Map<
string,
{ online: boolean; last_seen_at: string | null | undefined }
@@ -498,7 +506,7 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) {
if (tool.default) return t('agents.form.toolsPopup.groupDefault');
return t('agents.form.toolsPopup.groupCustom');
};
const tools: MultiSelectPopoverItem[] = data.tools.map(
const tools: MultiSelectPopoverItem[] = visibleTools.map(
(tool: UserToolType) => {
const base: MultiSelectPopoverItem = {
id: tool.id,
@@ -537,7 +545,7 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) {
groupOrder.indexOf(a.group || '') - groupOrder.indexOf(b.group || ''),
);
setUserTools(tools);
setRawUserTools(data.tools as UserToolType[]);
setRawUserTools(visibleTools);
};
const getModels = async () => {
const response = await modelService.getModels(token);
+10 -1
View File
@@ -4,7 +4,16 @@ export type NodeType =
| 'agent'
| 'note'
| 'state'
| 'condition';
| 'condition'
| 'code';
export interface CodeNodeConfig {
code: string;
inputs: string[];
output_variable?: string;
timeout?: number;
json_schema?: Record<string, unknown>;
}
export interface ConditionCase {
name?: string;
+385 -87
View File
@@ -3,6 +3,7 @@ import 'reactflow/dist/style.css';
import {
AlertCircle,
Bot,
Code2,
Database,
Flag,
GitBranch,
@@ -67,11 +68,29 @@ import { getToolDisplayName } from '../../utils/toolUtils';
import AgentPageHeader from '../AgentPageHeader';
import { Agent } from '../types';
import { ConditionCase, WorkflowNode } from '../types/workflow';
import {
createDefaultCodeConfig,
normalizeCodeConfig,
parseCodeJsonSchemaDraft,
serializeCodeConfig,
validateCodeJsonSchema,
} from './codeNodeConfig';
import MobileBlocker from './components/MobileBlocker';
import PromptTextArea from './components/PromptTextArea';
import { buildSimpleCel, parseSimpleCel } from './simpleCel';
import NodeDocumentsControl from './components/NodeDocumentsControl';
import PromptTextArea, {
extractUpstreamVariables,
} from './components/PromptTextArea';
import {
FILE_PASSING_OPTIONS,
FilePassing,
normalizeFilePassing,
toDocumentVariableOptions,
} from './documentConfig';
import { useUndoRedo, WorkflowSnapshot } from './hooks/useUndoRedo';
import {
AgentNode,
CodeNode,
ConditionNode,
EndNode,
NoteNode,
@@ -97,6 +116,8 @@ interface AgentNodeConfig {
chunks?: string;
retriever?: string;
json_schema?: Record<string, unknown>;
input_documents?: string[];
file_passing?: FilePassing;
}
interface UserTool {
@@ -104,6 +125,9 @@ interface UserTool {
name: string;
displayName: string;
customName?: string;
// Workflow-only builtins (e.g. read_document) are kept here; the classic
// agent picker filters them out.
workflow_only?: boolean;
}
function validateJsonSchemaConfig(schema: unknown): string | null {
@@ -158,79 +182,6 @@ function canReachEnd(
.some((e) => canReachEnd(e.target, edges, nodeIds, endIds, visited));
}
function parseSimpleCel(expression: string): {
variable: string;
operator: string;
value: string;
} {
const trimmedExpression = expression.trim();
let match = trimmedExpression.match(
/^(\w+)\.(contains|startsWith)\(["'](.*)["']\)$/,
);
if (match) return { variable: match[1], operator: match[2], value: match[3] };
match = trimmedExpression.match(/^(\w+)\.(contains|startsWith)\((.*)\)$/);
if (match) {
const rawValue = match[3].trim();
const unquotedValue = rawValue.replace(/^["'](.*)["']$/, '$1');
return {
variable: match[1],
operator: match[2],
value: unquotedValue,
};
}
match = trimmedExpression.match(/^(contains|startsWith)\(["'](.*)["']\)$/);
if (match) return { variable: '', operator: match[1], value: match[2] };
match = trimmedExpression.match(/^(contains|startsWith)\((.*)\)$/);
if (match) {
const rawValue = match[2].trim();
const unquotedValue = rawValue.replace(/^["'](.*)["']$/, '$1');
return { variable: '', operator: match[1], value: unquotedValue };
}
match = trimmedExpression.match(/^(\w+)\s*(==|!=|>=|<=|>|<)\s*["'](.*)["']$/);
if (match) return { variable: match[1], operator: match[2], value: match[3] };
match = trimmedExpression.match(/^(==|!=|>=|<=|>|<)\s*["'](.*)["']$/);
if (match) return { variable: '', operator: match[1], value: match[2] };
match = trimmedExpression.match(/^(\w+)\s*(==|!=|>=|<=|>|<)\s*(.*)$/);
if (match) return { variable: match[1], operator: match[2], value: match[3] };
match = trimmedExpression.match(/^(==|!=|>=|<=|>|<)\s*(.*)$/);
if (match) return { variable: '', operator: match[1], value: match[2] };
return { variable: '', operator: '==', value: '' };
}
function buildSimpleCel(
variable: string,
operator: string,
value: string,
): string {
const trimmedValue = value.trim();
const isNumeric = trimmedValue !== '' && !isNaN(Number(trimmedValue));
const isBool = trimmedValue === 'true' || trimmedValue === 'false';
const literalValue =
isNumeric || isBool ? trimmedValue : JSON.stringify(value);
const stringValue = JSON.stringify(value);
if (operator === 'contains') {
return variable
? `${variable}.contains(${stringValue})`
: `contains(${stringValue})`;
}
if (operator === 'startsWith') {
return variable
? `${variable}.startsWith(${stringValue})`
: `startsWith(${stringValue})`;
}
if (!variable) return `${operator} ${literalValue}`;
return `${variable} ${operator} ${literalValue}`;
}
function normalizeConditionCases(cases: ConditionCase[]): ConditionCase[] {
const usedHandles = new Set<string>();
let nextIndex = 0;
@@ -294,15 +245,18 @@ function createWorkflowPayload(
| 'agent'
| 'note'
| 'state'
| 'condition',
| 'condition'
| 'code',
title: node.data.title || node.data.label || node.type,
position: node.position,
data:
node.type === 'agent' ||
node.type === 'condition' ||
node.type === 'state'
? node.data.config
: node.data,
node.type === 'code'
? serializeCodeConfig(node.data.config)
: node.type === 'agent' ||
node.type === 'condition' ||
node.type === 'state'
? node.data.config
: node.data,
})),
edges: workflowEdges.map((edge) => ({
id: edge.id,
@@ -321,6 +275,7 @@ const NODE_TYPES: NodeTypes = {
note: NoteNode,
state: SetStateNode,
condition: ConditionNode,
code: CodeNode,
};
function WorkflowBuilderInner() {
@@ -564,6 +519,10 @@ function WorkflowBuilderInner() {
mode: 'simple',
cases: [{ name: '', expression: '', sourceHandle: 'case_0' }],
};
} else if (type === 'code') {
baseNode.data.title = 'Code';
baseNode.data.label = 'Code';
baseNode.data.config = createDefaultCodeConfig();
} else if (type === 'note') {
baseNode.data.title = 'Note';
baseNode.data.label = 'Note';
@@ -688,6 +647,27 @@ function WorkflowBuilderInner() {
[handleUpdateNodeData, selectedNode],
);
const handleCodeJsonSchemaChange = useCallback(
(text: string) => {
if (!selectedNode || selectedNode.type !== 'code') return;
const nodeId = selectedNode.id;
setAgentJsonSchemaDrafts((prev) => ({ ...prev, [nodeId]: text }));
const { schema, error } = parseCodeJsonSchemaDraft(text);
setAgentJsonSchemaErrors((prev) => ({ ...prev, [nodeId]: error }));
if (!error) {
handleUpdateNodeData({
config: {
...(selectedNode.data.config || {}),
json_schema: schema,
},
});
}
},
[handleUpdateNodeData, selectedNode],
);
const handleUpload = useCallback((files: File[]) => {
if (files && files.length > 0) {
setImageFile(files[0]);
@@ -744,9 +724,19 @@ function WorkflowBuilderInner() {
}, [isPublishing]);
useEffect(() => {
// Shared guard for the canvas shortcuts (undo/redo, Delete/Backspace to
// remove the selection, Escape to close the config panel): ignore the
// keystroke while typing in a field, while the Preview Sheet is open (its
// own inputs own the keys — a stray Delete there must not delete the node
// behind it), or once another handler has already consumed the event. Kept
// in one place so the branches can't drift apart.
const shouldIgnoreShortcut = (e: KeyboardEvent): boolean =>
e.defaultPrevented || showPreview || isEditableTarget(e.target);
const handleKeyDown = (e: KeyboardEvent) => {
const isEditable = isEditableTarget(e.target);
if ((e.ctrlKey || e.metaKey) && !isEditable) {
if (shouldIgnoreShortcut(e)) return;
if (e.ctrlKey || e.metaKey) {
const key = e.key.toLowerCase();
if (key === 'z') {
e.preventDefault();
@@ -760,7 +750,7 @@ function WorkflowBuilderInner() {
return;
}
}
if ((e.key === 'Delete' || e.key === 'Backspace') && !isEditable) {
if (e.key === 'Delete' || e.key === 'Backspace') {
e.preventDefault();
deleteNodesAndEdges(
nodes.filter((n) => n.selected || n.id === selectedNode?.id),
@@ -774,7 +764,15 @@ function WorkflowBuilderInner() {
};
window.addEventListener('keydown', handleKeyDown);
return () => window.removeEventListener('keydown', handleKeyDown);
}, [nodes, edges, selectedNode, deleteNodesAndEdges, undo, redo]);
}, [
nodes,
edges,
selectedNode,
deleteNodesAndEdges,
undo,
redo,
showPreview,
]);
const handlePanelBackdropClick = useCallback(() => {
setShowNodeConfig(false);
@@ -876,6 +874,29 @@ function WorkflowBuilderInner() {
});
}, [selectedNode]);
useEffect(() => {
if (!selectedNode || selectedNode.type !== 'code') return;
const nodeId = selectedNode.id;
const rawSchema = selectedNode.data.config?.json_schema;
setAgentJsonSchemaDrafts((prev) => {
if (prev[nodeId] !== undefined) return prev;
if (rawSchema === undefined || rawSchema === null) {
return { ...prev, [nodeId]: '' };
}
try {
return { ...prev, [nodeId]: JSON.stringify(rawSchema, null, 2) };
} catch {
return { ...prev, [nodeId]: String(rawSchema) };
}
});
setAgentJsonSchemaErrors((prev) => {
if (prev[nodeId] !== undefined) return prev;
return { ...prev, [nodeId]: validateCodeJsonSchema(rawSchema) };
});
}, [selectedNode]);
useEffect(() => {
const loadAgentDetails = async () => {
if (!agentId) return;
@@ -929,6 +950,10 @@ function WorkflowBuilderInner() {
};
} else if (n.type === 'state' && n.data) {
nodeData.config = n.data;
} else if (n.type === 'code') {
nodeData.config = normalizeCodeConfig(
n.data as Record<string, unknown> | undefined,
);
} else if (n.data) {
Object.assign(nodeData, n.data);
}
@@ -1185,6 +1210,27 @@ function WorkflowBuilderInner() {
});
});
const codeNodes = nodes.filter((n) => n.type === 'code');
codeNodes.forEach((node) => {
const codeTitle = node.data?.title || node.id;
const config = node.data?.config;
if (!(config?.code || '').trim()) {
errors.push(`Code node "${codeTitle}" must have code to run`);
}
const schemaValidationError = validateCodeJsonSchema(config?.json_schema);
const draftSchemaError = agentJsonSchemaErrors[node.id];
const effectiveSchemaError =
draftSchemaError !== undefined
? draftSchemaError
: schemaValidationError;
if (effectiveSchemaError) {
errors.push(
`Code node "${codeTitle}" JSON schema ${effectiveSchemaError}`,
);
}
});
return errors;
}, [workflowName, nodes, edges, agentJsonSchemaErrors, availableModels]);
@@ -1463,6 +1509,45 @@ function WorkflowBuilderInner() {
return selectedModel.supports_structured_output;
}, [selectedNode, availableModels]);
const selectedAgentDocumentOptions = useMemo(() => {
if (!selectedNode || selectedNode.type !== 'agent') return [];
return toDocumentVariableOptions(
extractUpstreamVariables(nodes, edges, selectedNode.id),
);
}, [selectedNode, nodes, edges]);
const selectedCodeDocumentOptions = useMemo(() => {
if (!selectedNode || selectedNode.type !== 'code') return [];
return toDocumentVariableOptions(
extractUpstreamVariables(nodes, edges, selectedNode.id),
);
}, [selectedNode, nodes, edges]);
const selectedCodeJsonSchemaText = useMemo(() => {
if (!selectedNode || selectedNode.type !== 'code') return '';
const draft = agentJsonSchemaDrafts[selectedNode.id];
if (draft !== undefined) return draft;
const schema = selectedNode.data.config?.json_schema;
if (schema === undefined || schema === null) return '';
try {
return JSON.stringify(schema, null, 2);
} catch {
return String(schema);
}
}, [selectedNode, agentJsonSchemaDrafts]);
const selectedCodeJsonSchemaError = useMemo(() => {
if (!selectedNode || selectedNode.type !== 'code') return null;
const cachedError = agentJsonSchemaErrors[selectedNode.id];
if (cachedError !== undefined) return cachedError;
return validateCodeJsonSchema(selectedNode.data.config?.json_schema);
}, [selectedNode, agentJsonSchemaErrors]);
return (
<>
<MobileBlocker />
@@ -1820,6 +1905,23 @@ function WorkflowBuilderInner() {
</span>
</div>
</div>
<div
className="group border-border bg-card flex cursor-move items-center gap-3 rounded-full border px-4 py-3 shadow-sm transition-all hover:shadow-md"
draggable
onDragStart={(e) => handleNodeDragStart(e, 'code')}
>
<div className="flex h-8 w-8 shrink-0 items-center justify-center rounded-full bg-indigo-100 text-indigo-600 transition-colors group-hover:bg-indigo-600 group-hover:text-white dark:bg-indigo-900/40 dark:text-indigo-300">
<Code2 size={18} />
</div>
<div className="flex flex-col">
<span className="text-foreground text-sm font-medium">
Code
</span>
<span className="text-muted-foreground text-xs">
Run code in a sandbox
</span>
</div>
</div>
</div>
</div>
</div>
@@ -1890,6 +1992,7 @@ function WorkflowBuilderInner() {
{selectedNode.type === 'note' && 'Note'}
{selectedNode.type === 'state' && 'Set global variables'}
{selectedNode.type === 'condition' && 'If / Else'}
{selectedNode.type === 'code' && 'Code'}
</h3>
<Button
type="button"
@@ -2185,6 +2288,60 @@ function WorkflowBuilderInner() {
emptyText="No sources available"
/>
</div>
<NodeDocumentsControl
key={selectedNode.id}
value={
selectedNode.data.config?.input_documents ??
[]
}
onChange={(nextInputDocuments) =>
handleUpdateNodeData({
config: {
...(selectedNode.data.config || {}),
input_documents: nextInputDocuments,
},
})
}
options={selectedAgentDocumentOptions}
label="Documents"
helpText="Documents passed to this agent from uploads or upstream nodes."
/>
<div>
<label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
File passing
</label>
<Select
value={normalizeFilePassing(
selectedNode.data.config?.file_passing,
)}
onValueChange={(value) =>
handleUpdateNodeData({
config: {
...(selectedNode.data.config || {}),
file_passing: value as FilePassing,
},
})
}
>
<SelectTrigger className="w-full">
<SelectValue />
</SelectTrigger>
<SelectContent>
{FILE_PASSING_OPTIONS.map((option) => (
<SelectItem
key={option.value}
value={option.value}
>
{option.label}
</SelectItem>
))}
</SelectContent>
</Select>
<p className="text-muted-foreground mt-1 text-xs">
Auto: send native when the model supports
it, otherwise extract text.
</p>
</div>
<div>
<label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
Structured Output (JSON Schema)
@@ -2717,6 +2874,136 @@ function WorkflowBuilderInner() {
</Button>
</>
)}
{selectedNode.type === 'code' && (
<>
<p className="text-xs text-gray-500 dark:text-gray-400">
Run code in the workflow sandbox. Produced
files are saved as artifacts.
</p>
<div>
<label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
Code
</label>
<textarea
value={selectedNode.data.config?.code ?? ''}
onChange={(e) =>
handleUpdateNodeData({
config: {
...(selectedNode.data.config || {}),
code: e.target.value,
},
})
}
className="border-border focus-visible:ring-ring/50 focus-visible:border-ring bg-card w-full rounded-xl border px-3 py-2 font-mono text-xs transition-all outline-none focus-visible:ring-2 dark:text-white"
rows={10}
spellCheck={false}
placeholder={'print("hello world")'}
/>
</div>
<NodeDocumentsControl
key={selectedNode.id}
value={selectedNode.data.config?.inputs ?? []}
onChange={(nextInputs) =>
handleUpdateNodeData({
config: {
...(selectedNode.data.config || {}),
inputs: nextInputs,
},
})
}
options={selectedCodeDocumentOptions}
label="Input files"
helpText="Artifacts/upstream refs staged as files in the sandbox (one becomes inputs/<name>)."
/>
<div>
<label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
Output Variable
</label>
<Input
type="text"
value={
selectedNode.data.config
?.output_variable || ''
}
onChange={(e) =>
handleUpdateNodeData({
config: {
...(selectedNode.data.config || {}),
output_variable: e.target.value,
},
})
}
className="bg-card h-auto rounded-xl px-3 py-2 text-sm shadow-none"
placeholder="Variable name for output"
/>
</div>
<div>
<label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
Timeout (seconds)
</label>
<Input
type="number"
min={1}
value={
selectedNode.data.config?.timeout ?? ''
}
onChange={(e) => {
const raw = e.target.value;
const parsed =
raw.trim() === ''
? undefined
: Number.parseInt(raw, 10);
handleUpdateNodeData({
config: {
...(selectedNode.data.config || {}),
timeout:
parsed !== undefined &&
Number.isFinite(parsed)
? parsed
: undefined,
},
});
}}
className="bg-card h-auto rounded-xl px-3 py-2 text-sm shadow-none"
placeholder="Optional"
/>
</div>
<div>
<label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
Structured Output (JSON Schema)
</label>
<textarea
value={selectedCodeJsonSchemaText}
onChange={(e) =>
handleCodeJsonSchemaChange(e.target.value)
}
className="border-border focus-visible:ring-ring/50 focus-visible:border-ring bg-card w-full rounded-xl border px-3 py-2 font-mono text-xs transition-all outline-none focus-visible:ring-2 dark:text-white"
rows={6}
placeholder={`{
"type": "object",
"properties": {
"result": { "type": "string" }
},
"required": ["result"]
}`}
/>
{selectedCodeJsonSchemaText.trim() !== '' && (
<p
className={`mt-2 text-xs ${
selectedCodeJsonSchemaError
? 'text-red-600 dark:text-red-400'
: 'text-green-600 dark:text-green-400'
}`}
>
{selectedCodeJsonSchemaError
? `Invalid JSON schema: ${selectedCodeJsonSchemaError}`
: 'Valid JSON schema'}
</p>
)}
</div>
</>
)}
</>
)}
</div>
@@ -2743,10 +3030,11 @@ function WorkflowBuilderInner() {
<Sheet open={showPreview} onOpenChange={setShowPreview}>
<SheetContent
side="right"
showCloseButton={false}
title="Workflow preview"
className="bg-card w-full max-w-none p-0 sm:max-w-[600px] md:max-w-[700px] lg:max-w-[800px]"
>
<WorkflowPreview
workflowId={workflowId}
workflowData={{
name: workflowName,
description: workflowDescription,
@@ -2754,10 +3042,20 @@ function WorkflowBuilderInner() {
.filter((n) => n.type !== 'note')
.map((n) => ({
id: n.id,
type: n.type as 'start' | 'end' | 'agent' | 'state',
type: n.type as
| 'start'
| 'end'
| 'agent'
| 'state'
| 'code',
title: n.data.title || n.data.label || n.type,
position: n.position,
data: n.type === 'agent' ? n.data.config : n.data,
data:
n.type === 'code'
? serializeCodeConfig(n.data.config)
: n.type === 'agent'
? n.data.config
: n.data,
})),
edges: edges.map((e) => ({
id: e.id,
@@ -2,7 +2,9 @@ import {
Bot,
CheckCircle2,
Circle,
Code2,
Database,
FileBox,
Flag,
GitBranch,
Loader2,
@@ -23,14 +25,18 @@ import MessageInput from '../../components/MessageInput';
import ConversationBubble from '../../conversation/ConversationBubble';
import { Query } from '../../conversation/conversationModels';
import { AppDispatch } from '../../store';
import { selectCompletedAttachments } from '../../upload/uploadSlice';
import { WorkflowEdge, WorkflowNode } from '../types/workflow';
import WorkflowRunArtifacts from './WorkflowRunArtifacts';
import {
addQuery,
fetchWorkflowPreviewAnswer,
handleWorkflowPreviewAbort,
previewSendBlockReason,
resendQuery,
resetWorkflowPreview,
selectActiveNodeId,
setPreviewOpen,
selectWorkflowExecutionSteps,
selectWorkflowPreviewQueries,
selectWorkflowPreviewStatus,
@@ -47,6 +53,9 @@ interface WorkflowData {
interface WorkflowPreviewProps {
workflowData: WorkflowData;
// Saved workflow id (when the draft has been persisted); enables run-artifact
// listing by persisting a ``workflow_runs`` row for the preview run.
workflowId?: string | null;
}
const NODE_ICONS: Record<string, React.ReactNode> = {
@@ -56,6 +65,7 @@ const NODE_ICONS: Record<string, React.ReactNode> = {
note: <StickyNote className="h-3 w-3" />,
state: <Database className="h-3 w-3" />,
condition: <GitBranch className="h-3 w-3" />,
code: <Code2 className="h-3 w-3" />,
};
const NODE_COLORS: Record<string, string> = {
@@ -65,6 +75,7 @@ const NODE_COLORS: Record<string, string> = {
note: 'text-yellow-600 dark:text-yellow-400',
state: 'text-blue-600 dark:text-blue-400',
condition: 'text-orange-600 dark:text-orange-400',
code: 'text-indigo-600 dark:text-indigo-400',
};
function ExecutionDetails({
@@ -134,8 +145,8 @@ function ExecutionDetails({
const node = nodes.find((n) => n.id === step.nodeId);
const displayName =
node?.title || node?.data?.title || step.nodeTitle;
const stateVars = step.stateSnapshot
? Object.entries(step.stateSnapshot).filter(
const stateVars = step.stateDelta
? Object.entries(step.stateDelta).filter(
([key]) => !['query', 'chat_history'].includes(key),
)
: [];
@@ -238,6 +249,61 @@ function ExecutionDetails({
);
}
function RunArtifactsSection({
workflowRunId,
isOpen,
onToggle,
runInProgress = false,
}: {
workflowRunId: string;
isOpen: boolean;
onToggle: () => void;
runInProgress?: boolean;
}) {
return (
<div className="mb-4 flex w-full flex-col flex-wrap items-start self-start lg:flex-nowrap">
<div className="my-2 flex flex-row items-center justify-center gap-3">
<div className="flex h-[26px] w-[30px] items-center justify-center">
<FileBox className="h-5 w-5 text-gray-600 dark:text-gray-400" />
</div>
<Button
type="button"
variant="ghost"
onClick={onToggle}
className="h-auto gap-2 px-0 py-0 hover:bg-transparent"
>
<p className="text-base font-semibold">Artifacts</p>
<img
src={ChevronDownIcon}
alt="ChevronDown"
className={cn(
'h-4 w-4 transform transition-transform duration-200 dark:invert',
isOpen ? 'rotate-180' : '',
)}
/>
</Button>
</div>
<div
className={cn(
'ml-3 grid w-full transition-all duration-300 ease-in-out',
isOpen ? 'grid-rows-[1fr] opacity-100' : 'grid-rows-[0fr] opacity-0',
)}
>
<div className="overflow-hidden">
<div className="max-h-[480px] overflow-y-auto pr-2">
{isOpen && (
<WorkflowRunArtifacts
workflowRunId={workflowRunId}
inProgress={runInProgress}
/>
)}
</div>
</div>
</div>
</div>
);
}
function WorkflowMiniMap({
nodes,
activeNodeId,
@@ -373,6 +439,7 @@ function WorkflowMiniMap({
export default function WorkflowPreview({
workflowData,
workflowId,
}: WorkflowPreviewProps) {
const dispatch = useDispatch<AppDispatch>();
@@ -380,9 +447,17 @@ export default function WorkflowPreview({
const status = useSelector(selectWorkflowPreviewStatus);
const executionSteps = useSelector(selectWorkflowExecutionSteps);
const activeNodeId = useSelector(selectActiveNodeId);
const completedAttachments = useSelector(selectCompletedAttachments);
const hasCompletedAttachment = completedAttachments.length > 0;
const [lastQueryReturnedErr, setLastQueryReturnedErr] = useState(false);
const [sendBlockedMessage, setSendBlockedMessage] = useState<string | null>(
null,
);
const [openDetailsIndex, setOpenDetailsIndex] = useState<number | null>(null);
const [openArtifactsIndex, setOpenArtifactsIndex] = useState<number | null>(
null,
);
const fetchStream = useRef<{ abort: () => void } | null>(null);
const stepRefs = useRef<Map<string, HTMLDivElement>>(new Map());
@@ -411,11 +486,12 @@ export default function WorkflowPreview({
question,
workflowData,
indx: index,
workflowId,
}),
);
fetchStream.current = promise;
},
[dispatch, workflowData],
[dispatch, workflowData, workflowId],
);
const handleQuestion = useCallback(
@@ -428,8 +504,24 @@ export default function WorkflowPreview({
isRetry?: boolean;
index?: number;
}) => {
// An unsaved draft can't bridge uploaded documents into the run (no
// persisted workflow_id), so block the send rather than run with the docs
// silently dropped. The uploads are kept (not cleared) for a retry after
// the workflow is saved.
const blockReason = previewSendBlockReason(
workflowId,
hasCompletedAttachment,
);
if (blockReason) {
setSendBlockedMessage(blockReason);
return;
}
setSendBlockedMessage(null);
const trimmedQuestion = question.trim();
if (trimmedQuestion === '') return;
// Doc-driven nodes read ``input_documents`` rather than the query, so an
// attachment-only run is allowed to proceed with an empty question.
if (trimmedQuestion === '' && !hasCompletedAttachment) return;
if (index !== undefined) {
if (!isRetry) dispatch(resendQuery({ index, prompt: trimmedQuestion }));
@@ -442,9 +534,15 @@ export default function WorkflowPreview({
handleFetchAnswer({ question: trimmedQuestion, index: undefined });
}
},
[dispatch, handleFetchAnswer],
[dispatch, handleFetchAnswer, hasCompletedAttachment, workflowId],
);
// Clear the block message once it no longer applies (the workflow was saved,
// or the attachments were removed) so a stale warning doesn't linger.
useEffect(() => {
if (workflowId || !hasCompletedAttachment) setSendBlockedMessage(null);
}, [workflowId, hasCompletedAttachment]);
const handleQuestionSubmission = (
question?: string,
updated?: boolean,
@@ -456,8 +554,8 @@ export default function WorkflowPreview({
index: indx,
isRetry: false,
});
} else if (question && status !== 'loading') {
const currentInput = question.trim();
} else if ((question || hasCompletedAttachment) && status !== 'loading') {
const currentInput = (question ?? '').trim();
if (lastQueryReturnedErr && queries.length > 0) {
const lastQueryIndex = queries.length - 1;
handleQuestion({
@@ -477,10 +575,12 @@ export default function WorkflowPreview({
useEffect(() => {
dispatch(resetWorkflowPreview());
dispatch(setPreviewOpen(true));
return () => {
if (fetchStream.current) fetchStream.current.abort();
handleWorkflowPreviewAbort();
dispatch(resetWorkflowPreview());
dispatch(setPreviewOpen(false));
};
}, [dispatch]);
@@ -594,6 +694,22 @@ export default function WorkflowPreview({
/>
)}
{/* Run artifacts — shown as soon as a persisted run id
is known; mid-run it lists artifacts as nodes produce
them and refetches when the run completes. */}
{query.workflowRunId && (
<RunArtifactsSection
workflowRunId={query.workflowRunId}
isOpen={openArtifactsIndex === index}
onToggle={() =>
setOpenArtifactsIndex(
openArtifactsIndex === index ? null : index,
)
}
runInProgress={isStreamingLastQuery}
/>
)}
{/* Response bubble */}
{(query.response ||
shouldShowThought ||
@@ -627,12 +743,18 @@ export default function WorkflowPreview({
)}
</div>
<div className="bg-card absolute right-0 bottom-0 left-0 flex w-full flex-col gap-2 px-4 pt-2 pb-4">
{sendBlockedMessage && (
<p className="text-xs text-red-500" role="alert">
{sendBlockedMessage}
</p>
)}
<MessageInput
onSubmit={(text) => handleQuestionSubmission(text)}
loading={status === 'loading'}
showSourceButton={false}
showToolButton={false}
autoFocus={true}
allowSendWithoutText={true}
/>
</div>
</div>
@@ -0,0 +1,223 @@
import { ChevronLeft, FileBox } from 'lucide-react';
import { useCallback, useEffect, useState } from 'react';
import { useSelector } from 'react-redux';
import userService from '../../api/services/userService';
import DocumentArtifactView from '../../components/DocumentArtifactView';
import Spinner from '../../components/Spinner';
import {
isDocumentArtifact,
type DocumentArtifact,
} from '../../components/artifactViewUtils';
import { Button } from '@/components/ui/button';
import { selectToken } from '../../preferences/preferenceSlice';
interface RunArtifactSummary {
id: string;
kind: string | null;
title: string | null;
current_version: number | null;
}
interface WorkflowRunArtifactsProps {
workflowRunId: string;
/** True while the run is still streaming; flipping to false refetches so
* a panel opened mid-run picks up the final artifact list. */
inProgress?: boolean;
}
/** List a workflow run's produced artifacts, with click-through preview + download. */
export default function WorkflowRunArtifacts({
workflowRunId,
inProgress = false,
}: WorkflowRunArtifactsProps) {
const token = useSelector(selectToken);
const [artifacts, setArtifacts] = useState<RunArtifactSummary[] | null>(null);
const [loading, setLoading] = useState(true);
const [error, setError] = useState<string | null>(null);
const [selectedId, setSelectedId] = useState<string | null>(null);
const [detail, setDetail] = useState<DocumentArtifact | null>(null);
const [detailLoading, setDetailLoading] = useState(false);
const [detailError, setDetailError] = useState<string | null>(null);
const loadList = useCallback(() => {
if (!workflowRunId) return;
let cancelled = false;
setLoading(true);
setError(null);
userService
.listWorkflowRunArtifacts(workflowRunId, token)
.then(async (res: Response) => {
if (cancelled) return;
if (!res.ok) {
// A 403 is the expected unsaved-draft / unauthorized case (no
// persisted run row): show the informational empty state. Any other
// non-OK status (500, an expired 401, ...) is a real failure — surface
// an error + Retry instead of masquerading as "no artifacts".
if (res.status === 403) {
setArtifacts([]);
} else {
setError('Failed to load artifacts');
}
setLoading(false);
return;
}
const data = await res.json().catch(() => null);
if (cancelled) return;
setArtifacts(data?.success ? (data.artifacts ?? []) : []);
setLoading(false);
})
.catch(() => {
if (cancelled) return;
setError('Failed to load artifacts');
setLoading(false);
});
return () => {
cancelled = true;
};
}, [workflowRunId, token]);
useEffect(() => {
const cleanup = loadList();
return cleanup;
}, [loadList, inProgress]);
const fetchDetail = useCallback(
(artifactId: string) => {
let cancelled = false;
setDetailLoading(true);
setDetailError(null);
userService
.getDocumentArtifact(artifactId, token)
.then(async (res: Response) => {
if (cancelled) return;
if (!res.ok) {
setDetailError('Failed to load artifact');
setDetailLoading(false);
return;
}
const data = await res.json().catch(() => null);
if (cancelled) return;
if (data?.success && isDocumentArtifact(data.artifact)) {
setDetail(data.artifact);
setDetailLoading(false);
} else {
setDetailError('This artifact cannot be previewed');
setDetailLoading(false);
}
})
.catch(() => {
if (cancelled) return;
setDetailError('Failed to load artifact');
setDetailLoading(false);
});
return () => {
cancelled = true;
};
},
[token],
);
useEffect(() => {
if (!selectedId) {
setDetail(null);
setDetailError(null);
return;
}
const cleanup = fetchDetail(selectedId);
return cleanup;
}, [selectedId, fetchDetail]);
if (loading) {
return (
<div className="flex items-center gap-2 px-3 py-4 text-sm text-gray-500 dark:text-gray-400">
<Spinner size="small" /> Loading artifacts...
</div>
);
}
if (error) {
return (
<div className="flex items-center justify-between gap-2 px-3 py-3 text-sm text-red-500">
<span>{error}</span>
<Button type="button" variant="outline" size="sm" onClick={loadList}>
Retry
</Button>
</div>
);
}
if (!artifacts || artifacts.length === 0) {
return (
<div className="px-3 py-3 text-sm text-gray-500 dark:text-gray-400">
{inProgress
? 'Run in progress — artifacts appear here as nodes produce them.'
: 'No artifacts produced by this run.'}
</div>
);
}
if (selectedId) {
return (
<div className="flex h-full min-h-0 flex-col">
<div className="mb-2 flex items-center">
<Button
type="button"
variant="ghost"
size="sm"
className="gap-1 px-2"
onClick={() => setSelectedId(null)}
>
<ChevronLeft className="h-4 w-4" />
Back to artifacts
</Button>
</div>
<div className="min-h-0 flex-1 overflow-hidden">
{detailLoading ? (
<div className="flex h-full items-center justify-center">
<Spinner />
</div>
) : detailError ? (
<div className="flex h-full items-center justify-center">
<p className="text-sm text-red-500">{detailError}</p>
</div>
) : detail ? (
<DocumentArtifactView
artifact={detail}
onRefresh={() => fetchDetail(selectedId)}
/>
) : null}
</div>
</div>
);
}
return (
<ul className="space-y-2">
{artifacts.map((artifact) => (
<li key={artifact.id}>
<Button
type="button"
variant="outline"
onClick={() => setSelectedId(artifact.id)}
className="h-auto w-full justify-start gap-3 px-3 py-2 text-left"
>
<FileBox className="h-4 w-4 shrink-0 text-gray-500 dark:text-gray-400" />
<div className="min-w-0 flex-1">
<div className="truncate text-sm font-medium text-gray-900 dark:text-white">
{artifact.title || `Artifact ${artifact.id.slice(0, 8)}`}
</div>
<div className="truncate text-xs text-gray-500 dark:text-gray-400">
{artifact.kind || 'file'}
{artifact.current_version != null
? ` · v${artifact.current_version}`
: ''}
</div>
</div>
</Button>
</li>
))}
</ul>
);
}
@@ -0,0 +1,163 @@
import { describe, expect, it } from 'vitest';
import {
createDefaultCodeConfig,
normalizeCodeConfig,
parseCodeInputs,
parseCodeJsonSchemaDraft,
serializeCodeConfig,
stringifyCodeInputs,
validateCodeJsonSchema,
} from './codeNodeConfig';
describe('validateCodeJsonSchema', () => {
it('treats undefined/null as valid (no schema)', () => {
expect(validateCodeJsonSchema(undefined)).toBeNull();
expect(validateCodeJsonSchema(null)).toBeNull();
});
it('rejects arrays and non-objects', () => {
expect(validateCodeJsonSchema([])).toMatch(/valid JSON object/);
expect(validateCodeJsonSchema('x')).toMatch(/valid JSON object/);
});
it('requires a type or schema field', () => {
expect(validateCodeJsonSchema({})).toMatch(/"type" or "schema"/);
expect(validateCodeJsonSchema({ type: 'object' })).toBeNull();
expect(validateCodeJsonSchema({ schema: {} })).toBeNull();
});
});
describe('parseCodeJsonSchemaDraft', () => {
it('returns undefined schema and no error for blank input', () => {
expect(parseCodeJsonSchemaDraft(' ')).toEqual({
schema: undefined,
error: null,
});
});
it('flags invalid JSON', () => {
const result = parseCodeJsonSchemaDraft('{ not json');
expect(result.schema).toBeUndefined();
expect(result.error).toMatch(/valid JSON/);
});
it('flags structurally invalid schema', () => {
const result = parseCodeJsonSchemaDraft('{"foo": 1}');
expect(result.schema).toBeUndefined();
expect(result.error).toMatch(/"type" or "schema"/);
});
it('parses a valid schema', () => {
const result = parseCodeJsonSchemaDraft(
'{"type":"object","properties":{}}',
);
expect(result.error).toBeNull();
expect(result.schema).toEqual({ type: 'object', properties: {} });
});
});
describe('inputs (de)serialization', () => {
it('splits on newlines and commas and trims', () => {
expect(parseCodeInputs('A1\nfoo, bar\n\n baz ')).toEqual([
'A1',
'foo',
'bar',
'baz',
]);
});
it('round-trips through stringify', () => {
const inputs = ['A1', 'state_var'];
expect(parseCodeInputs(stringifyCodeInputs(inputs))).toEqual(inputs);
});
it('stringifies undefined to empty string', () => {
expect(stringifyCodeInputs(undefined)).toBe('');
});
});
describe('createDefaultCodeConfig', () => {
it('returns python defaults with empty code and inputs', () => {
expect(createDefaultCodeConfig()).toEqual({
code: '',
inputs: [],
});
});
});
describe('serializeCodeConfig', () => {
it('omits optional fields when empty', () => {
const serialized = serializeCodeConfig({
code: 'print(1)',
inputs: ['', ' '],
output_variable: ' ',
});
expect(serialized).toEqual({
code: 'print(1)',
inputs: [],
});
expect('output_variable' in serialized).toBe(false);
expect('timeout' in serialized).toBe(false);
expect('json_schema' in serialized).toBe(false);
});
it('keeps all CodeNodeConfig fields when present', () => {
const serialized = serializeCodeConfig({
code: 'print(1)',
inputs: ['A1', 'state_var'],
output_variable: 'result',
timeout: 30,
json_schema: { type: 'object' },
});
expect(serialized).toEqual({
code: 'print(1)',
inputs: ['A1', 'state_var'],
output_variable: 'result',
timeout: 30,
json_schema: { type: 'object' },
});
});
it('falls back to python defaults for an undefined config', () => {
expect(serializeCodeConfig(undefined)).toEqual({
code: '',
inputs: [],
});
});
});
describe('normalizeCodeConfig', () => {
it('tolerates a missing/empty payload', () => {
expect(normalizeCodeConfig(undefined)).toEqual({
code: '',
inputs: [],
});
});
it('coerces bad field types without throwing', () => {
const config = normalizeCodeConfig({
code: 42,
inputs: ['A1', 5, 'B2'],
output_variable: '',
timeout: 'soon',
json_schema: [],
} as unknown as Record<string, unknown>);
expect(config).toEqual({
code: '',
inputs: ['A1', 'B2'],
});
});
it('preserves a fully populated saved config (round-trip with serialize)', () => {
const saved = {
code: 'print(1)',
inputs: ['A1'],
output_variable: 'result',
timeout: 30,
json_schema: { type: 'object' },
};
const normalized = normalizeCodeConfig(saved);
expect(serializeCodeConfig(normalized)).toEqual(saved);
});
});
@@ -0,0 +1,125 @@
import type { CodeNodeConfig } from '../types/workflow';
/**
* Validate a parsed JSON schema object the same way the agent node does:
* it must be a non-array object that declares a "type" or "schema" field.
* Returns null when valid (or absent), otherwise a short error fragment.
*/
export function validateCodeJsonSchema(schema: unknown): string | null {
if (schema === undefined || schema === null) return null;
if (typeof schema !== 'object' || Array.isArray(schema)) {
return 'must be a valid JSON object';
}
const schemaObject = schema as Record<string, unknown>;
if (!('schema' in schemaObject) && !('type' in schemaObject)) {
return 'must include either a "type" or "schema" field';
}
return null;
}
/**
* Parse the raw json_schema textarea content. Returns the parsed schema (or
* undefined when blank) and a validation error message when invalid.
*/
export function parseCodeJsonSchemaDraft(text: string): {
schema: Record<string, unknown> | undefined;
error: string | null;
} {
if (text.trim() === '') {
return { schema: undefined, error: null };
}
let parsed: unknown;
try {
parsed = JSON.parse(text);
} catch {
return { schema: undefined, error: 'must be valid JSON' };
}
const error = validateCodeJsonSchema(parsed);
if (error) {
return { schema: undefined, error };
}
return { schema: parsed as Record<string, unknown>, error: null };
}
/** Split a newline/comma separated inputs textarea into a clean ref list. */
export function parseCodeInputs(text: string): string[] {
return text
.split(/[\n,]/)
.map((entry) => entry.trim())
.filter((entry) => entry.length > 0);
}
/** Render an inputs list back into the textarea (one ref per line). */
export function stringifyCodeInputs(inputs: string[] | undefined): string {
return (inputs || []).join('\n');
}
/** Default config applied when a fresh code node is dropped onto the canvas. */
export function createDefaultCodeConfig(): CodeNodeConfig {
return {
code: '',
inputs: [],
};
}
/**
* Normalize an arbitrary loaded config (from a saved workflow) into the editor
* shape, tolerating missing/extra fields so existing workflows keep loading.
*/
export function normalizeCodeConfig(
raw: Record<string, unknown> | undefined,
): CodeNodeConfig {
const source = raw || {};
const config: CodeNodeConfig = {
code: typeof source.code === 'string' ? source.code : '',
inputs: Array.isArray(source.inputs)
? source.inputs.filter(
(entry): entry is string => typeof entry === 'string',
)
: [],
};
if (typeof source.output_variable === 'string' && source.output_variable) {
config.output_variable = source.output_variable;
}
if (typeof source.timeout === 'number' && Number.isFinite(source.timeout)) {
config.timeout = source.timeout;
}
if (
source.json_schema &&
typeof source.json_schema === 'object' &&
!Array.isArray(source.json_schema)
) {
config.json_schema = source.json_schema as Record<string, unknown>;
}
return config;
}
/**
* Serialize the editor config into the exact CodeNodeConfig shape the backend
* expects. Optional fields are omitted when empty so the payload stays clean.
*/
export function serializeCodeConfig(
config: CodeNodeConfig | undefined,
): CodeNodeConfig {
const source = config || createDefaultCodeConfig();
const serialized: CodeNodeConfig = {
code: source.code || '',
inputs: (source.inputs || []).filter(
(entry) => typeof entry === 'string' && entry.trim() !== '',
),
};
if (source.output_variable && source.output_variable.trim() !== '') {
serialized.output_variable = source.output_variable.trim();
}
if (typeof source.timeout === 'number' && Number.isFinite(source.timeout)) {
serialized.timeout = source.timeout;
}
if (
source.json_schema &&
typeof source.json_schema === 'object' &&
!Array.isArray(source.json_schema)
) {
serialized.json_schema = source.json_schema;
}
return serialized;
}
@@ -0,0 +1,128 @@
import { Plus } from 'lucide-react';
import { useState } from 'react';
import { Button } from '@/components/ui/button';
import { Input } from '@/components/ui/input';
import { MultiSelect } from '@/components/ui/multi-select';
import {
appendDocumentRef,
documentsModeToInputDocuments,
DocumentsMode,
getDocumentsMode,
withChosenDocumentOptions,
} from '../documentConfig';
interface NodeDocumentsControlProps {
value: string[];
onChange: (next: string[]) => void;
options: { value: string; label: string }[];
label: string;
helpText?: string;
}
const MODE_OPTIONS: { mode: DocumentsMode; label: string }[] = [
{ mode: 'all', label: 'All input docs' },
{ mode: 'none', label: 'None' },
{ mode: 'choose', label: 'Choose…' },
];
/** Shared All/None/Choose documents picker for agent and code workflow nodes. */
export default function NodeDocumentsControl({
value,
onChange,
options,
label,
helpText,
}: NodeDocumentsControlProps) {
// Track mode in component state so "Choose" stays reachable even when the
// chosen list is empty (an empty list otherwise reads back as "None").
const [mode, setMode] = useState<DocumentsMode>(() =>
getDocumentsMode(value),
);
const [refDraft, setRefDraft] = useState('');
const chosen = documentsModeToInputDocuments('choose', value);
const showChoose = mode === 'choose';
const selectMode = (next: DocumentsMode) => {
setMode(next);
onChange(documentsModeToInputDocuments(next, chosen));
};
const addRef = () => {
const next = appendDocumentRef(chosen, refDraft);
setRefDraft('');
if (next === chosen) return;
setMode('choose');
onChange(documentsModeToInputDocuments('choose', next));
};
return (
<div>
<label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
{label}
</label>
<div className="border-border bg-card flex gap-1 rounded-xl border p-1">
{MODE_OPTIONS.map(({ mode: optionMode, label: modeLabel }) => (
<Button
key={optionMode}
type="button"
variant="ghost"
size="sm"
onClick={() => selectMode(optionMode)}
className={`h-auto flex-1 rounded-lg px-2 py-1.5 text-xs font-medium ${
mode === optionMode
? 'bg-primary text-white'
: 'text-gray-600 dark:text-gray-300'
}`}
>
{modeLabel}
</Button>
))}
</div>
{showChoose && (
<div className="mt-2 flex flex-col gap-2">
<MultiSelect
options={withChosenDocumentOptions(options, chosen)}
selected={chosen}
onChange={(next) =>
onChange(documentsModeToInputDocuments('choose', next))
}
placeholder="Select documents..."
searchPlaceholder="Search variables..."
emptyText="No upstream documents"
/>
<div className="flex gap-2">
<Input
type="text"
value={refDraft}
onChange={(e) => setRefDraft(e.target.value)}
onKeyDown={(e) => {
if (e.key === 'Enter') {
e.preventDefault();
addRef();
}
}}
className="bg-card h-auto rounded-xl px-3 py-2 text-sm shadow-none"
placeholder="Add ref (e.g. A1)"
/>
<Button
type="button"
variant="ghost"
size="sm"
onClick={addRef}
className="h-auto shrink-0 gap-1 rounded-lg px-3 py-2 text-sm font-medium text-gray-600 dark:text-gray-300"
>
<Plus size={14} />
Add
</Button>
</div>
</div>
)}
{helpText && (
<p className="text-muted-foreground mt-1 text-xs">{helpText}</p>
)}
</div>
);
}
@@ -10,10 +10,16 @@ import {
PopoverTrigger,
} from '@/components/ui/popover';
interface WorkflowVariable {
export interface WorkflowVariable {
label: string;
templatePath: string;
section: string;
// True when the variable resolves to artifact reference(s) at run time
// (uploaded input_documents, code-node outputs) rather than plain LLM/state
// TEXT. Only artifact-bearing variables may be picked as node Documents;
// selecting a text output there makes the engine append the literal variable
// name and the node hard-fails. Heuristic keyed on the producing node type.
producesArtifact?: boolean;
}
const GLOBAL_CONTEXT_VARIABLES: WorkflowVariable[] = [
@@ -62,6 +68,11 @@ const GLOBAL_CONTEXT_VARIABLES: WorkflowVariable[] = [
templatePath: 'system.user_id',
section: 'Global context',
},
{
label: 'artifacts.artifact(id)',
templatePath: 'artifacts.artifact(id)',
section: 'Global context',
},
];
function toAgentTemplatePath(variableName: string): string {
@@ -93,7 +104,7 @@ function getUpstreamNodeIds(nodeId: string, edges: Edge[]): Set<string> {
return upstream;
}
function extractUpstreamVariables(
export function extractUpstreamVariables(
nodes: Node[],
edges: Edge[],
selectedNodeId: string,
@@ -109,40 +120,58 @@ function extractUpstreamVariables(
templatePath: 'agent.chat_history',
section: 'Workflow input',
},
{
label: 'agent.input_documents',
templatePath: 'agent.input_documents',
section: 'Workflow input',
// Uploaded documents are artifact references.
producesArtifact: true,
},
...GLOBAL_CONTEXT_VARIABLES,
];
const seen = new Set(variables.map((variable) => variable.templatePath));
const upstreamIds = getUpstreamNodeIds(selectedNodeId, edges);
const pushNodeOutput = (
node: Node,
outputName: string,
sectionFallback: string,
producesArtifact: boolean,
) => {
const templatePath = toAgentTemplatePath(outputName);
if (seen.has(templatePath)) return;
seen.add(templatePath);
variables.push({
label: templatePath,
templatePath,
section: node.data.title || node.data.label || sectionFallback,
producesArtifact,
});
};
for (const node of nodes) {
if (!upstreamIds.has(node.id)) continue;
if (node.type === 'agent') {
const defaultOutputTemplatePath = toAgentTemplatePath(
if (node.type === 'agent' || node.type === 'code') {
// Agent and code nodes both expose `node_<id>_output` and an optional
// `output_variable`, but only code-node outputs resolve to artifact
// references in the engine — agent outputs are LLM TEXT. So both feed the
// prompt-variable popover, while only code outputs are offered as
// Documents (see toDocumentVariableOptions / producesArtifact).
const producesArtifact = node.type === 'code';
const sectionFallback = node.type === 'code' ? 'Code' : 'Agent';
pushNodeOutput(
node,
`node_${node.id}_output`,
sectionFallback,
producesArtifact,
);
if (!seen.has(defaultOutputTemplatePath)) {
seen.add(defaultOutputTemplatePath);
variables.push({
label: defaultOutputTemplatePath,
templatePath: defaultOutputTemplatePath,
section: node.data.title || node.data.label || 'Agent',
});
}
const outputVariable = String(
node.data?.config?.output_variable || '',
).trim();
if (outputVariable) {
const templatePath = toAgentTemplatePath(outputVariable);
if (!seen.has(templatePath)) {
seen.add(templatePath);
variables.push({
label: templatePath,
templatePath,
section: node.data.title || node.data.label || 'Agent',
});
}
pushNodeOutput(node, outputVariable, sectionFallback, producesArtifact);
}
}
@@ -0,0 +1,189 @@
import { describe, expect, it } from 'vitest';
import type { WorkflowVariable } from './components/PromptTextArea';
import {
ALL_INPUT_DOCUMENTS_TOKEN,
appendDocumentRef,
DEFAULT_FILE_PASSING,
documentsModeToInputDocuments,
getDocumentsMode,
normalizeFilePassing,
stripAgentPrefix,
toDocumentVariableOptions,
withChosenDocumentOptions,
} from './documentConfig';
describe('getDocumentsMode', () => {
it('treats missing/empty as none', () => {
expect(getDocumentsMode(undefined)).toBe('none');
expect(getDocumentsMode([])).toBe('none');
});
it('treats the wildcard token as all', () => {
expect(getDocumentsMode([ALL_INPUT_DOCUMENTS_TOKEN])).toBe('all');
});
it('treats a bare-name list as choose', () => {
expect(getDocumentsMode(['wire_doc'])).toBe('choose');
expect(getDocumentsMode(['wire_doc', 'A1'])).toBe('choose');
});
});
describe('documentsModeToInputDocuments', () => {
it('maps all to the wildcard token regardless of chosen names', () => {
expect(documentsModeToInputDocuments('all', ['wire_doc'])).toEqual([
ALL_INPUT_DOCUMENTS_TOKEN,
]);
});
it('maps none to an empty list', () => {
expect(documentsModeToInputDocuments('none', ['wire_doc'])).toEqual([]);
});
it('maps choose to the chosen bare names, dropping blanks and the wildcard', () => {
expect(
documentsModeToInputDocuments('choose', [
'wire_doc',
' ',
ALL_INPUT_DOCUMENTS_TOKEN,
'A1',
]),
).toEqual(['wire_doc', 'A1']);
});
it('round-trips a chosen selection through getDocumentsMode', () => {
const next = documentsModeToInputDocuments('choose', ['wire_doc']);
expect(getDocumentsMode(next)).toBe('choose');
expect(next).toEqual(['wire_doc']);
});
});
describe('normalizeFilePassing', () => {
it('passes through supported values', () => {
expect(normalizeFilePassing('auto')).toBe('auto');
expect(normalizeFilePassing('native')).toBe('native');
expect(normalizeFilePassing('extract')).toBe('extract');
});
it('falls back to the default for unknown values', () => {
expect(normalizeFilePassing(undefined)).toBe(DEFAULT_FILE_PASSING);
expect(normalizeFilePassing('bogus')).toBe(DEFAULT_FILE_PASSING);
});
});
describe('stripAgentPrefix', () => {
it('strips a dotted agent prefix to a bare name', () => {
expect(stripAgentPrefix('agent.wire_doc')).toBe('wire_doc');
expect(stripAgentPrefix('agent.node_abc_output')).toBe('node_abc_output');
});
it('strips a bracketed agent prefix and unescapes quotes', () => {
expect(stripAgentPrefix("agent['weird name']")).toBe('weird name');
expect(stripAgentPrefix("agent['it\\'s']")).toBe("it's");
});
it('leaves bare names untouched', () => {
expect(stripAgentPrefix('A1')).toBe('A1');
});
});
describe('toDocumentVariableOptions', () => {
const make = (
templatePath: string,
producesArtifact = false,
): WorkflowVariable => ({
label: templatePath,
templatePath,
section: 'x',
producesArtifact,
});
it('offers only artifact-producing variables (input_documents, code outputs)', () => {
const options = toDocumentVariableOptions([
make('agent.input_documents', true),
make('agent.node_code1_output', true),
make("agent['weird name']", true),
]);
expect(options).toEqual([
{ value: 'input_documents', label: 'input_documents' },
{ value: 'node_code1_output', label: 'node_code1_output' },
{ value: 'weird name', label: 'weird name' },
]);
});
it('excludes plain agent/state TEXT outputs that are not artifact refs', () => {
const options = toDocumentVariableOptions([
make('agent.node_agent1_output'),
make('agent.some_state_var'),
make('agent.query'),
make('agent.chat_history'),
make('source.content'),
make('system.date'),
make('artifacts.artifact(id)'),
]);
expect(options).toEqual([]);
});
it('deduplicates by bare name', () => {
const options = toDocumentVariableOptions([
make('agent.wire_doc', true),
make('agent.wire_doc', true),
]);
expect(options).toEqual([{ value: 'wire_doc', label: 'wire_doc' }]);
});
});
describe('appendDocumentRef', () => {
it('appends a trimmed literal ref', () => {
expect(appendDocumentRef(['wire_doc'], ' A1 ')).toEqual([
'wire_doc',
'A1',
]);
});
it('ignores blanks and the wildcard token', () => {
expect(appendDocumentRef(['wire_doc'], ' ')).toEqual(['wire_doc']);
expect(appendDocumentRef(['wire_doc'], ALL_INPUT_DOCUMENTS_TOKEN)).toEqual([
'wire_doc',
]);
});
it('does not append an already-present ref', () => {
expect(appendDocumentRef(['A1'], 'A1')).toEqual(['A1']);
});
});
describe('withChosenDocumentOptions', () => {
const base = [{ value: 'wire_doc', label: 'wire_doc' }];
it('keeps known options unchanged when all chosen are present', () => {
expect(withChosenDocumentOptions(base, ['wire_doc'])).toEqual(base);
});
it('appends an orphaned chosen name (renamed/deleted upstream var) as removable', () => {
expect(
withChosenDocumentOptions(base, ['wire_doc', 'renamed_doc']),
).toEqual([
{ value: 'wire_doc', label: 'wire_doc' },
{ value: 'renamed_doc', label: 'renamed_doc' },
]);
});
it('surfaces a literal ref that has no upstream option', () => {
expect(withChosenDocumentOptions([], ['A1'])).toEqual([
{ value: 'A1', label: 'A1' },
]);
});
it('ignores blanks and the wildcard token', () => {
expect(
withChosenDocumentOptions(base, ['', ' ', ALL_INPUT_DOCUMENTS_TOKEN]),
).toEqual(base);
});
it('does not duplicate an already-known chosen name', () => {
expect(withChosenDocumentOptions(base, ['wire_doc', 'wire_doc'])).toEqual(
base,
);
});
});
@@ -0,0 +1,117 @@
import type { WorkflowVariable } from './components/PromptTextArea';
/** Token meaning "all uploaded input documents" in the input_documents list. */
export const ALL_INPUT_DOCUMENTS_TOKEN = '*';
/** How an agent node receives selected documents. */
export type FilePassing = 'auto' | 'native' | 'extract';
export const DEFAULT_FILE_PASSING: FilePassing = 'auto';
export const FILE_PASSING_OPTIONS: { value: FilePassing; label: string }[] = [
{ value: 'auto', label: 'Auto' },
{ value: 'native', label: 'Native' },
{ value: 'extract', label: 'Extract' },
];
/** Three-way Documents selector position derived from input_documents. */
export type DocumentsMode = 'all' | 'none' | 'choose';
/** Normalize an unknown file_passing value to a supported option. */
export function normalizeFilePassing(value: unknown): FilePassing {
return value === 'native' || value === 'extract' || value === 'auto'
? value
: DEFAULT_FILE_PASSING;
}
/** Derive the toggle position from the stored input_documents list. */
export function getDocumentsMode(
inputDocuments: string[] | undefined,
): DocumentsMode {
const docs = inputDocuments ?? [];
if (docs.length === 0) return 'none';
if (docs.length === 1 && docs[0] === ALL_INPUT_DOCUMENTS_TOKEN) return 'all';
return 'choose';
}
/** Build the input_documents list for a chosen toggle position. */
export function documentsModeToInputDocuments(
mode: DocumentsMode,
chosen: string[],
): string[] {
if (mode === 'all') return [ALL_INPUT_DOCUMENTS_TOKEN];
if (mode === 'none') return [];
return chosen.filter(
(name) => name.trim() !== '' && name !== ALL_INPUT_DOCUMENTS_TOKEN,
);
}
/** Strip a leading `agent.` / `agent['...']` wrapper to a bare variable name. */
export function stripAgentPrefix(templatePath: string): string {
const trimmed = templatePath.trim();
const bracketMatch = trimmed.match(/^agent\[(?:'|")(.*)(?:'|")\]$/);
if (bracketMatch) {
return bracketMatch[1].replace(/\\(['"\\])/g, '$1');
}
if (trimmed.startsWith('agent.')) {
return trimmed.slice('agent.'.length);
}
return trimmed;
}
/**
* Multiselect options of upstream document variables, stored as bare names.
*
* Only variables that carry artifact references at run time are offered
* (`producesArtifact`): the uploaded `input_documents` and code-node outputs.
* Plain agent/state TEXT outputs are excluded — the engine can't resolve them
* to artifacts, so picking one appends the literal variable name and the node
* hard-fails. The distinction is assigned upstream in `extractUpstreamVariables`
* keyed on the producing node type.
*/
export function toDocumentVariableOptions(
variables: WorkflowVariable[],
): { value: string; label: string }[] {
const options: { value: string; label: string }[] = [];
const seen = new Set<string>();
for (const variable of variables) {
if (!variable.producesArtifact) continue;
const bareName = stripAgentPrefix(variable.templatePath);
if (!bareName || seen.has(bareName)) continue;
seen.add(bareName);
options.push({ value: bareName, label: bareName });
}
return options;
}
/** Append a typed literal ref to a chosen list, ignoring blanks, the wildcard and duplicates. */
export function appendDocumentRef(chosen: string[], ref: string): string[] {
const value = ref.trim();
if (
value === '' ||
value === ALL_INPUT_DOCUMENTS_TOKEN ||
chosen.includes(value)
) {
return chosen;
}
return [...chosen, value];
}
/** Append stored chosen names lacking an upstream option so they stay visible and removable. */
export function withChosenDocumentOptions(
options: { value: string; label: string }[],
chosen: string[],
): { value: string; label: string }[] {
const known = new Set(options.map((option) => option.value));
const merged = [...options];
for (const name of chosen) {
const value = name.trim();
if (!value || value === ALL_INPUT_DOCUMENTS_TOKEN || known.has(value)) {
continue;
}
known.add(value);
merged.push({ value, label: value });
}
return merged;
}
@@ -5,7 +5,7 @@ interface BaseNodeProps {
title: string;
children?: ReactNode;
selected?: boolean;
type?: 'start' | 'end' | 'default' | 'state' | 'agent' | 'condition';
type?: 'start' | 'end' | 'default' | 'state' | 'agent' | 'condition' | 'code';
icon?: ReactNode;
handles?: {
source?: boolean;
@@ -42,6 +42,9 @@ export const BaseNode: React.FC<BaseNodeProps> = ({
} else if (type === 'condition') {
iconBg = 'bg-orange-100 dark:bg-orange-900/30';
iconColor = 'text-orange-600 dark:text-orange-400';
} else if (type === 'code') {
iconBg = 'bg-indigo-100 dark:bg-indigo-900/30';
iconColor = 'text-indigo-600 dark:text-indigo-400';
}
return (
@@ -0,0 +1,56 @@
import { Code2 } from 'lucide-react';
import { memo } from 'react';
import { NodeProps } from 'reactflow';
import { CodeNodeConfig } from '../../types/workflow';
import { BaseNode } from './BaseNode';
type CodeNodeData = {
title?: string;
label?: string;
config?: Partial<CodeNodeConfig>;
};
const CodeNode = ({ data, selected }: NodeProps<CodeNodeData>) => {
const title = data.title || data.label || 'Code';
const config = data.config || {};
const code = (config.code || '').trim();
const firstLine = code.split('\n').find((line) => line.trim() !== '') || '';
const codeHint =
firstLine.length > 40 ? `${firstLine.slice(0, 40)}…` : firstLine;
return (
<BaseNode
title={title}
type="code"
selected={selected}
icon={<Code2 size={16} />}
handles={{ source: true, target: true }}
>
<div className="flex flex-col gap-1">
{codeHint ? (
<div
className="truncate font-mono text-xs text-gray-500 dark:text-gray-400"
title={code}
>
{codeHint}
</div>
) : (
<div className="text-xs text-gray-500 dark:text-gray-400">
No code yet
</div>
)}
{config.output_variable && (
<div
className="truncate text-xs text-gray-500 dark:text-gray-400"
title={`Output: ${config.output_variable}`}
>
Output: {config.output_variable}
</div>
)}
</div>
</BaseNode>
);
};
export default memo(CodeNode);
@@ -2,6 +2,7 @@ import { Bot, Flag, Play, StickyNote } from 'lucide-react';
import { memo } from 'react';
import { BaseNode } from './BaseNode';
import CodeNode from './CodeNode';
import ConditionNode from './ConditionNode';
import SetStateNode from './SetStateNode';
@@ -145,3 +146,4 @@ export const NoteNode = memo(function NoteNode({
export { SetStateNode };
export { ConditionNode };
export { CodeNode };
@@ -0,0 +1,86 @@
import { describe, expect, it } from 'vitest';
import { buildSimpleCel, parseSimpleCel } from './simpleCel';
describe('parseSimpleCel', () => {
it('parses a plain variable comparison', () => {
expect(parseSimpleCel('status == "open"')).toEqual({
variable: 'status',
operator: '==',
value: 'open',
});
});
it('parses a dotted variable path the backend CEL accepts', () => {
// Regression: Simple mode blocked `audit.high_risk_accounts >= 3` at
// Preview ("must specify a variable") while createWorkflow accepted it.
expect(parseSimpleCel('audit.high_risk_accounts >= 3')).toEqual({
variable: 'audit.high_risk_accounts',
operator: '>=',
value: '3',
});
});
it('parses dotted paths with contains/startsWith', () => {
expect(parseSimpleCel('audit.summary.contains("risk")')).toEqual({
variable: 'audit.summary',
operator: 'contains',
value: 'risk',
});
expect(parseSimpleCel('doc.name.startsWith("Q3")')).toEqual({
variable: 'doc.name',
operator: 'startsWith',
value: 'Q3',
});
});
it('parses quoted comparisons on dotted paths', () => {
expect(parseSimpleCel('meta.region == "EMEA"')).toEqual({
variable: 'meta.region',
operator: '==',
value: 'EMEA',
});
});
it('keeps variable-less forms working', () => {
expect(parseSimpleCel('>= 4')).toEqual({
variable: '',
operator: '>=',
value: '4',
});
expect(parseSimpleCel('contains("x")')).toEqual({
variable: '',
operator: 'contains',
value: 'x',
});
});
it('falls back to an empty template for unparseable input', () => {
expect(parseSimpleCel('!!not cel!!')).toEqual({
variable: '',
operator: '==',
value: '',
});
});
});
describe('buildSimpleCel round-trip', () => {
it('round-trips a dotted numeric comparison', () => {
const built = buildSimpleCel('audit.high_risk_accounts', '>=', '3');
expect(built).toBe('audit.high_risk_accounts >= 3');
expect(parseSimpleCel(built)).toEqual({
variable: 'audit.high_risk_accounts',
operator: '>=',
value: '3',
});
});
it('round-trips a dotted contains', () => {
const built = buildSimpleCel('audit.summary', 'contains', 'risk');
expect(parseSimpleCel(built)).toEqual({
variable: 'audit.summary',
operator: 'contains',
value: 'risk',
});
});
});
+81
View File
@@ -0,0 +1,81 @@
/**
* Parse/build the Simple-mode condition expressions. The variable side accepts
* dotted paths (`audit.high_risk_accounts`) because the engine's CEL evaluator
* resolves them — Simple mode rejecting what the backend happily runs left
* saved workflows blocked at Preview with "must specify a variable".
*/
export function parseSimpleCel(expression: string): {
variable: string;
operator: string;
value: string;
} {
const trimmedExpression = expression.trim();
let match = trimmedExpression.match(
/^([\w.]+)\.(contains|startsWith)\(["'](.*)["']\)$/,
);
if (match) return { variable: match[1], operator: match[2], value: match[3] };
match = trimmedExpression.match(/^([\w.]+)\.(contains|startsWith)\((.*)\)$/);
if (match) {
const rawValue = match[3].trim();
const unquotedValue = rawValue.replace(/^["'](.*)["']$/, '$1');
return {
variable: match[1],
operator: match[2],
value: unquotedValue,
};
}
match = trimmedExpression.match(/^(contains|startsWith)\(["'](.*)["']\)$/);
if (match) return { variable: '', operator: match[1], value: match[2] };
match = trimmedExpression.match(/^(contains|startsWith)\((.*)\)$/);
if (match) {
const rawValue = match[2].trim();
const unquotedValue = rawValue.replace(/^["'](.*)["']$/, '$1');
return { variable: '', operator: match[1], value: unquotedValue };
}
match = trimmedExpression.match(
/^([\w.]+)\s*(==|!=|>=|<=|>|<)\s*["'](.*)["']$/,
);
if (match) return { variable: match[1], operator: match[2], value: match[3] };
match = trimmedExpression.match(/^(==|!=|>=|<=|>|<)\s*["'](.*)["']$/);
if (match) return { variable: '', operator: match[1], value: match[2] };
match = trimmedExpression.match(/^([\w.]+)\s*(==|!=|>=|<=|>|<)\s*(.*)$/);
if (match) return { variable: match[1], operator: match[2], value: match[3] };
match = trimmedExpression.match(/^(==|!=|>=|<=|>|<)\s*(.*)$/);
if (match) return { variable: '', operator: match[1], value: match[2] };
return { variable: '', operator: '==', value: '' };
}
export function buildSimpleCel(
variable: string,
operator: string,
value: string,
): string {
const trimmedValue = value.trim();
const isNumeric = trimmedValue !== '' && !isNaN(Number(trimmedValue));
const isBool = trimmedValue === 'true' || trimmedValue === 'false';
const literalValue =
isNumeric || isBool ? trimmedValue : JSON.stringify(value);
const stringValue = JSON.stringify(value);
if (operator === 'contains') {
return variable
? `${variable}.contains(${stringValue})`
: `contains(${stringValue})`;
}
if (operator === 'startsWith') {
return variable
? `${variable}.startsWith(${stringValue})`
: `startsWith(${stringValue})`;
}
if (!variable) return `${operator} ${literalValue}`;
return `${variable} ${operator} ${literalValue}`;
}
@@ -0,0 +1,119 @@
import { describe, expect, it } from 'vitest';
import { Attachment } from '../../upload/uploadSlice';
import reducer, {
addQuery,
collectCompletedAttachmentIds,
collectRunAttachmentIds,
previewSendBlockReason,
resetWorkflowPreview,
setPreviewOpen,
setWorkflowRunId,
UNSAVED_DRAFT_ATTACHMENTS_MESSAGE,
} from './workflowPreviewSlice';
const seedState = () => reducer(undefined, { type: '@@INIT' });
const att = (over: Partial<Attachment>): Attachment => ({
id: 'a1',
fileName: 'f.pdf',
progress: 100,
status: 'completed',
taskId: 't1',
...over,
});
describe('collectCompletedAttachmentIds', () => {
it('returns ids of completed attachments only', () => {
const ids = collectCompletedAttachmentIds([
att({ id: 'done', status: 'completed' }),
att({ id: 'busy', status: 'processing' }),
att({ id: 'up', status: 'uploading' }),
att({ id: 'bad', status: 'failed' }),
]);
expect(ids).toEqual(['done']);
});
it('drops completed rows with no server id and returns [] when none', () => {
expect(
collectCompletedAttachmentIds([att({ id: '', status: 'completed' })]),
).toEqual([]);
expect(collectCompletedAttachmentIds([])).toEqual([]);
});
});
describe('collectRunAttachmentIds', () => {
const completed = [att({ id: 'done', status: 'completed' })];
it('returns completed ids for a saved workflow', () => {
expect(collectRunAttachmentIds(completed, 'wf-1')).toEqual(['done']);
});
it('returns [] for an unsaved draft so uploads are neither sent nor cleared', () => {
expect(collectRunAttachmentIds(completed, null)).toEqual([]);
expect(collectRunAttachmentIds(completed, undefined)).toEqual([]);
expect(collectRunAttachmentIds(completed, '')).toEqual([]);
});
});
describe('previewSendBlockReason', () => {
it('blocks an unsaved draft that has completed attachments', () => {
expect(previewSendBlockReason(null, true)).toBe(
UNSAVED_DRAFT_ATTACHMENTS_MESSAGE,
);
expect(previewSendBlockReason(undefined, true)).toBe(
UNSAVED_DRAFT_ATTACHMENTS_MESSAGE,
);
expect(previewSendBlockReason('', true)).toBe(
UNSAVED_DRAFT_ATTACHMENTS_MESSAGE,
);
});
it('allows a saved workflow, or an unsaved draft with no attachments', () => {
expect(previewSendBlockReason('wf-1', true)).toBeNull();
expect(previewSendBlockReason('wf-1', false)).toBeNull();
expect(previewSendBlockReason(null, false)).toBeNull();
});
});
describe('setWorkflowRunId', () => {
it('stores the run id on the addressed query', () => {
let state = seedState();
state = reducer(state, addQuery({ prompt: 'run it' }));
state = reducer(
state,
setWorkflowRunId({ index: 0, workflowRunId: 'run-1' }),
);
expect(state.queries[0].workflowRunId).toBe('run-1');
});
it('ignores an out-of-range index without throwing', () => {
let state = seedState();
state = reducer(state, addQuery({ prompt: 'q' }));
state = reducer(
state,
setWorkflowRunId({ index: 5, workflowRunId: 'run-x' }),
);
expect(state.queries[0].workflowRunId).toBeUndefined();
});
});
describe('setPreviewOpen', () => {
it('defaults to closed and toggles the preview-open flag', () => {
let state = seedState();
expect(state.previewOpen).toBe(false);
state = reducer(state, setPreviewOpen(true));
expect(state.previewOpen).toBe(true);
state = reducer(state, setPreviewOpen(false));
expect(state.previewOpen).toBe(false);
});
it('survives a resetWorkflowPreview so the overlay flag stays set', () => {
let state = seedState();
state = reducer(state, setPreviewOpen(true));
state = reducer(state, addQuery({ prompt: 'q' }));
state = reducer(state, resetWorkflowPreview());
expect(state.queries).toEqual([]);
expect(state.previewOpen).toBe(true);
});
});
@@ -2,6 +2,7 @@ import { createAsyncThunk, createSlice, PayloadAction } from '@reduxjs/toolkit';
import conversationService from '../../api/services/conversationService';
import { Query, Status } from '../../conversation/conversationModels';
import { Attachment, clearAttachments } from '../../upload/uploadSlice';
import { WorkflowEdge, WorkflowNode } from '../types/workflow';
export interface WorkflowExecutionStep {
@@ -12,7 +13,7 @@ export interface WorkflowExecutionStep {
reasoning?: string;
startedAt?: number;
completedAt?: number;
stateSnapshot?: Record<string, unknown>;
stateDelta?: Record<string, unknown>;
output?: unknown;
error?: string;
}
@@ -26,6 +27,9 @@ interface WorkflowData {
export interface WorkflowQuery extends Query {
executionSteps?: WorkflowExecutionStep[];
// The run id for this query's execution; lets the artifacts panel list the
// run's produced artifacts (GET /api/artifacts?workflow_run_id=).
workflowRunId?: string;
}
export interface WorkflowPreviewState {
@@ -33,6 +37,9 @@ export interface WorkflowPreviewState {
status: Status;
executionSteps: WorkflowExecutionStep[];
activeNodeId: string | null;
// True while the Preview drawer is mounted; lets global overlays (the
// upload toast) shift off the drawer's right-side input.
previewOpen: boolean;
}
const initialState: WorkflowPreviewState = {
@@ -40,6 +47,7 @@ const initialState: WorkflowPreviewState = {
status: 'idle',
executionSteps: [],
activeNodeId: null,
previewOpen: false,
};
let abortController: AbortController | null = null;
@@ -56,6 +64,50 @@ interface ThunkState {
token: string | null;
};
workflowPreview: WorkflowPreviewState;
upload: { attachments: Attachment[] };
}
/** Server-assigned ids of completed attachments, for the run's ``attachments``. */
export function collectCompletedAttachmentIds(
attachments: Attachment[],
): string[] {
return attachments
.filter((att) => att.status === 'completed' && att.id)
.map((att) => att.id);
}
/** User-facing message shown when a Preview send is blocked. */
export const UNSAVED_DRAFT_ATTACHMENTS_MESSAGE =
'Save the workflow before attaching documents in Preview.';
/**
* Reason a Preview send must be blocked, or null when it may proceed. An
* unsaved draft has no persisted ``workflow_id``, so the backend can't bridge
* uploaded documents into the run — it would execute with the docs silently
* dropped. Block instead, keeping the uploads for a retry after the save.
*/
export function previewSendBlockReason(
workflowId: string | null | undefined,
hasCompletedAttachment: boolean,
): string | null {
if (!workflowId && hasCompletedAttachment) {
return UNSAVED_DRAFT_ATTACHMENTS_MESSAGE;
}
return null;
}
/**
* Completed attachment ids eligible for a run. Attachments only survive as
* run-scoped artifacts when the run persists a ``workflow_runs`` row, which
* needs a saved ``workflow_id``; on an unsaved draft none are eligible (the UI
* blocks that send upstream), so they are neither sent nor cleared.
*/
export function collectRunAttachmentIds(
attachments: Attachment[],
workflowId: string | null | undefined,
): string[] {
if (!workflowId) return [];
return collectCompletedAttachmentIds(attachments);
}
export const fetchWorkflowPreviewAnswer = createAsyncThunk<
@@ -64,11 +116,15 @@ export const fetchWorkflowPreviewAnswer = createAsyncThunk<
question: string;
workflowData: WorkflowData;
indx?: number;
workflowId?: string | null;
},
{ state: ThunkState }
>(
'workflowPreview/fetchAnswer',
async ({ question, workflowData, indx }, { dispatch, getState }) => {
async (
{ question, workflowData, indx, workflowId },
{ dispatch, getState },
) => {
if (abortController) abortController.abort();
abortController = new AbortController();
const { signal } = abortController;
@@ -76,9 +132,24 @@ export const fetchWorkflowPreviewAnswer = createAsyncThunk<
const state = getState();
if (state.preference) {
// Doc-driven workflows read uploaded files as run-scoped artifacts; pass
// the completed attachment ids so the backend bridges them (user-scoped).
// Gated on a saved workflow_id: an unsaved draft can't bridge them, so we
// neither send nor clear the uploads (the Preview UI blocks that send and
// keeps them for a retry after saving).
const attachmentIds = collectRunAttachmentIds(
state.upload.attachments,
workflowId,
);
if (attachmentIds.length > 0) dispatch(clearAttachments());
// A saved workflow id lets the run persist a ``workflow_runs`` row so its
// produced artifacts are listable + authz'd; omitted for unsaved drafts.
const payload = {
question,
workflow: workflowData,
...(workflowId ? { workflow_id: workflowId } : {}),
...(attachmentIds.length > 0 ? { attachments: attachmentIds } : {}),
save_conversation: false,
visibility: 'hidden',
};
@@ -117,6 +188,13 @@ export const fetchWorkflowPreviewAnswer = createAsyncThunk<
if (data.type === 'end') {
dispatch(workflowPreviewSlice.actions.setStatus('idle'));
} else if (data.type === 'workflow_run') {
dispatch(
setWorkflowRunId({
index: targetIndex,
workflowRunId: data.workflow_run_id,
}),
);
} else if (data.type === 'thought') {
dispatch(
updateThought({
@@ -134,7 +212,8 @@ export const fetchWorkflowPreviewAnswer = createAsyncThunk<
nodeTitle: data.node_title,
status: data.status,
reasoning: data.reasoning,
stateSnapshot: data.state_snapshot,
// Older streams (pre-rename) still say state_snapshot.
stateDelta: data.state_delta ?? data.state_snapshot,
output: data.output,
error: data.error,
},
@@ -332,7 +411,7 @@ export const workflowPreviewSlice = createSlice({
nodeTitle: step.nodeTitle,
status: step.status,
reasoning: step.reasoning,
stateSnapshot: step.stateSnapshot,
stateDelta: step.stateDelta,
output: step.output,
error: step.error,
startedAt:
@@ -348,8 +427,8 @@ export const workflowPreviewSlice = createSlice({
};
if (existingIndex !== -1) {
updatedStep.stateSnapshot =
step.stateSnapshot ?? querySteps[existingIndex].stateSnapshot;
updatedStep.stateDelta =
step.stateDelta ?? querySteps[existingIndex].stateDelta;
updatedStep.output = step.output ?? querySteps[existingIndex].output;
updatedStep.error = step.error ?? querySteps[existingIndex].error;
querySteps[existingIndex] = updatedStep;
@@ -366,6 +445,15 @@ export const workflowPreviewSlice = createSlice({
state.executionSteps.push(updatedStep);
}
},
setWorkflowRunId(
state,
action: PayloadAction<{ index: number; workflowRunId: string }>,
) {
const { index, workflowRunId } = action.payload;
if (state.queries[index]) {
state.queries[index].workflowRunId = workflowRunId;
}
},
setActiveNodeId(state, action: PayloadAction<string | null>) {
state.activeNodeId = action.payload;
},
@@ -393,6 +481,9 @@ export const workflowPreviewSlice = createSlice({
state.executionSteps = [];
state.activeNodeId = null;
},
setPreviewOpen: (state, action: PayloadAction<boolean>) => {
state.previewOpen = action.payload;
},
},
extraReducers(builder) {
builder
@@ -430,6 +521,9 @@ export const selectWorkflowExecutionSteps = (
) => state.workflowPreview.executionSteps;
export const selectActiveNodeId = (state: RootStateWithWorkflowPreview) =>
state.workflowPreview.activeNodeId;
export const selectWorkflowPreviewOpen = (
state: RootStateWithWorkflowPreview,
) => state.workflowPreview.previewOpen;
export const {
addQuery,
@@ -440,11 +534,13 @@ export const {
updateStreamingSource,
updateToolCall,
updateExecutionStep,
setWorkflowRunId,
setActiveNodeId,
setStatus,
raiseError,
resetWorkflowPreview,
clearExecutionSteps,
setPreviewOpen,
} = workflowPreviewSlice.actions;
export default workflowPreviewSlice.reducer;
+17
View File
@@ -110,6 +110,23 @@ const endpoints = {
AGENT_FOLDER: (id: string) => `/api/agents/folders/${id}`,
MOVE_AGENT_TO_FOLDER: '/api/agents/folders/move_agent',
GET_ARTIFACT: (artifactId: string) => `/api/artifact/${artifactId}`,
GET_DOCUMENT_ARTIFACT: (artifactId: string) =>
`/api/artifacts/${artifactId}`,
LIST_WORKFLOW_RUN_ARTIFACTS: (workflowRunId: string) =>
`/api/artifacts?workflow_run_id=${encodeURIComponent(workflowRunId)}`,
DOWNLOAD_ARTIFACT: (
artifactId: string,
version?: number,
disposition?: string,
) => {
const params = new URLSearchParams();
if (version != null) params.set('version', String(version));
if (disposition) params.set('disposition', disposition);
const qs = params.toString();
return `/api/artifacts/${artifactId}/download${qs ? `?${qs}` : ''}`;
},
RESTORE_ARTIFACT: (artifactId: string) =>
`/api/artifacts/${artifactId}/restore`,
WORKFLOWS: '/api/workflows',
WORKFLOW: (id: string) => `/api/workflows/${id}`,
CUSTOM_MODELS: '/api/user/models',
+37
View File
@@ -322,6 +322,43 @@ const userService = {
apiClient.post(endpoints.USER.MOVE_AGENT_TO_FOLDER, data, token),
getArtifact: (artifactId: string, token: string | null): Promise<any> =>
apiClient.get(endpoints.USER.GET_ARTIFACT(artifactId), token),
getDocumentArtifact: (
artifactId: string,
token: string | null,
): Promise<Response> =>
apiClient.get(endpoints.USER.GET_DOCUMENT_ARTIFACT(artifactId), token),
listWorkflowRunArtifacts: (
workflowRunId: string,
token: string | null,
): Promise<Response> =>
apiClient.get(
endpoints.USER.LIST_WORKFLOW_RUN_ARTIFACTS(workflowRunId),
token,
),
downloadArtifact: (
artifactId: string,
token: string | null,
version?: number,
disposition?: 'url',
): Promise<Response> =>
apiClient.get(
endpoints.USER.DOWNLOAD_ARTIFACT(artifactId, version, disposition),
token,
// ?disposition=url asks the s3 strategy to return the presigned URL as
// JSON (for a top-level navigation) instead of a CORS-blocked 302; the
// Accept header is the same opt-in via content negotiation.
disposition === 'url' ? { Accept: 'application/json' } : {},
),
restoreArtifactVersion: (
artifactId: string,
version: number,
token: string | null,
): Promise<Response> =>
apiClient.post(
endpoints.USER.RESTORE_ARTIFACT(artifactId),
{ version },
token,
),
getWorkflow: (id: string, token: string | null): Promise<any> =>
apiClient.get(endpoints.USER.WORKFLOW(id), token),
createWorkflow: (data: any, token: string | null): Promise<any> =>
@@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" height="24px" viewBox="0 -960 960 960" width="24px" fill="currentColor"><path d="M320-240h320v-80H320v80Zm0-160h320v-80H320v80ZM240-80q-33 0-56.5-23.5T160-160v-640q0-33 23.5-56.5T240-880h320l240 240v480q0 33-23.5 56.5T720-80H240Zm280-520v-200H240v640h480v-440H520Z"/></svg>

After

Width:  |  Height:  |  Size: 313 B

@@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" height="24px" viewBox="0 -960 960 960" width="24px" fill="currentColor"><path d="M160-160q-33 0-56.5-23.5T80-240v-480q0-33 23.5-56.5T160-800h640q33 0 56.5 23.5T880-720v480q0 33-23.5 56.5T800-160H160Zm0-80h640v-400H160v400Zm140-40-56-56 103-104-104-104 57-56 160 160-160 160Zm180 0v-80h200v80H480Z"/></svg>

After

Width:  |  Height:  |  Size: 345 B

@@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" height="24px" viewBox="0 -960 960 960" width="24px" fill="currentColor"><path d="M80-680v-200h200v80H160v120H80Zm600 0v-120H680v-80h200v200h-80ZM80-80v-200h80v120h120v80H80Zm600 0v-80h120v-120h80v200H680ZM280-280h400v-80H280v80Zm0-160h400v-80H280v80Z"/></svg>

After

Width:  |  Height:  |  Size: 299 B

+116 -59
View File
@@ -12,7 +12,9 @@ import remarkGfm from 'remark-gfm';
import userService from '../api/services/userService';
import { useDarkTheme } from '../hooks';
import { selectToken } from '../preferences/preferenceSlice';
import { isDocumentArtifact, type DocumentArtifact } from './artifactViewUtils';
import CopyButton from './CopyButton';
import DocumentArtifactView from './DocumentArtifactView';
import Spinner from './Spinner';
import { Button } from './ui/button';
import { Sheet, SheetContent } from './ui/sheet';
@@ -282,13 +284,16 @@ export default function ArtifactSidebar({
const currentFetchIdRef = React.useRef<string | null>(null);
const token = useSelector(selectToken);
const [artifact, setArtifact] = useState<ArtifactData | null>(null);
const [documentArtifact, setDocumentArtifact] =
useState<DocumentArtifact | null>(null);
const [loading, setLoading] = useState(false);
const [error, setError] = useState<string | null>(null);
const [refreshNonce, setRefreshNonce] = useState(0);
const [effectiveArtifactId, setEffectiveArtifactId] = useState<string | null>(
artifactId,
);
const title = getArtifactTitle(artifact, toolName);
const title = documentArtifact?.title || getArtifactTitle(artifact, toolName);
// Reset last successful todo artifact ID when conversation changes
useEffect(() => {
@@ -308,6 +313,7 @@ export default function ArtifactSidebar({
useEffect(() => {
if (!isOpen || !effectiveArtifactId) {
setArtifact(null);
setDocumentArtifact(null);
setError(null);
setLoading(false);
currentFetchIdRef.current = null;
@@ -315,76 +321,118 @@ export default function ArtifactSidebar({
}
// Generate a unique ID for this fetch
const fetchId = `${effectiveArtifactId}-${Date.now()}`;
const artifactIdToFetch = effectiveArtifactId;
const fetchId = `${artifactIdToFetch}-${Date.now()}`;
currentFetchIdRef.current = fetchId;
setLoading(true);
setError(null);
// Note: For todo artifacts, the endpoint always returns all todos for the tool; will be coversation scoped later
// Document/file artifacts live behind the generalized /api/artifacts/<id>
// endpoint; notes/todos behind the legacy /api/artifact/<id>. Try the
// document path first and fall back to the legacy flow on a non-document
// shape so a single id resolves cleanly either way.
userService
.getArtifact(effectiveArtifactId, token)
.then(async (res: any) => {
// Ignore if this is not the current fetch
if (currentFetchIdRef.current !== fetchId) return;
const isResponseLike = res && typeof res.json === 'function';
const status = isResponseLike ? res.status : undefined;
const ok = isResponseLike ? Boolean(res.ok) : true;
let data: any = res;
if (isResponseLike) {
try {
data = await res.json();
} catch {
data = null;
}
.getDocumentArtifact(artifactIdToFetch, token)
.then(async (res: Response) => {
if (currentFetchIdRef.current !== fetchId) return true;
if (!res.ok) return false;
let data: any = null;
try {
data = await res.json();
} catch {
data = null;
}
// Check again after async operation
if (currentFetchIdRef.current !== fetchId) return;
if (ok && data?.success && data?.artifact) {
setArtifact(data.artifact);
// Remember the last successful todo artifact id so we can fallback if a newer id 404s.
if (data.artifact?.artifact_type === 'todo_list') {
lastSuccessfulTodoArtifactIdRef.current = effectiveArtifactId;
}
if (currentFetchIdRef.current !== fetchId) return true;
if (data?.success && isDocumentArtifact(data.artifact)) {
setDocumentArtifact(data.artifact);
setArtifact(null);
setLoading(false);
return;
return true;
}
const isTodoTool = (toolName ?? '').toLowerCase().includes('todo');
// If the latest todo artifact id is missing (404), fall back to the last known good one
// so the backend can still resolve `tool_id` for the todo list.
if (
status === 404 &&
isTodoTool &&
lastSuccessfulTodoArtifactIdRef.current &&
lastSuccessfulTodoArtifactIdRef.current !== effectiveArtifactId
) {
// Update effectiveArtifactId to trigger a new fetch with the fallback id
setEffectiveArtifactId(lastSuccessfulTodoArtifactIdRef.current);
setLoading(false);
return;
}
// Ensure we show a visible error state instead of rendering nothing.
const message =
data?.message ||
(status === 404 ? 'Artifact not found' : null) ||
'Failed to load artifact';
setError(message);
setLoading(false);
return false;
})
.catch((err) => {
// Ignore if this is not the current fetch
.catch(() => false)
.then((handled) => {
if (handled) return;
if (currentFetchIdRef.current !== fetchId) return;
setError('Failed to fetch artifact');
setLoading(false);
setDocumentArtifact(null);
fetchLegacyArtifact(fetchId);
});
}, [isOpen, effectiveArtifactId, token, toolName, conversationId]);
// Legacy notes/todo fetch (unchanged behavior).
function fetchLegacyArtifact(fetchId: string) {
userService
.getArtifact(artifactIdToFetch, token)
.then(async (res: any) => {
// Ignore if this is not the current fetch
if (currentFetchIdRef.current !== fetchId) return;
const isResponseLike = res && typeof res.json === 'function';
const status = isResponseLike ? res.status : undefined;
const ok = isResponseLike ? Boolean(res.ok) : true;
let data: any = res;
if (isResponseLike) {
try {
data = await res.json();
} catch {
data = null;
}
}
// Check again after async operation
if (currentFetchIdRef.current !== fetchId) return;
if (ok && data?.success && data?.artifact) {
setArtifact(data.artifact);
// Remember the last successful todo artifact id so we can fallback if a newer id 404s.
if (data.artifact?.artifact_type === 'todo_list') {
lastSuccessfulTodoArtifactIdRef.current = effectiveArtifactId;
}
setLoading(false);
return;
}
const isTodoTool = (toolName ?? '').toLowerCase().includes('todo');
// If the latest todo artifact id is missing (404), fall back to the last known good one
// so the backend can still resolve `tool_id` for the todo list.
if (
status === 404 &&
isTodoTool &&
lastSuccessfulTodoArtifactIdRef.current &&
lastSuccessfulTodoArtifactIdRef.current !== effectiveArtifactId
) {
// Update effectiveArtifactId to trigger a new fetch with the fallback id
setEffectiveArtifactId(lastSuccessfulTodoArtifactIdRef.current);
setLoading(false);
return;
}
// Ensure we show a visible error state instead of rendering nothing.
const message =
data?.message ||
(status === 404 ? 'Artifact not found' : null) ||
'Failed to load artifact';
setError(message);
setLoading(false);
})
.catch(() => {
// Ignore if this is not the current fetch
if (currentFetchIdRef.current !== fetchId) return;
setError('Failed to fetch artifact');
setLoading(false);
});
}
}, [
isOpen,
effectiveArtifactId,
token,
toolName,
conversationId,
refreshNonce,
]);
const renderContent = () => {
if (loading) {
@@ -401,6 +449,14 @@ export default function ArtifactSidebar({
</div>
);
}
if (documentArtifact) {
return (
<DocumentArtifactView
artifact={documentArtifact}
onRefresh={() => setRefreshNonce((n) => n + 1)}
/>
);
}
// Avoid rendering an empty panel if the artifact couldn't be loaded for any reason.
if (!artifact) {
return (
@@ -465,6 +521,7 @@ export default function ArtifactSidebar({
<SheetContent
side="right"
showCloseButton={false}
title={title || 'Artifact preview'}
className="dark:bg-card bg-card flex h-full w-80 flex-col gap-0 border-l border-[#9ca3af]/10 p-0 sm:w-96 sm:max-w-none"
>
<div className="flex w-full items-center justify-between border-b border-gray-200 px-4 py-3 dark:border-gray-700">
@@ -0,0 +1,435 @@
import mermaid from 'mermaid';
import { Download, FileText, History, RotateCcw } from 'lucide-react';
import { useEffect, useMemo, useRef, useState, type ReactNode } from 'react';
import { useSelector } from 'react-redux';
import userService from '../api/services/userService';
import { useDarkTheme } from '../hooks';
import { selectToken } from '../preferences/preferenceSlice';
import MarkdownPreview from './MarkdownPreview';
import Spinner from './Spinner';
import {
buildPreviewDocument,
bytesPreviewModeForMime,
displayFilename,
findCurrentVersion,
formatBytes,
previewModeForKind,
sortVersionsDesc,
triggerResponseDownload,
type ArtifactVersion,
type BytesPreviewMode,
type DocumentArtifact,
} from './artifactViewUtils';
import { Button } from './ui/button';
import { useArtifactBytes } from './useArtifactBytes';
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from './ui/select';
function asPreviewString(spec: unknown): string {
if (typeof spec === 'string') return spec;
if (spec && typeof spec === 'object') {
const rec = spec as Record<string, unknown>;
for (const key of ['html', 'svg', 'content', 'source', 'code']) {
if (typeof rec[key] === 'string') return rec[key] as string;
}
}
return '';
}
function FramePreview({
artifactKind,
source,
}: {
artifactKind: string;
source: string;
}) {
const [isDarkTheme] = useDarkTheme();
const [rendered, setRendered] = useState<string | null>(null);
const [renderError, setRenderError] = useState(false);
const isMermaid = artifactKind.toLowerCase() === 'mermaid';
// `securityLevel: 'sandbox'` makes mermaid build the diagram's live DOM inside
// its OWN opaque-origin `sandbox=""` iframe (not the app document) and return
// an `<iframe src="data:text/html;base64,…">` wrapper, so untrusted diagram
// text never becomes live nodes in the app origin. The wrapper is embedded
// into the outer scriptless `sandbox=""` preview iframe below.
useEffect(() => {
if (!isMermaid) {
setRendered(source);
return;
}
let cancelled = false;
setRendered(null);
setRenderError(false);
mermaid.initialize({
startOnLoad: false,
securityLevel: 'sandbox',
theme: isDarkTheme ? 'dark' : 'default',
});
mermaid
.render(`artifact-mermaid-${Date.now()}`, source)
.then(({ svg }) => {
if (!cancelled) setRendered(svg);
})
.catch(() => {
if (!cancelled) setRenderError(true);
});
return () => {
cancelled = true;
};
}, [isMermaid, source, isDarkTheme]);
if (renderError) {
return (
<pre className="overflow-auto p-4 text-xs text-gray-600 dark:text-gray-400">
{source}
</pre>
);
}
if (rendered === null) {
return (
<div className="flex h-full items-center justify-center">
<Spinner />
</div>
);
}
// For mermaid, `rendered` is mermaid's sandboxed-iframe wrapper; the 'mermaid'
// preview document carries a CSP that permits its inner `data:` frame.
const previewKind = isMermaid ? 'mermaid' : artifactKind;
return (
<iframe
// No allow-same-origin: content is fully isolated from the app origin.
sandbox=""
title="Artifact preview"
className="h-full w-full rounded-md border border-gray-200 bg-white dark:border-gray-700"
srcDoc={buildPreviewDocument(previewKind, rendered)}
/>
);
}
function DownloadCard({
filename,
size,
onDownload,
downloading,
}: {
filename: string;
size: number | null;
onDownload: () => void;
downloading: boolean;
}) {
return (
<div className="flex h-full flex-col items-center justify-center gap-4">
<div className="flex flex-col items-center gap-3 rounded-xl border border-gray-200 p-8 text-center dark:border-gray-700">
<FileText className="h-12 w-12 text-gray-400" />
<div>
<p className="text-sm font-medium break-all text-gray-800 dark:text-gray-200">
{filename}
</p>
{size != null && (
<p className="text-xs text-gray-500 dark:text-gray-400">
{formatBytes(size)}
</p>
)}
</div>
<Button
type="button"
variant="outline"
size="sm"
onClick={onDownload}
disabled={downloading}
>
{downloading ? (
<Spinner size="small" />
) : (
<Download className="h-4 w-4" />
)}
Download
</Button>
</div>
</div>
);
}
/**
* Render a version's downloadable BYTES inline by fetching them through the
* authed download path. html/svg are embedded ONLY inside the scriptless
* `sandbox=""` + CSP iframe (never the app DOM); images are shown via an inert
* blob-URL `<img>`; text/markdown render as escaped React content. On fetch
* error (or oversized text) it falls back to the supplied download card.
*/
function BytesPreview({
artifactId,
version,
mode,
token,
fallback,
}: {
artifactId: string;
version: number;
mode: Exclude<BytesPreviewMode, 'card'>;
token: string | null;
fallback: ReactNode;
}) {
const state = useArtifactBytes(artifactId, version, mode, token);
if (state.status === 'loading') {
return (
<div className="flex h-full items-center justify-center">
<Spinner />
</div>
);
}
if (state.status === 'error') {
return <>{fallback}</>;
}
if (state.status === 'image') {
return (
<div className="flex h-full items-center justify-center overflow-auto p-4">
<img
src={state.url}
alt="Artifact preview"
className="max-h-full max-w-full object-contain"
/>
</div>
);
}
// state.status === 'text'
if (mode === 'iframe-html' || mode === 'iframe-svg') {
return (
<iframe
// No allow-same-origin / no scripts: bytes are isolated from the app origin.
sandbox=""
title="Artifact preview"
className="h-full w-full rounded-md border border-gray-200 bg-white dark:border-gray-700"
srcDoc={buildPreviewDocument(
mode === 'iframe-svg' ? 'svg' : 'html',
state.text,
)}
/>
);
}
if (mode === 'text-markdown') {
return <MarkdownPreview content={state.text} />;
}
return (
<pre className="h-full overflow-auto p-4 text-xs whitespace-pre-wrap text-gray-700 dark:text-gray-300">
{state.text}
</pre>
);
}
export default function DocumentArtifactView({
artifact,
onRefresh,
}: {
artifact: DocumentArtifact;
onRefresh: () => void;
}) {
const token = useSelector(selectToken);
const versionsDesc = useMemo(
() => sortVersionsDesc(artifact.versions),
[artifact.versions],
);
const [selectedVersion, setSelectedVersion] = useState<number>(
artifact.current_version,
);
const [downloading, setDownloading] = useState(false);
const [restoring, setRestoring] = useState(false);
const [actionError, setActionError] = useState<string | null>(null);
// Guards async setState after the sidebar closes mid-request (the download /
// restore promises can resolve after unmount), mirroring the sidebar's
// currentFetchIdRef discipline.
const isMountedRef = useRef(true);
useEffect(() => {
isMountedRef.current = true;
return () => {
isMountedRef.current = false;
};
}, []);
useEffect(() => {
setSelectedVersion(artifact.current_version);
}, [artifact.id, artifact.current_version]);
const selectedRow: ArtifactVersion | null = useMemo(() => {
return (
artifact.versions.find((v) => v.version === selectedVersion) ??
findCurrentVersion(artifact)
);
}, [artifact, selectedVersion]);
const isCurrent = selectedVersion === artifact.current_version;
const mode = previewModeForKind(artifact.kind, selectedRow?.mime_type);
const filename = displayFilename(artifact, selectedRow);
const handleDownload = async () => {
setActionError(null);
setDownloading(true);
try {
const response = await userService.downloadArtifact(
artifact.id,
token,
selectedVersion,
'url',
);
const ok = await triggerResponseDownload(response, filename);
if (!ok && isMountedRef.current) setActionError('Download failed');
} catch {
if (isMountedRef.current) setActionError('Download failed');
} finally {
if (isMountedRef.current) setDownloading(false);
}
};
const handleRestore = async () => {
setActionError(null);
setRestoring(true);
try {
const response = await userService.restoreArtifactVersion(
artifact.id,
selectedVersion,
token,
);
if (!response.ok) {
if (isMountedRef.current) setActionError('Restore failed');
return;
}
if (isMountedRef.current) onRefresh();
} catch {
if (isMountedRef.current) setActionError('Restore failed');
} finally {
if (isMountedRef.current) setRestoring(false);
}
};
const downloadCard = (
<DownloadCard
filename={filename}
size={selectedRow?.size ?? null}
onDownload={handleDownload}
downloading={downloading}
/>
);
const renderPreview = () => {
// Preferred path: the version's `spec` already carries renderable
// html/svg/mermaid markup — render it in the sandboxed frame, no fetch.
if (mode === 'frame') {
const source = asPreviewString(
isCurrent ? artifact.spec : selectedRow?.preview_text,
);
if (source) {
return <FramePreview artifactKind={artifact.kind} source={source} />;
}
// No usable spec content: fall through to the bytes path below.
}
// Inline text already on hand (code/data kinds expose it as preview_text /
// spec) — render it without a fetch.
if (mode === 'text') {
const text = isCurrent
? asPreviewString(artifact.spec) || (selectedRow?.preview_text ?? '')
: (selectedRow?.preview_text ?? '');
if (text) {
return (
<pre className="h-full overflow-auto p-4 text-xs whitespace-pre-wrap text-gray-700 dark:text-gray-300">
{text}
</pre>
);
}
// No inline text: try fetching the bytes below.
}
// Bytes path: no usable spec/inline content, so decide the preview by the
// version's mime (kind fallback) and fetch the stored bytes for an inline
// render. Anything not previewable (office docs, pdf, octet-stream) and any
// fetch failure fall back to the download card.
const bytesMode: BytesPreviewMode = bytesPreviewModeForMime(
selectedRow?.mime_type,
artifact.kind,
);
if (bytesMode === 'card' || !selectedRow) {
return downloadCard;
}
return (
<BytesPreview
// Refetch on artifact/version/mode change via the key + hook deps.
key={`${artifact.id}:${selectedVersion}:${bytesMode}`}
artifactId={artifact.id}
version={selectedVersion}
mode={bytesMode}
token={token}
fallback={downloadCard}
/>
);
};
return (
<div className="flex h-full w-full flex-col gap-3 overflow-hidden">
<div className="flex flex-wrap items-center gap-2">
<Select
value={String(selectedVersion)}
onValueChange={(v) => setSelectedVersion(Number(v))}
>
<SelectTrigger className="h-8 w-auto gap-2 rounded-md px-3 text-xs">
<History className="h-3.5 w-3.5" />
<SelectValue />
</SelectTrigger>
<SelectContent>
{versionsDesc.map((v) => (
<SelectItem key={v.version} value={String(v.version)}>
Version {v.version}
{v.version === artifact.current_version ? ' (current)' : ''}
</SelectItem>
))}
</SelectContent>
</Select>
<Button
type="button"
variant="outline"
size="sm"
onClick={handleDownload}
disabled={downloading}
>
{downloading ? (
<Spinner size="small" />
) : (
<Download className="h-4 w-4" />
)}
Download
</Button>
{!isCurrent && (
<Button
type="button"
variant="ghost"
size="sm"
onClick={handleRestore}
disabled={restoring}
>
{restoring ? (
<Spinner size="small" />
) : (
<RotateCcw className="h-4 w-4" />
)}
Restore
</Button>
)}
</div>
{actionError && <p className="text-xs text-red-500">{actionError}</p>}
<div className="min-h-0 flex-1 overflow-hidden">{renderPreview()}</div>
</div>
);
}
@@ -0,0 +1,74 @@
import ReactMarkdown from 'react-markdown';
import remarkGfm from 'remark-gfm';
/**
* Minimal GFM markdown renderer for artifact text previews. ReactMarkdown
* parses to React elements (no `dangerouslySetInnerHTML`), so untrusted
* markdown bytes never become raw DOM in the app origin.
*/
export default function MarkdownPreview({ content }: { content: string }) {
return (
<div className="h-full overflow-auto p-4">
<ReactMarkdown
className="flex flex-col gap-3 text-sm leading-normal break-words whitespace-pre-wrap text-gray-800 dark:text-gray-200"
remarkPlugins={[remarkGfm]}
components={{
a({ children, href }) {
return (
<a
href={href}
target="_blank"
rel="noopener noreferrer"
className="text-blue-600 hover:underline dark:text-blue-400"
>
{children}
</a>
);
},
ul({ children }) {
return (
<ul className="list-inside list-disc pl-4 whitespace-normal">
{children}
</ul>
);
},
ol({ children }) {
return (
<ol className="list-inside list-decimal pl-4 whitespace-normal">
{children}
</ol>
);
},
h1({ children }) {
return <h1 className="text-xl font-bold">{children}</h1>;
},
h2({ children }) {
return <h2 className="text-lg font-bold">{children}</h2>;
},
h3({ children }) {
return <h3 className="text-base font-bold">{children}</h3>;
},
code({ children, ...rest }) {
return (
<code
className="dark:bg-accent dark:text-foreground rounded-md bg-gray-200 px-2 py-1 text-xs font-normal"
{...rest}
>
{children}
</code>
);
},
blockquote({ children }) {
return (
<blockquote className="border-l-4 border-gray-300 pl-4 italic dark:border-gray-600">
{children}
</blockquote>
);
},
}}
>
{content}
</ReactMarkdown>
</div>
);
}
Loaded 100 of 159 files, more files were not shown because too many files have changed in this diff. Show more