DocsGPT/tests/agents/test_workflow_run_scoped_artifacts.py

"""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 docsgpt.agents.tool_executor import ToolExecutor
from docsgpt.agents.workflows.node_agent import WorkflowNodeAgentFactory
from docsgpt.agents.workflows.schemas import (
    NodeType,
    Workflow,
    WorkflowGraph,
    WorkflowNode,
)
from docsgpt.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(
        "docsgpt.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


def test_agent_node_skips_run_scope_when_not_persisted(monkeypatch):
    """An unpersisted run must NOT stamp workflow_run_id (it would orphan artifacts)."""
    engine = _create_engine(workflow_run_id="22222222-2222-2222-2222-222222222222")
    engine.run_persisted = False
    node = _agent_node()
    stub = _StubNodeAgent([{"answer": "done"}])

    monkeypatch.setattr(
        WorkflowNodeAgentFactory, "create", staticmethod(lambda **kwargs: stub)
    )
    monkeypatch.setattr(
        "docsgpt.core.model_utils.get_api_key_for_provider", lambda _provider: None
    )

    list(engine._execute_agent_node(node))

    # Left unset: the run-scoped tools then persist under a conversation parent or
    # cleanly error, never orphaning an artifact under a nonexistent run row.
    assert stub.tool_executor.workflow_run_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(
        "docsgpt.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 docsgpt.agents.tools.artifact_generator import ArtifactGeneratorTool
    from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository
    from docsgpt.storage.local import LocalStorage
    from docsgpt.storage.storage_creator import StorageCreator

    storage = LocalStorage(base_dir=str(tmp_path))
    monkeypatch.setattr(StorageCreator, "_instance", storage, raising=False)
    monkeypatch.setattr(
        "docsgpt.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