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 ──────────────────────────────────────────────────────