DocsGPT/tests/agents/test_workflow_code_node.py

"""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.
"""

import json
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/put_file/exec and returns a fixed result; no real sandbox."""

    def __init__(self, result):
        self._result = result
        self.opened = []
        self.closed = []
        self.put_files = []
        self.last_code = None

    def open(self, session_id, ttl=None):
        self.opened.append(session_id)
        return session_id

    def put_file(self, session_id, dest_path, data):
        self.put_files.append((dest_path, data))

    def exec(self, session_id, code, timeout=None):
        self.last_timeout = timeout
        self.last_code = code
        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_reads_prior_state_from_state_json(patch_sandbox):
    # Prior state is staged as DATA in state.json (workspace root = kernel cwd) so
    # node code reads it with json.load(open("state.json")) -- e.g. state["decision"].
    engine = _engine()
    engine.state["decision"] = {"pass": True, "score": 7}
    node = _code_node(
        output_variable="out",
        code="import json\nd = json.load(open('state.json'))\nprint(d['decision'])\n",
    )

    list(engine._execute_code_node(node))

    manager = patch_sandbox["manager_holder"]["manager"]
    staged = dict(manager.put_files)
    assert "state.json" in staged
    payload = json.loads(staged["state.json"].decode("utf-8"))
    assert payload["decision"] == {"pass": True, "score": 7}


def test_code_node_literal_braces_passed_verbatim_not_templated(patch_sandbox):
    # Proves code nodes are NOT Jinja-rendered: a literal ``{{ ... }}`` in the code
    # reaches exec() byte-for-byte (no injection path that interpolates state).
    engine = _engine()
    engine.state["decision"] = "INJECTED"
    literal = "x = '{{ agent.decision }}'\nprint(x)\n"
    node = _code_node(output_variable="out", code=literal)

    list(engine._execute_code_node(node))

    manager = patch_sandbox["manager_holder"]["manager"]
    assert manager.last_code == literal
    assert "INJECTED" not in manager.last_code


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