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.
This commit is contained in:
Alex committed 2026-06-24 12:30:18 +01:00
1 parent 1fe6236281
commit 1ac1f793b7
13 files changed
+1337 -156

No files matched your search

+30 -154
View File
@@ -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]:
+3
View File
@@ -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],
+12
View File
@@ -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
+192 -2
View File
@@ -1,5 +1,7 @@
import json
import logging
import re
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, Generator, List, Optional, TYPE_CHECKING
@@ -8,6 +10,7 @@ from application.agents.workflows.node_agent import WorkflowNodeAgentFactory
from application.agents.workflows.schemas import (
AgentNodeConfig,
AgentType,
CodeNodeConfig,
ConditionNodeConfig,
ExecutionStatus,
NodeExecutionLog,
@@ -34,15 +37,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:
+13
View File
@@ -357,6 +357,19 @@ def validate_workflow_structure(
if schema_error:
errors.append(f"Agent node '{agent_title}' JSON schema {schema_error}")
code_nodes = [n for n in nodes if n.get("type") == "code"]
for code_node in code_nodes:
code_title = code_node.get("title", code_node.get("id", "unknown"))
raw_config = code_node.get("data", {}) or {}
if not isinstance(raw_config, dict):
errors.append(f"Code node '{code_title}' has invalid configuration")
continue
if not str(raw_config.get("code", "")).strip():
errors.append(f"Code node '{code_title}' must have code to execute")
_, schema_error = validate_json_schema_payload(raw_config.get("json_schema"))
if schema_error:
errors.append(f"Code node '{code_title}' JSON schema {schema_error}")
for node in nodes:
if not node.get("id"):
errors.append("All nodes must have an id")
+294
View File
@@ -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,
}
@@ -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:
+103
View File
@@ -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.<name>.id }}`` is the documented accessor; mirror artifact_id to id.
view["id"] = view.get("artifact_id")
return view
class ArtifactsNamespace(NamespaceBuilder):
"""Artifact references namespace: {{ artifacts.<name>.id }} / .mime_type / .filename + artifact(id)."""
@property
def namespace_name(self) -> str:
return "artifacts"
def build(
self,
artifacts_data: Optional[Dict[str, Any]] = None,
artifact_parent: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Dict[str, Any]:
"""Build the artifacts namespace from named references plus a parent-scoped ``artifact(id)`` lookup.
Args:
artifacts_data: Map of output-variable name -> artifact reference dict.
artifact_parent: Parent scope ({"conversation_id"|"workflow_run_id"}) for ``artifact(id)``.
Returns:
A dict of named metadata views plus an ``artifact`` callable; never any bytes.
"""
context: Dict[str, Any] = {}
for name, ref in (artifacts_data or {}).items():
view = _artifact_view(ref)
if view is not None:
context[str(name)] = view
context["artifact"] = self._make_lookup(artifact_parent or {})
return context
def _make_lookup(self, parent: Dict[str, Any]):
"""Return an ``artifact(id)`` helper that resolves metadata scoped to the run/conversation parent."""
conversation_id = parent.get("conversation_id")
workflow_run_id = parent.get("workflow_run_id")
def artifact(artifact_id: Any) -> Dict[str, Any]:
"""Resolve one artifact's metadata by id, scoped to this parent (never cross-tenant)."""
if not artifact_id or (conversation_id is None and workflow_run_id is None):
return {}
try:
from application.storage.db.repositories.artifacts import (
ArtifactsRepository,
)
from application.storage.db.session import db_readonly
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
row = repo.get_artifact_in_parent(
str(artifact_id),
conversation_id=conversation_id,
workflow_run_id=workflow_run_id,
)
if row is None:
return {}
version = repo.get_version(str(artifact_id), row.get("current_version", 1))
except Exception:
logger.exception("Failed to resolve artifact %s in namespace", artifact_id)
return {}
view = {
"artifact_id": str(row["id"]),
"id": str(row["id"]),
"version": row.get("current_version"),
"kind": row.get("kind"),
"title": row.get("title"),
}
if version is not None:
for key in ("mime_type", "filename", "size"):
value = version.get(key)
if isinstance(value, (str, int, float, bool, type(None))):
view[key] = value
return view
return artifact
class 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]:
@@ -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")
+262
View File
@@ -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
+2
View File
@@ -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,
}
@@ -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"]) == {}
+136
View File
@@ -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 ──────────────────────────────────────────────────────