Spaces:
Running
Running
| 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() | |
| 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)) | |