mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 18:15:33 +00:00
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:
159 files changed
+23593
-732
No files matched your search
@@ -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
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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__
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
@@ -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"])
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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, {})
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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;")
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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:"
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
@@ -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 %}
|
||||
@@ -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 %}
|
||||
@@ -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 %}
|
||||
@@ -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
|
||||
@@ -0,0 +1 @@
|
||||
"""Pluggable code-execution sandbox abstraction for running untrusted LLM code."""
|
||||
@@ -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
|
||||
@@ -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)."""
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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: {}
|
||||
@@ -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"]
|
||||
@@ -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.
|
||||
Executable
+38
@@ -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
|
||||
Executable
+20
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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',
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
@@ -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',
|
||||
|
||||
@@ -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 |
@@ -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
Reference in new issue
Block a user