tests/test_runtime.py

import asyncio

import pytest

from trace2mcp.compiler import compile_trace
from trace2mcp.models import (
    PolicyViolationError,
    Reference,
    Trace,
    TraceEvent,
    Workflow,
    WorkflowNode,
)
from trace2mcp.policy import ToolContract
from trace2mcp.runtime import (
    ToolNotFoundError,
    ToolRegistry,
    WorkflowExecutionError,
    execute_workflow,
)


def read_event(event_id, tool, arguments, result):
    return TraceEvent(event_id, tool, arguments, result, side_effect="read")


def test_runtime_resolves_refs_and_returns_sink_output():
    workflow = compile_trace(
        Trace(
            inputs={"value": 4},
            events=[
                read_event("double", "double", {"number": 4}, {"number": 8}),
                read_event("format", "format", {"number": 8}, "answer=8"),
            ],
        )
    )
    registry = ToolRegistry(
        {
            "double": lambda number: {"number": number * 2},
            "format": lambda number: f"answer={number}",
        }
    )

    result = asyncio.run(execute_workflow(workflow, registry, inputs={"value": 5}))

    assert result.outputs == {"double": {"number": 10}, "format": "answer=10"}
    assert result.output == "answer=10"
    assert result.metrics.nodes_executed == 2
    assert result.metrics.layer_count == 2
    assert all(metric.status == "success" for metric in result.metrics.node_metrics)


def test_independent_nodes_execute_concurrently_in_one_layer():
    workflow = Workflow(
        name="parallel",
        inputs={},
        nodes=[
            WorkflowNode("a", "wait", {"value": "a"}),
            WorkflowNode("b", "wait", {"value": "b"}),
        ],
    )
    arrived = 0
    both_arrived = asyncio.Event()

    async def wait(value):
        nonlocal arrived
        arrived += 1
        if arrived == 2:
            both_arrived.set()
        await asyncio.wait_for(both_arrived.wait(), timeout=0.5)
        return value

    result = asyncio.run(execute_workflow(workflow, {"wait": wait}))
    assert result.output == {"a": "a", "b": "b"}
    assert result.metrics.max_parallelism == 2
    assert result.metrics.layer_count == 1


def test_registry_decorator_and_missing_tool_errors_are_clear():
    registry = ToolRegistry()

    @registry.register("echo")
    async def echo(value):
        return value

    assert registry.names() == ("echo",)
    workflow = Workflow("missing", {}, [WorkflowNode("a", "absent", {})])
    with pytest.raises(ToolNotFoundError):
        asyncio.run(execute_workflow(workflow, registry))


def test_unresolvable_reference_is_wrapped_with_node_context():
    workflow = Workflow(
        "bad-ref",
        {},
        [
            WorkflowNode(
                "a",
                "echo",
                {"value": Reference("input", ("missing",))},
            )
        ],
    )
    with pytest.raises(PolicyViolationError, match="UNRESOLVED_INPUT_REFERENCE"):
        asyncio.run(execute_workflow(workflow, {"echo": lambda value: value}))


def test_runtime_denies_unsafe_workflow_before_calling_handler():
    calls = []

    def destructive():
        calls.append("called")

    workflow = Workflow(
        "unsafe",
        {},
        [WorkflowNode("delete", "delete", {}, side_effect="delete")],
        safety="unsafe",
    )
    registry = ToolRegistry(
        {"delete": destructive},
        contracts={"delete": ToolContract("delete", side_effect="delete")},
    )
    with pytest.raises(PolicyViolationError, match="denied"):
        asyncio.run(execute_workflow(workflow, registry))
    assert calls == []


def test_registered_contract_overrides_untrusted_trace_label():
    calls = []
    workflow = Workflow(
        "spoofed",
        {},
        [WorkflowNode("send", "gmail.send", {}, side_effect="read")],
    )
    registry = ToolRegistry(
        {"gmail.send": lambda: calls.append("sent")},
        contracts={
            "gmail.send": ToolContract(
                "gmail.send",
                side_effect="external_side_effect",
                idempotency="non_idempotent",
            )
        },
    )
    with pytest.raises(PolicyViolationError, match="approval"):
        asyncio.run(execute_workflow(workflow, registry))
    assert calls == []

    result = asyncio.run(execute_workflow(workflow, registry, approvals={"send"}))
    assert calls == ["sent"]
    assert result.output is None


def test_runtime_respects_bounded_concurrency():
    active = 0
    peak = 0

    async def work(value):
        nonlocal active, peak
        active += 1
        peak = max(peak, active)
        await asyncio.sleep(0.01)
        active -= 1
        return value

    workflow = Workflow(
        "bounded",
        {},
        [WorkflowNode(f"n{index}", "work", {"value": index}) for index in range(20)],
    )
    result = asyncio.run(execute_workflow(workflow, {"work": work}, max_concurrency=3))
    assert result.metrics.nodes_executed == 20
    assert peak == 3


def test_retry_requires_reviewed_idempotent_contract():
    attempts = 0

    async def flaky():
        nonlocal attempts
        attempts += 1
        if attempts < 3:
            raise RuntimeError("temporary")
        return "ok"

    workflow = Workflow(
        "retry",
        {},
        [WorkflowNode("flaky", "flaky", {}, idempotency="idempotent", max_attempts=3)],
    )
    registry = ToolRegistry(
        {"flaky": flaky},
        contracts={"flaky": ToolContract("flaky", side_effect="read", idempotency="idempotent")},
    )
    assert asyncio.run(execute_workflow(workflow, registry)).output == "ok"
    assert attempts == 3

    unsafe = Workflow(
        "unsafe-retry",
        {},
        [WorkflowNode("send", "send", {}, max_attempts=2)],
    )
    unsafe_registry = ToolRegistry(
        {"send": lambda: None},
        contracts={
            "send": ToolContract(
                "send",
                side_effect="external_side_effect",
                idempotency="non_idempotent",
            )
        },
    )
    with pytest.raises(PolicyViolationError, match="requests retry"):
        asyncio.run(execute_workflow(unsafe, unsafe_registry, approvals={"send"}))


def test_async_siblings_are_cancelled_and_awaited_after_failure():
    completed = False

    async def fail():
        await asyncio.sleep(0.01)
        raise RuntimeError("boom")

    async def slow():
        nonlocal completed
        await asyncio.sleep(0.2)
        completed = True

    workflow = Workflow(
        "cancel",
        {},
        [WorkflowNode("fail", "fail", {}), WorkflowNode("slow", "slow", {})],
    )

    async def scenario():
        with pytest.raises(WorkflowExecutionError):
            await execute_workflow(workflow, {"fail": fail, "slow": slow})
        await asyncio.sleep(0.25)

    asyncio.run(scenario())
    assert completed is False


def test_step_and_global_timeouts_stop_async_work():
    async def slow():
        await asyncio.sleep(1)

    step_timeout = Workflow(
        "step-timeout",
        {},
        [WorkflowNode("slow", "slow", {}, timeout_seconds=0.01)],
    )
    with pytest.raises(WorkflowExecutionError):
        asyncio.run(execute_workflow(step_timeout, {"slow": slow}))

    global_timeout = Workflow(
        "global-timeout",
        {},
        [WorkflowNode("slow", "slow", {})],
    )
    with pytest.raises(TimeoutError):
        asyncio.run(execute_workflow(global_timeout, {"slow": slow}, timeout_seconds=0.01))