diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py index 7cc49ef9..d168ba9a 100644 --- a/application/agents/tool_executor.py +++ b/application/agents/tool_executor.py @@ -197,6 +197,10 @@ class ToolExecutor: self.tool_calls: List[Dict] = [] self._loaded_tools: Dict[str, object] = {} self.conversation_id: Optional[str] = None + # Set by the workflow engine for agent nodes so run-scoped tools + # (artifact_generator / code_executor) address artifacts by the + # workflow run rather than a conversation. + self.workflow_run_id: Optional[str] = None self.message_id: Optional[str] = None self.client_tools: Optional[List[Dict]] = None self._name_to_tool: Dict[str, Tuple[str, str]] = {} @@ -945,6 +949,10 @@ class ToolExecutor: tool_config["tool_id"] = str(row_id) if self.conversation_id: tool_config["conversation_id"] = self.conversation_id + # Workflow agent nodes run-scope their artifact tools so a short + # ref (A1) and edit_artifact resolve against the workflow run. + if self.workflow_run_id: + tool_config["workflow_run_id"] = self.workflow_run_id if tool_data["name"] == "scheduler": # Agent-bound: stamp schedules.agent_id. Agentless: the tool # falls back to ``origin_conversation_id`` as the schedule's diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index 78a0b0b9..fb6b8156 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -283,6 +283,11 @@ class WorkflowEngine: } node_agent = WorkflowNodeAgentFactory.create(**factory_kwargs) + # Run-scope the node agent's tools so artifact_generator / code_executor + # address artifacts by this workflow run: a short ref (A1) created by one + # node resolves for edit_artifact in a later node within the same run. + if getattr(node_agent, "tool_executor", None) is not None: + node_agent.tool_executor.workflow_run_id = self.workflow_run_id full_response_parts: List[str] = [] structured_response_parts: List[str] = [] diff --git a/tests/agents/test_workflow_run_scoped_artifacts.py b/tests/agents/test_workflow_run_scoped_artifacts.py new file mode 100644 index 00000000..678758c8 --- /dev/null +++ b/tests/agents/test_workflow_run_scoped_artifacts.py @@ -0,0 +1,209 @@ +"""Regression tests: workflow AGENT nodes run-scope their artifact tools. + +The engine stamps each node agent's ``ToolExecutor.workflow_run_id`` so run-aware +tools (artifact_generator / code_executor) address artifacts by the workflow run. +``ToolExecutor._get_or_load_tool`` then stamps ``workflow_run_id`` into the loaded +tool's ``tool_config`` only when set, so a short ref (A1) created by one node +resolves for ``edit_artifact`` in a later node within the same run. +""" + +from __future__ import annotations + +import uuid +from types import SimpleNamespace +from typing import Any, Dict, Optional +from unittest.mock import Mock + +import pytest + +from application.agents.tool_executor import ToolExecutor +from application.agents.workflows.node_agent import WorkflowNodeAgentFactory +from application.agents.workflows.schemas import ( + NodeType, + Workflow, + WorkflowGraph, + WorkflowNode, +) +from application.agents.workflows.workflow_engine import WorkflowEngine + + +class _StubNodeAgent: + """Node-agent stub carrying a real ToolExecutor and an LLM-free gen().""" + + def __init__(self, events: list) -> None: + self.events = events + self.tool_executor = ToolExecutor(user="user-1") + + def gen(self, _prompt: str): + yield from self.events + + +def _create_engine(workflow_run_id: Optional[str] = None) -> WorkflowEngine: + """Build an engine bound to a bare agent stub (no LLM).""" + graph = WorkflowGraph(workflow=Workflow(name="Run Scope Test"), nodes=[], edges=[]) + agent = SimpleNamespace( + endpoint="stream", + llm_name="openai", + model_id="gpt-4o-mini", + api_key="test-key", + chat_history=[], + decoded_token={"sub": "user-1"}, + ) + return WorkflowEngine(graph, agent, workflow_run_id=workflow_run_id) + + +def _agent_node(node_id: str = "agent_1") -> WorkflowNode: + """Minimal classic AGENT node config.""" + return WorkflowNode( + id=node_id, + workflow_id="workflow-1", + type=NodeType.AGENT, + title="Agent", + position={"x": 0, "y": 0}, + config={ + "agent_type": "classic", + "system_prompt": "You are a helpful assistant.", + "prompt_template": "", + "stream_to_user": False, + "tools": [], + }, + ) + + +# --------------------------------------------------------------------------- +# Engine wiring +# --------------------------------------------------------------------------- + + +def test_agent_node_run_scopes_tool_executor(monkeypatch): + """The node agent's ToolExecutor inherits the engine's workflow_run_id.""" + engine = _create_engine(workflow_run_id="22222222-2222-2222-2222-222222222222") + node = _agent_node() + stub = _StubNodeAgent([{"answer": "done"}]) + + monkeypatch.setattr( + WorkflowNodeAgentFactory, "create", staticmethod(lambda **kwargs: stub) + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", lambda _provider: None + ) + + list(engine._execute_agent_node(node)) + + assert stub.tool_executor.workflow_run_id == engine.workflow_run_id + # Both-parents safety: a workflow node addresses artifacts by run, never by a + # conversation. If a conversation_id ever leaks into the node agent, the + # run-scoped parent gate would have two parents — fail here instead. + assert stub.tool_executor.conversation_id is None + + +# --------------------------------------------------------------------------- +# ToolExecutor stamping +# --------------------------------------------------------------------------- + + +def _capture_tool_config(monkeypatch) -> Dict[str, Any]: + """Patch ToolManager so _get_or_load_tool's tool_config is captured.""" + captured: Dict[str, Any] = {} + mock_tm = Mock() + mock_tm.load_tool.return_value = Mock() + + def _load_tool(name, tool_config, user_id=None): + captured["tool_config"] = tool_config + return mock_tm.load_tool.return_value + + mock_tm.load_tool.side_effect = _load_tool + monkeypatch.setattr( + "application.agents.tool_executor.ToolManager", lambda config: mock_tm + ) + return captured + + +def test_get_or_load_tool_stamps_workflow_run_id_when_set(monkeypatch): + """A run-scoped executor stamps workflow_run_id into a tool's tool_config.""" + captured = _capture_tool_config(monkeypatch) + executor = ToolExecutor(user="user-1") + executor.workflow_run_id = "33333333-3333-3333-3333-333333333333" + tool_data = { + "id": "00000000-0000-0000-0000-0000000000aa", + "name": "artifact_generator", + "config": {}, + } + + executor._get_or_load_tool(tool_data, "t1", "create_artifact") + + assert captured["tool_config"]["workflow_run_id"] == executor.workflow_run_id + # A run-scoped node has no conversation parent. + assert "conversation_id" not in captured["tool_config"] + + +def test_get_or_load_tool_omits_workflow_run_id_for_chat(monkeypatch): + """The chat case (no workflow_run_id) leaves it off the tool_config.""" + captured = _capture_tool_config(monkeypatch) + executor = ToolExecutor(user="user-1") + executor.conversation_id = "conv-1" + tool_data = { + "id": "00000000-0000-0000-0000-0000000000bb", + "name": "artifact_generator", + "config": {}, + } + + executor._get_or_load_tool(tool_data, "t2", "create_artifact") + + assert "workflow_run_id" not in captured["tool_config"] + assert captured["tool_config"]["conversation_id"] == "conv-1" + + +# --------------------------------------------------------------------------- +# Run-scoped resolution against real Postgres + storage (no LLM, no sandbox) +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_run_scoped_ref_resolves_for_edit(pg_engine, tmp_path, monkeypatch): + """A1 created under a workflow_run_id resolves for edit_artifact -> v2.""" + pytest.importorskip("jsonschema") + from application.agents.tools.artifact_generator import ArtifactGeneratorTool + from application.storage.db.repositories.artifacts import ArtifactsRepository + from application.storage.local import LocalStorage + from application.storage.storage_creator import StorageCreator + + 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 + ) + # Skip the Jupyter-gateway renderer: the run-scoping under test is the ref + # resolution + version append, not the rendered bytes. + monkeypatch.setattr( + ArtifactGeneratorTool, "_render", lambda self, kind, spec: {"data": b"%PDF-1.4 stub"} + ) + + workflow_run_id = str(uuid.uuid4()) + tool = ArtifactGeneratorTool( + tool_config={"workflow_run_id": workflow_run_id, "tool_id": str(uuid.uuid4())}, + user_id="user-run", + ) + + created = tool.execute_action( + "create_artifact", kind="presentation", title="Deck", spec={"slides": [{"title": "a"}]} + ) + assert created["status"] == "ok", created + assert created["version"] == 1 + assert created["ref"] == "A1" + artifact_id = created["artifact_id"] + + # edit_artifact by the run-scoped short ref resolves the same artifact -> v2. + edited = tool.execute_action( + "edit_artifact", id="A1", spec_patch={"slides": [{"title": "a"}, {"title": "b"}]} + ) + assert edited["status"] == "ok", edited + assert edited["version"] == 2 + assert edited["artifact_id"] == artifact_id + + with pg_engine.connect() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact_in_parent(artifact_id, workflow_run_id=workflow_run_id) + v2 = repo.get_version(artifact_id, 2) + assert artifact["current_version"] == 2 + assert len(v2["spec"]["slides"]) == 2