From 1ac1f793b7c39c316d4536c5e9b544d8a6509052 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 24 Jun 2026 12:30:18 +0100 Subject: [PATCH] Add workflow code node and artifacts templating namespace Add a code workflow node that runs code in the run-scoped sandbox session and writes produced files as artifact references into workflow state, passing them by reference (only id and metadata, never bytes) so downstream nodes and CEL conditions can branch on them. Add an artifacts.* templating namespace that resolves those references to metadata via a run-scoped lookup, available to both the workflow engine and the prompt renderer. Extract the sandbox-to- artifact persistence into a shared helper reused by the code node and the code_executor tool. --- application/agents/tools/code_executor.py | 184 ++--------- application/agents/workflow_agent.py | 3 + application/agents/workflows/schemas.py | 12 + .../agents/workflows/workflow_engine.py | 194 +++++++++++- application/api/user/workflows/routes.py | 13 + application/sandbox/artifacts_capture.py | 294 ++++++++++++++++++ .../storage/db/repositories/workflow_runs.py | 5 + application/templates/namespaces.py | 103 ++++++ tests/agents/test_workflow_agent_graph.py | 7 + tests/agents/test_workflow_code_node.py | 262 ++++++++++++++++ tests/agents/test_workflow_schemas.py | 2 + .../test_workflow_code_node_e2e.py | 278 +++++++++++++++++ tests/test_namespaces.py | 136 ++++++++ 13 files changed, 1337 insertions(+), 156 deletions(-) create mode 100644 application/sandbox/artifacts_capture.py create mode 100644 tests/agents/test_workflow_code_node.py create mode 100644 tests/integration/test_workflow_code_node_e2e.py diff --git a/application/agents/tools/code_executor.py b/application/agents/tools/code_executor.py index a815f6f2..095505a2 100644 --- a/application/agents/tools/code_executor.py +++ b/application/agents/tools/code_executor.py @@ -2,26 +2,32 @@ from __future__ import annotations -import hashlib -import io import logging -import mimetypes import re from typing import Any, Dict, List, Optional, Tuple -from sqlalchemy import text - from application.agents.tools.base import Tool from application.core.settings import settings +from application.sandbox.artifacts_capture import ( + MAX_CAPTURED_FILES, + capture_artifacts, + infer_mime as _infer_mime, + kind_for_mime as _kind_for_mime, + snapshot_signatures, +) 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, db_session +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 @@ -30,22 +36,6 @@ _OUTPUT_TAIL_BYTES = 4000 # [A-Za-z0-9_-]+, so any disallowed character is stripped before binding. _SESSION_ID_RE = re.compile(r"[^A-Za-z0-9_-]+") -_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 _tail(stream: Optional[str]) -> str: """Return the trailing slice of ``stream`` bounded by ``_OUTPUT_TAIL_BYTES``.""" @@ -56,20 +46,6 @@ def _tail(stream: Optional[str]) -> str: return stream[-_OUTPUT_TAIL_BYTES:] -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 - - class CodeExecutorTool(Tool): """Code Executor Run Python (or other) code in a sandboxed, semi-persistent session and capture produced files as artifacts. @@ -198,7 +174,7 @@ class CodeExecutorTool(Tool): if not isinstance(code, str) or not code.strip(): return {"status": "error", "error": "code is required."} - capture_artifacts = kwargs.get("capture_artifacts", True) + should_capture = kwargs.get("capture_artifacts", True) ttl = self._coerce_int(kwargs.get("ttl")) timeout = self._resolve_timeout(kwargs.get("timeout")) inputs = kwargs.get("inputs") or [] @@ -216,7 +192,7 @@ class CodeExecutorTool(Tool): return {"status": "error", "error": materialized["error"]} pre_signatures: Dict[str, Tuple[int, Optional[str]]] = {} - if capture_artifacts: + if should_capture: pre_signatures = self._snapshot_signatures(manager, session_id) try: @@ -228,7 +204,7 @@ class CodeExecutorTool(Tool): # Capture even on error/timeout so partial outputs aren't lost; a # capture failure must never mask the run's real status. artifacts: List[Dict[str, Any]] = [] - if capture_artifacts: + if should_capture: try: artifacts = self._capture_artifacts(manager, session_id, pre_signatures) except Exception: @@ -290,133 +266,33 @@ class CodeExecutorTool(Tool): # 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 + _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.""" - signatures: Dict[str, Tuple[int, Optional[str]]] = {} - try: - files = manager.list_files(session_id) - except Exception: - logger.exception("code_executor: pre-exec listing failed") - return signatures - for rel_path in files: - if rel_path.startswith("inputs/"): - continue - try: - data = manager.get_file(session_id, rel_path) - except Exception: - logger.exception("code_executor: pre-exec signature read failed") - continue - signatures[rel_path] = (len(data), hashlib.sha256(data).hexdigest()) - return signatures + return snapshot_signatures(manager, session_id) def _capture_artifacts( self, manager: Any, session_id: str, pre_signatures: Dict[str, Tuple[int, Optional[str]]] ) -> List[Dict[str, Any]]: """Persist each non-input workspace file that is new or whose content changed.""" - try: - post_files = set(manager.list_files(session_id)) - except Exception: - logger.exception("code_executor: post-exec listing failed") - return [] - - candidates = sorted(f for f in post_files if not f.startswith("inputs/")) - storage = StorageCreator.get_storage() - captured: List[Dict[str, Any]] = [] - for rel_path in candidates: - if len(captured) >= self._MAX_CAPTURED_FILES: - logger.warning("code_executor: capture cap reached; remaining files skipped") - break - try: - data = manager.get_file(session_id, rel_path) - except Exception: - logger.exception("code_executor: 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 - ref = self._persist_artifact(storage, rel_path, data, session_id) - if ref is not None: - captured.append(ref) + 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, + produced_by={ + "tool": "code_executor", + "action": "run_code", + "session_id": session_id, + }, + ) if captured: self._last_artifact_id = captured[0]["artifact_id"] return captured - def _persist_artifact( - self, storage: Any, rel_path: str, data: bytes, session_id: str - ) -> Optional[Dict[str, Any]]: - """Store ``data`` and create an artifact row; size/sha256 are computed server-side. - - 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, but a crash between save and commit can still leak. - """ - # The sandbox filename is display-only; the storage key is derived from - # server-controlled values so a hostile name can't redirect the write. - display_name = rel_path.rsplit("/", 1)[-1] - filename = safe_filename(display_name) - size = len(data) - sha256 = hashlib.sha256(data).hexdigest() - mime_type = _infer_mime(filename) - kind = _kind_for_mime(mime_type) - saved_key: Optional[str] = None - try: - with db_session() as conn: - repo = ArtifactsRepository(conn) - artifact = repo.create_artifact( - self.user_id, - kind, - conversation_id=self.conversation_id, - workflow_run_id=self.workflow_run_id, - title=display_name, - mime_type=mime_type, - filename=filename, - storage_path=None, - size=size, - sha256=sha256, - produced_by={ - "tool": "code_executor", - "action": "run_code", - "session_id": session_id, - }, - ) - artifact_id = str(artifact["id"]) - # ``inputs/{user}/artifacts/...`` is the project storage-namespace - # convention (matches attachments + spec §4); the ``inputs/`` prefix - # is the user's namespace root, not an "input file" marker. - storage_path = f"inputs/{self.user_id}/artifacts/{artifact_id}/v1/{filename}" - # Set the server-derived key on version 1, then write the bytes as - # the LAST statement so a save failure rolls the whole row back. - conn.execute( - text( - "UPDATE artifact_versions SET storage_path = :p " - "WHERE artifact_id = CAST(:aid AS uuid) AND version = 1" - ), - {"p": storage_path, "aid": artifact_id}, - ) - storage.save_file(io.BytesIO(data), storage_path) - saved_key = storage_path - except Exception: - logger.exception("code_executor: failed to persist artifact") - # The bytes landed but the commit failed: drop the now-orphaned key. - if saved_key is not None: - try: - storage.delete_file(saved_key) - except Exception: - logger.exception("code_executor: orphaned-key cleanup failed for %s", saved_key) - return None - return { - "artifact_id": artifact_id, - "version": 1, - "filename": filename, - "mime_type": mime_type, - "size": size, - } - def _shape_payload( self, result: ExecResult, artifacts: List[Dict[str, Any]], inputs_loaded: List[str] ) -> Dict[str, Any]: diff --git a/application/agents/workflow_agent.py b/application/agents/workflow_agent.py index ea780975..2920b5b0 100644 --- a/application/agents/workflow_agent.py +++ b/application/agents/workflow_agent.py @@ -214,6 +214,9 @@ class WorkflowAgent(BaseAgent): str(workflow_row["id"]), owner_id, run.status.value, + # Persist under the engine's run id so any run-scoped + # artifacts produced by code nodes resolve their parent. + run_id=self._engine.workflow_run_id, inputs=run.inputs, result=run.outputs, steps=[step.model_dump(mode="json") for step in run.steps], diff --git a/application/agents/workflows/schemas.py b/application/agents/workflows/schemas.py index 29ced1be..d5dda148 100644 --- a/application/agents/workflows/schemas.py +++ b/application/agents/workflows/schemas.py @@ -12,6 +12,7 @@ class NodeType(str, Enum): NOTE = "note" STATE = "state" CONDITION = "condition" + CODE = "code" class AgentType(str, Enum): @@ -50,6 +51,17 @@ class AgentNodeConfig(BaseModel): json_schema: Optional[Dict[str, Any]] = None +class CodeNodeConfig(BaseModel): + model_config = ConfigDict(extra="allow") + code: str = "" + language: str = "python" + inputs: List[str] = Field(default_factory=list) + libraries: 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): model_config = ConfigDict(extra="forbid", populate_by_name=True) name: Optional[str] = None diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index cf06bb7e..dd07c71d 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -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,29 @@ 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_-]+") 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()) self.state: WorkflowState = {} self.execution_log: List[Dict[str, Any]] = [] self._condition_result: Optional[str] = None @@ -170,6 +187,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, @@ -314,6 +332,166 @@ 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.') + + 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() + manager.open(session_id) + try: + loaded = self._materialize_code_inputs(manager, session_id, config.inputs, user_id) + pre_signatures = snapshot_signatures(manager, session_id) + result = manager.exec(session_id, code, timeout=timeout) + 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}, + ) + finally: + try: + manager.close(session_id) + except Exception: + logger.exception("Code node failed to close sandbox session") + + 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.storage.db.base_repository import looks_like_uuid + 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] = [] + artifact_ids = self._resolve_input_artifact_ids(inputs) + if not artifact_ids: + return loaded + storage = StorageCreator.get_storage() + for artifact_id in artifact_ids: + if not looks_like_uuid(artifact_id): + raise ValueError(f"input artifact {artifact_id} not found in this run.") + with db_readonly() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact_in_parent( + artifact_id, workflow_run_id=self.workflow_run_id + ) + if artifact is None: + raise ValueError(f"input artifact {artifact_id} 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.") + filename = safe_filename(version.get("filename") or artifact_id) + data = storage.get_file(version["storage_path"]).read() + manager.put_file(session_id, f"inputs/{filename}", data) + loaded.append(f"inputs/{filename}") + return loaded + + def _resolve_input_artifact_ids(self, inputs: List[str]) -> List[str]: + """Resolve node ``inputs`` (state-var names holding refs, or raw ids) to artifact ids.""" + resolved: List[str] = [] + for raw in inputs or []: + 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(raw, str) and raw.strip(): + resolved.append(raw.strip()) + return resolved + + 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 _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]: @@ -425,6 +603,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 +628,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: diff --git a/application/api/user/workflows/routes.py b/application/api/user/workflows/routes.py index 8d7bda40..47a9798a 100644 --- a/application/api/user/workflows/routes.py +++ b/application/api/user/workflows/routes.py @@ -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") diff --git a/application/sandbox/artifacts_capture.py b/application/sandbox/artifacts_capture.py new file mode 100644 index 00000000..e779278a --- /dev/null +++ b/application/sandbox/artifacts_capture.py @@ -0,0 +1,294 @@ +"""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 hashlib +import io +import logging +import mimetypes +from typing import Any, Dict, List, Optional, Tuple + +from sqlalchemy import text + +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__) + +# 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 + +_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 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 snapshot_signatures(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.""" + 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 + for rel_path in files: + if rel_path.startswith("inputs/"): + continue + 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, + produced_by: Optional[Dict[str, Any]] = None, +) -> List[Dict[str, Any]]: + """Persist each non-input workspace file that is new or whose content changed. + + 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 [] + + candidates = sorted(f for f in post_files if not f.startswith("inputs/")) + captured: List[Dict[str, Any]] = [] + for rel_path in candidates: + if len(captured) >= MAX_CAPTURED_FILES: + logger.warning("artifacts_capture: capture cap reached; remaining files skipped") + break + 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 + ref = persist_artifact( + rel_path, + data, + user_id=user_id, + conversation_id=conversation_id, + workflow_run_id=workflow_run_id, + produced_by=produced_by, + ) + 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, + 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, + 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 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, + 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 + try: + with db_session() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.create_artifact( + user_id, + kind, + conversation_id=conversation_id, + workflow_run_id=workflow_run_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) + storage.save_file(io.BytesIO(data), storage_path) + saved_key = storage_path + except Exception: + logger.exception("artifacts_capture: failed to persist new artifact") + _cleanup_orphan(storage, saved_key) + return None + return { + "artifact_id": artifact_id, + "version": 1, + "filename": safe_name, + "mime_type": mime_type, + "size": size, + } + + +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, +) -> 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 + try: + with db_session() as conn: + repo = ArtifactsRepository(conn) + 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) + storage.save_file(io.BytesIO(data), storage_path) + saved_key = storage_path + except Exception: + logger.exception("artifacts_capture: failed to append artifact version") + _cleanup_orphan(storage, saved_key) + return None + return { + "artifact_id": str(artifact_id), + "version": version_number, + "filename": safe_name, + "mime_type": mime_type, + "size": size, + } diff --git a/application/storage/db/repositories/workflow_runs.py b/application/storage/db/repositories/workflow_runs.py index 964a9e09..59389981 100644 --- a/application/storage/db/repositories/workflow_runs.py +++ b/application/storage/db/repositories/workflow_runs.py @@ -25,6 +25,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 +38,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: diff --git a/application/templates/namespaces.py b/application/templates/namespaces.py index c2e6f8db..9f1992db 100644 --- a/application/templates/namespaces.py +++ b/application/templates/namespaces.py @@ -151,6 +151,108 @@ class ToolsNamespace(NamespaceBuilder): 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..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..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 NamespaceManager: """Manages all namespace builders and context assembly""" @@ -160,6 +262,7 @@ class NamespaceManager: "passthrough": PassthroughNamespace(), "source": SourceNamespace(), "tools": ToolsNamespace(), + "artifacts": ArtifactsNamespace(), } def build_context(self, **kwargs) -> Dict[str, Any]: diff --git a/tests/agents/test_workflow_agent_graph.py b/tests/agents/test_workflow_agent_graph.py index 8d714fbe..72244518 100644 --- a/tests/agents/test_workflow_agent_graph.py +++ b/tests/agents/test_workflow_agent_graph.py @@ -216,6 +216,8 @@ class TestSaveWorkflowRun: WorkflowRunsRepository, ) + import uuid as _uuid + user = "u-saverun" wf = WorkflowsRepository(pg_conn).create(user, "wf") agent = _make_agent( @@ -225,12 +227,17 @@ class TestSaveWorkflowRun: agent._engine.execution_log = [] agent._engine.state = {"output": "hello"} agent._engine.get_execution_summary.return_value = [] + # The run row is persisted under the engine's run id (the parent for any + # run-scoped artifacts); a MagicMock id can't be adapted by the driver. + run_id = str(_uuid.uuid4()) + agent._engine.workflow_run_id = run_id with _patch_db(pg_conn): agent._save_workflow_run("my query") runs = WorkflowRunsRepository(pg_conn).list_for_workflow(str(wf["id"])) assert len(runs) >= 1 + assert str(runs[0]["id"]) == run_id def test_exception_is_swallowed(self): agent = _make_agent(workflow_id="x", workflow_owner="u") diff --git a/tests/agents/test_workflow_code_node.py b/tests/agents/test_workflow_code_node.py new file mode 100644 index 00000000..64b5e51f --- /dev/null +++ b/tests/agents/test_workflow_code_node.py @@ -0,0 +1,262 @@ +"""Unit tests for the workflow ``code`` node and the pass-by-reference convention. + +These exercise the engine's code-node logic with a fake sandbox manager and a +patched persistence helper (no live gateway / DB / storage), plus the +serialization round-trip and CEL branching on an artifact reference. +""" + +from types import SimpleNamespace + +import pytest + +from application.agents.workflow_agent import WorkflowAgent +from application.agents.workflows.cel_evaluator import evaluate_cel +from application.agents.workflows.schemas import ( + NodeType, + Workflow, + WorkflowGraph, + WorkflowNode, +) +from application.agents.workflows.workflow_engine import WorkflowEngine + + +def _engine() -> WorkflowEngine: + graph = WorkflowGraph(workflow=Workflow(name="Code Node Test"), nodes=[], edges=[]) + agent = SimpleNamespace( + endpoint="stream", + llm_name="openai", + model_id="gpt-4o-mini", + api_key="test-key", + chat_history=[], + user="user-code", + decoded_token={"sub": "user-code"}, + ) + return WorkflowEngine(graph, agent, workflow_run_id="11111111-1111-1111-1111-111111111111") + + +def _code_node(node_id="code_1", **config) -> WorkflowNode: + base = {"code": "print('hi')"} + base.update(config) + return WorkflowNode( + id=node_id, + workflow_id="wf-1", + type=NodeType.CODE, + title="Code", + position={"x": 0, "y": 0}, + config=base, + ) + + +class _Result: + def __init__(self, ok=True, stdout="", error_name=None, error_value=None): + self.status = "ok" if ok else "error" + self.stdout = stdout + self.stderr = "" + self.error_name = error_name + self.error_value = error_value + + @property + def ok(self): + return self.status == "ok" + + +class _FakeManager: + """Records open/close and returns a fixed exec result; no real sandbox.""" + + def __init__(self, result): + self._result = result + self.opened = [] + self.closed = [] + + def open(self, session_id, ttl=None): + self.opened.append(session_id) + return session_id + + def exec(self, session_id, code, timeout=None): + self.last_timeout = timeout + return self._result + + def close(self, session_id): + self.closed.append(session_id) + + +@pytest.fixture() +def patch_sandbox(monkeypatch): + """Patch the sandbox manager + capture helper; return knobs to drive them.""" + state = {"result": _Result(ok=True, stdout="ok"), "captured": [], "snapshot_calls": 0} + + manager_holder = {} + + def _get_manager(): + manager = _FakeManager(state["result"]) + manager_holder["manager"] = manager + return manager + + monkeypatch.setattr( + "application.sandbox.sandbox_creator.SandboxCreator.get_manager", _get_manager + ) + monkeypatch.setattr( + "application.sandbox.artifacts_capture.snapshot_signatures", + lambda *a, **k: state.__setitem__("snapshot_calls", state["snapshot_calls"] + 1) or {}, + ) + monkeypatch.setattr( + "application.sandbox.artifacts_capture.capture_artifacts", + lambda *a, **k: list(state["captured"]), + ) + state["manager_holder"] = manager_holder + return state + + +def test_code_node_writes_artifact_reference_into_state(patch_sandbox): + engine = _engine() + ref = { + "artifact_id": "art-1", + "version": 1, + "filename": "report.pdf", + "mime_type": "application/pdf", + "size": 10, + } + patch_sandbox["captured"] = [ref] + node = _code_node(output_variable="report", code="open('report.pdf','wb').write(b'x')") + + list(engine._execute_code_node(node)) + + # The reference (JSON primitives only) lands under both keys; no bytes. + assert engine.state["node_code_1_output"] == ref + assert engine.state["report"] == ref + assert all(not isinstance(v, (bytes, bytearray)) for v in engine.state["report"].values()) + # The sandbox session is bound to the run id and closed after the run. + manager = patch_sandbox["manager_holder"]["manager"] + assert manager.opened == ["11111111-1111-1111-1111-111111111111"] + assert manager.closed == ["11111111-1111-1111-1111-111111111111"] + + +def test_code_node_no_artifacts_still_writes_status(patch_sandbox): + engine = _engine() + patch_sandbox["captured"] = [] + node = _code_node(output_variable="out", code="print('noop')") + + list(engine._execute_code_node(node)) + + assert engine.state["out"] == {"artifacts": [], "status": "ok"} + + +def test_code_node_failure_raises(patch_sandbox): + engine = _engine() + patch_sandbox["result"] = _Result(ok=False, error_name="ValueError", error_value="boom") + node = _code_node(code="raise ValueError('boom')") + + with pytest.raises(ValueError, match="failed: ValueError: boom"): + list(engine._execute_code_node(node)) + + +def test_code_node_empty_code_raises(patch_sandbox): + engine = _engine() + node = _code_node(code=" ") + with pytest.raises(ValueError, match="no code to execute"): + list(engine._execute_code_node(node)) + + +def test_code_node_json_schema_validates_decision(patch_sandbox): + engine = _engine() + patch_sandbox["result"] = _Result(ok=True, stdout='{"pass": true, "score": 5}') + patch_sandbox["captured"] = [] + node = _code_node( + output_variable="decision", + code="print('{...}')", + json_schema={ + "type": "object", + "properties": {"pass": {"type": "boolean"}, "score": {"type": "integer"}}, + "required": ["pass"], + }, + ) + + list(engine._execute_code_node(node)) + + assert engine.state["decision"] == {"pass": True, "score": 5} + + +def test_code_node_json_schema_rejects_non_json_stdout(patch_sandbox): + engine = _engine() + patch_sandbox["result"] = _Result(ok=True, stdout="not json") + node = _code_node( + json_schema={"type": "object"}, + code="print('x')", + ) + with pytest.raises(ValueError, match="stdout was not valid JSON"): + list(engine._execute_code_node(node)) + + +def test_code_node_json_schema_merges_artifact_reference(patch_sandbox): + engine = _engine() + ref = {"artifact_id": "art-9", "version": 1, "filename": "r.pdf", "mime_type": "application/pdf", "size": 3} + patch_sandbox["result"] = _Result(ok=True, stdout='{"pass": false}') + patch_sandbox["captured"] = [ref] + node = _code_node( + output_variable="decision", + json_schema={"type": "object", "properties": {"pass": {"type": "boolean"}}}, + ) + + list(engine._execute_code_node(node)) + + decision = engine.state["decision"] + assert decision["pass"] is False + assert decision["artifacts"] == [ref] + + +def test_code_node_timeout_clamped_to_cap(patch_sandbox, monkeypatch): + from application.core.settings import settings + + monkeypatch.setattr(settings, "SANDBOX_EXEC_TIMEOUT", 30, raising=False) + engine = _engine() + node = _code_node(timeout=9999, code="print('x')") + list(engine._execute_code_node(node)) + manager = patch_sandbox["manager_holder"]["manager"] + assert manager.last_timeout == 30.0 + + +def test_resolve_input_artifact_ids_from_state_refs_and_raw(): + engine = _engine() + engine.state["report"] = {"artifact_id": "art-from-ref"} + engine.state["not_a_ref"] = "plain string" + ids = engine._resolve_input_artifact_ids(["report", "art-raw-id", "not_a_ref"]) + # A state var holding a ref resolves to its artifact_id; any other entry is + # taken as a raw artifact id (the literal token, not a resolved value). + assert ids == ["art-from-ref", "art-raw-id", "not_a_ref"] + + +# --------------------------------------------------------------------------- +# Pass-by-reference: survives serialization + CEL branches on the metadata. +# --------------------------------------------------------------------------- + + +def test_artifact_reference_survives_serialize_state_value(): + agent = WorkflowAgent.__new__(WorkflowAgent) + ref = { + "artifact_id": "art-1", + "version": 2, + "filename": "report.pdf", + "mime_type": "application/pdf", + "size": 1234, + } + state = {"report": ref, "decision": {"pass": True, "artifacts": [ref]}} + + serialized = agent._serialize_state(state) + + # Every primitive survives untouched (no stringification of the dict ref). + assert serialized["report"] == ref + assert serialized["decision"]["pass"] is True + assert serialized["decision"]["artifacts"][0] == ref + + +def test_cel_branches_on_artifact_metadata(): + state = { + "report": { + "artifact_id": "art-1", + "size": 1234, + "mime_type": "application/pdf", + } + } + assert evaluate_cel('report.size > 0 && report.mime_type == "application/pdf"', state) is True + assert evaluate_cel('report.mime_type == "text/plain"', state) is False + assert evaluate_cel("report.size > 5000", state) is False diff --git a/tests/agents/test_workflow_schemas.py b/tests/agents/test_workflow_schemas.py index e98038ef..a4dd96e2 100644 --- a/tests/agents/test_workflow_schemas.py +++ b/tests/agents/test_workflow_schemas.py @@ -38,6 +38,7 @@ class TestNodeType: assert NodeType.NOTE == "note" assert NodeType.STATE == "state" assert NodeType.CONDITION == "condition" + assert NodeType.CODE == "code" @pytest.mark.unit def test_all_members(self): @@ -48,6 +49,7 @@ class TestNodeType: NodeType.NOTE, NodeType.STATE, NodeType.CONDITION, + NodeType.CODE, } diff --git a/tests/integration/test_workflow_code_node_e2e.py b/tests/integration/test_workflow_code_node_e2e.py new file mode 100644 index 00000000..ae7075b2 --- /dev/null +++ b/tests/integration/test_workflow_code_node_e2e.py @@ -0,0 +1,278 @@ +"""End-to-end workflow ``code`` node: live Jupyter gateway + live Postgres + local storage. + +Launches a real ``jupyter kernelgateway``, wires the sandbox manager + a temp-dir +``LocalStorage`` + the integration Postgres engine, and drives a ``WorkflowEngine`` +through a ``code`` node: the node writes a file -> an artifact REFERENCE lands in +``state`` (artifact row persisted, run-scoped; the ref is ``{artifact_id,...}`` with +NO bytes); the ref survives the ``workflow_runs`` state-snapshot serialization; a +downstream node re-reads it by reference; the ``artifacts.*`` namespace and CEL both +resolve its metadata. + +Skips gracefully when the gateway binary / websocket-client / POSTGRES_URI is absent. +""" + +from __future__ import annotations + +import shutil +import socket +import subprocess +import time +import uuid +from types import SimpleNamespace + +import pytest + +requests = pytest.importorskip("requests") +pytest.importorskip("websocket") # websocket-client + +from application.agents.workflow_agent import WorkflowAgent # noqa: E402 +from application.agents.workflows.cel_evaluator import evaluate_cel # noqa: E402 +from application.agents.workflows.schemas import ( # noqa: E402 + NodeType, + Workflow, + WorkflowGraph, + WorkflowNode, +) +from application.agents.workflows.workflow_engine import WorkflowEngine # noqa: E402 +from application.sandbox.jupyter_gateway import JupyterKernelGatewaySandbox # noqa: E402 +from application.sandbox.manager import SandboxManager # noqa: E402 +from application.sandbox.sandbox_creator import SandboxCreator # noqa: E402 +from application.storage.db.repositories.artifacts import ArtifactsRepository # noqa: E402 +from application.storage.local import LocalStorage # noqa: E402 +from application.storage.storage_creator import StorageCreator # noqa: E402 +from application.templates.namespaces import NamespaceManager # noqa: E402 + +_GATEWAY_BIN = shutil.which("jupyter-kernelgateway") or shutil.which("jupyter") + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + _GATEWAY_BIN is None, + reason="jupyter kernel gateway not installed (pip install jupyter-kernel-gateway)", + ), +] + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _gateway_cmd(port: int) -> list: + if _GATEWAY_BIN.endswith("jupyter-kernelgateway"): + base = [_GATEWAY_BIN] + else: + base = [_GATEWAY_BIN, "kernelgateway"] + return base + [ + "--KernelGatewayApp.ip=127.0.0.1", + f"--KernelGatewayApp.port={port}", + "--ZMQChannelsWebsocketConnection.limit_rate=False", + ] + + +@pytest.fixture(scope="module") +def gateway_url(): + port = _free_port() + proc = subprocess.Popen(_gateway_cmd(port), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + url = f"http://127.0.0.1:{port}" + deadline = time.time() + 30 + ready = False + try: + while time.time() < deadline: + if proc.poll() is not None: + pytest.skip("jupyter kernelgateway process exited during startup") + try: + if requests.get(f"{url}/api", timeout=1).status_code == 200: + ready = True + break + except requests.RequestException: + time.sleep(0.3) + if not ready: + pytest.skip("jupyter kernelgateway did not become ready in time") + yield url + finally: + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +@pytest.fixture() +def wired(gateway_url, pg_engine, tmp_path, monkeypatch): + """Sandbox manager + temp local storage + integration PG, with a fresh run id.""" + backend = JupyterKernelGatewaySandbox(gateway_url=gateway_url, default_timeout=30.0) + SandboxCreator._instance = SandboxManager(backend=backend, max_ttl=1200.0) + + storage = LocalStorage(base_dir=str(tmp_path)) + monkeypatch.setattr(StorageCreator, "_instance", storage, raising=False) + monkeypatch.setattr("application.storage.db.session.get_engine", lambda: pg_engine) + + run_id = str(uuid.uuid4()) + try: + yield run_id, pg_engine, storage + finally: + SandboxCreator.reset() + + +def _engine(run_id: str) -> WorkflowEngine: + graph = WorkflowGraph(workflow=Workflow(name="Code Node E2E"), nodes=[], edges=[]) + agent = SimpleNamespace( + endpoint="stream", + llm_name="openai", + model_id="gpt-4o-mini", + api_key="k", + chat_history=[], + user="user-wf-e2e", + decoded_token={"sub": "user-wf-e2e"}, + request_id="req-1", + retrieved_docs=None, + ) + return WorkflowEngine(graph, agent, workflow_run_id=run_id) + + +def _code_node(node_id, code, **config): + cfg = {"code": code} + cfg.update(config) + return WorkflowNode( + id=node_id, workflow_id="wf-e2e", type=NodeType.CODE, title="Code", + position={"x": 0, "y": 0}, config=cfg, + ) + + +def test_code_node_persists_artifact_reference_in_state(wired): + run_id, pg_engine, storage = wired + engine = _engine(run_id) + + node = _code_node( + "code_1", + "with open('report.txt', 'w') as f:\n f.write('compliance ok')\n", + output_variable="report", + ) + list(engine._execute_code_node(node)) + + ref = engine.state["report"] + # The state holds an artifact REFERENCE (JSON primitives), never bytes. + assert set(ref) == {"artifact_id", "version", "filename", "mime_type", "size"} + assert ref["filename"] == "report.txt" + assert ref["mime_type"] == "text/plain" + assert ref["size"] == len(b"compliance ok") + assert all(not isinstance(v, (bytes, bytearray)) for v in ref.values()) + assert engine.state["node_code_1_output"] == ref + + # The artifact row is persisted, parent-scoped to THIS run (not cross-tenant). + with pg_engine.connect() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact_in_parent(ref["artifact_id"], workflow_run_id=run_id) + assert artifact is not None + assert repo.get_artifact_in_parent(ref["artifact_id"], workflow_run_id=str(uuid.uuid4())) is None + version = repo.get_version(ref["artifact_id"], 1) + assert version["size"] == len(b"compliance ok") + assert version["produced_by"]["node_id"] == "code_1" + # The bytes live in storage, never in state. + assert storage.get_file(version["storage_path"]).read() == b"compliance ok" + + +def test_reference_survives_serialization_and_cel_branches(wired): + run_id, _pg_engine, _storage = wired + engine = _engine(run_id) + + node = _code_node( + "code_meta", + "open('out.json', 'w').write('{\"k\": 1}')\n", + output_variable="report", + ) + list(engine._execute_code_node(node)) + ref = engine.state["report"] + + # Survives the workflow_runs state-snapshot serialization unchanged. + agent = WorkflowAgent.__new__(WorkflowAgent) + serialized = agent._serialize_state(engine.state) + assert serialized["report"] == ref + assert isinstance(serialized["report"], dict) + + # CEL branches on the reference's metadata (nested-dict indexing). + assert evaluate_cel("report.size > 0", engine.state) is True + assert evaluate_cel('report.mime_type == "application/json"', engine.state) is True + assert evaluate_cel('report.mime_type == "application/pdf"', engine.state) is False + + +def test_condition_node_branches_on_artifact_metadata(wired): + run_id, _pg_engine, _storage = wired + engine = _engine(run_id) + + list(engine._execute_code_node( + _code_node("code_c", "open('r.json','w').write('{}')", output_variable="report") + )) + + cond = WorkflowNode( + id="cond_1", workflow_id="wf-e2e", type=NodeType.CONDITION, title="Branch", + position={"x": 0, "y": 0}, + config={ + "mode": "simple", + "cases": [ + {"name": "has json", "expression": 'report.mime_type == "application/json"', + "sourceHandle": "json_case"}, + ], + }, + ) + list(engine._execute_condition_node(cond)) + assert engine._condition_result == "json_case" + + +def test_downstream_node_reads_reference_via_inputs(wired): + run_id, _pg_engine, _storage = wired + engine = _engine(run_id) + + list(engine._execute_code_node( + _code_node("producer", "open('seed.txt','w').write('seed-bytes')", output_variable="seed") + )) + + # A downstream code node references the upstream ref by its state-var name; + # the engine re-fetches the bytes by id and stages them into the workspace. + consumer = _code_node( + "consumer", + "data = open('inputs/seed.txt', 'rb').read()\n" + "open('combined.txt', 'wb').write(data + b'-processed')\n", + inputs=["seed"], + output_variable="combined", + ) + list(engine._execute_code_node(consumer)) + + out = engine.state["combined"] + assert out["filename"] == "combined.txt" + assert out["size"] == len(b"seed-bytes-processed") + assert out["artifact_id"] != engine.state["seed"]["artifact_id"] + + +def test_artifacts_namespace_resolves_reference(wired): + run_id, _pg_engine, _storage = wired + engine = _engine(run_id) + + list(engine._execute_code_node( + _code_node("ns_node", "open('deck.txt','w').write('hello')", output_variable="report") + )) + ref = engine.state["report"] + + # The artifacts.* namespace (shared with the prompt renderer) exposes the + # reference's metadata by output-variable name, and artifact(id) resolves + # parent-scoped metadata — never bytes. + ctx = NamespaceManager().build_context( + artifacts_data={"report": ref}, + artifact_parent={"workflow_run_id": run_id}, + ) + assert ctx["artifacts"]["report"]["id"] == ref["artifact_id"] + assert ctx["artifacts"]["report"]["mime_type"] == "text/plain" + assert ctx["artifacts"]["report"]["filename"] == "deck.txt" + + looked_up = ctx["artifacts"]["artifact"](ref["artifact_id"]) + assert looked_up["id"] == ref["artifact_id"] + assert looked_up["filename"] == "deck.txt" + assert all(not isinstance(v, (bytes, bytearray)) for v in looked_up.values()) + + # A foreign run id never resolves this artifact (no cross-tenant leak). + foreign = NamespaceManager().build_context( + artifact_parent={"workflow_run_id": str(uuid.uuid4())}, + ) + assert foreign["artifacts"]["artifact"](ref["artifact_id"]) == {} diff --git a/tests/test_namespaces.py b/tests/test_namespaces.py index 0f849264..d0c5486a 100644 --- a/tests/test_namespaces.py +++ b/tests/test_namespaces.py @@ -4,6 +4,7 @@ from unittest.mock import patch import pytest from application.templates.namespaces import ( + ArtifactsNamespace, NamespaceBuilder, NamespaceManager, PassthroughNamespace, @@ -188,6 +189,141 @@ class TestToolsNamespace: assert result == {"good": "ok"} +# ── ArtifactsNamespace ───────────────────────────────────────────────────────── + + +def _ref(**overrides): + base = { + "artifact_id": "a-1", + "version": 2, + "mime_type": "application/pdf", + "filename": "report.pdf", + "size": 1234, + } + base.update(overrides) + return base + + +@pytest.mark.unit +class TestArtifactsNamespace: + def test_namespace_name(self): + assert ArtifactsNamespace().namespace_name == "artifacts" + + def test_named_ref_exposes_id_mime_filename(self): + ns = ArtifactsNamespace() + ctx = ns.build(artifacts_data={"report": _ref()}) + assert ctx["report"]["id"] == "a-1" + assert ctx["report"]["artifact_id"] == "a-1" + assert ctx["report"]["mime_type"] == "application/pdf" + assert ctx["report"]["filename"] == "report.pdf" + assert ctx["report"]["version"] == 2 + assert ctx["report"]["size"] == 1234 + + def test_non_artifact_values_dropped(self): + ns = ArtifactsNamespace() + ctx = ns.build( + artifacts_data={ + "good": _ref(), + "plain": "just a string", + "missing_id": {"mime_type": "text/plain"}, + "none": None, + } + ) + assert "good" in ctx + assert "plain" not in ctx + assert "missing_id" not in ctx + assert "none" not in ctx + + def test_bytes_never_exposed(self): + ns = ArtifactsNamespace() + ctx = ns.build( + artifacts_data={"report": _ref(content=b"\x00\x01binary", raw=bytearray(b"x"))} + ) + view = ctx["report"] + assert "content" not in view + assert "raw" not in view + for value in view.values(): + assert not isinstance(value, (bytes, bytearray)) + + def test_no_data_only_exposes_lookup(self): + ns = ArtifactsNamespace() + ctx = ns.build() + assert callable(ctx["artifact"]) + assert [k for k in ctx if k != "artifact"] == [] + + def test_artifact_lookup_without_parent_returns_empty(self): + ns = ArtifactsNamespace() + artifact = ns.build()["artifact"] + assert artifact("a-1") == {} + + def test_artifact_lookup_resolves_parent_scoped_metadata(self): + ns = ArtifactsNamespace() + + class _Repo: + def __init__(self, conn): + self.conn = conn + + def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None): + assert workflow_run_id == "run-9" + assert conversation_id is None + return {"id": artifact_id, "current_version": 3, "kind": "document", "title": "Deck"} + + def get_version(self, artifact_id, version): + assert version == 3 + return {"mime_type": "application/pdf", "filename": "deck.pdf", "size": 42, "spec": {"x": 1}} + + from contextlib import contextmanager + + @contextmanager + def _fake_readonly(): + yield object() + + with patch( + "application.storage.db.repositories.artifacts.ArtifactsRepository", _Repo + ), patch("application.storage.db.session.db_readonly", _fake_readonly): + artifact = ns.build(artifact_parent={"workflow_run_id": "run-9"})["artifact"] + meta = artifact("art-7") + + assert meta["id"] == "art-7" + assert meta["version"] == 3 + assert meta["mime_type"] == "application/pdf" + assert meta["filename"] == "deck.pdf" + assert meta["size"] == 42 + # The version's spec (a nested dict, potentially large) is never exposed. + assert "spec" not in meta + + def test_artifact_lookup_cross_tenant_returns_empty(self): + ns = ArtifactsNamespace() + + class _Repo: + def __init__(self, conn): + pass + + def get_artifact_in_parent(self, *a, **k): + return None + + def get_version(self, *a, **k): + raise AssertionError("must not fetch a version for a foreign artifact") + + from contextlib import contextmanager + + @contextmanager + def _fake_readonly(): + yield object() + + with patch( + "application.storage.db.repositories.artifacts.ArtifactsRepository", _Repo + ), patch("application.storage.db.session.db_readonly", _fake_readonly): + artifact = ns.build(artifact_parent={"workflow_run_id": "run-9"})["artifact"] + assert artifact("foreign") == {} + + def test_manager_includes_artifacts_namespace(self): + mgr = NamespaceManager() + ctx = mgr.build_context(artifacts_data={"report": _ref()}) + assert ctx["artifacts"]["report"]["id"] == "a-1" + assert callable(ctx["artifacts"]["artifact"]) + + # ── NamespaceBuilder ABC ──────────────────────────────────────────────────────