Spaces:
Running
Running
| import asyncio | |
| from collections.abc import Callable | |
| from typing import Any | |
| import pytest | |
| from trace2mcp.compiler import compile_trace | |
| from trace2mcp.hol_guard import ( | |
| HolGuardUnavailableError, | |
| _normalize_decision, | |
| hol_guard_from_environment, | |
| ) | |
| from trace2mcp.models import PolicyViolationError, Trace, TraceEvent | |
| from trace2mcp.policy import ToolContract | |
| from trace2mcp.runtime import ( | |
| ToolRegistry, | |
| WorkflowExecutionError, | |
| execute_workflow, | |
| invoke_capability, | |
| ) | |
| from trace2mcp.runtime_guard import ( | |
| RuntimeGuardAction, | |
| RuntimeGuardDecision, | |
| RuntimeGuardRequest, | |
| ) | |
| class RecordingGuard: | |
| def __init__( | |
| self, | |
| decide: Callable[[RuntimeGuardRequest], RuntimeGuardDecision], | |
| ) -> None: | |
| self.decide = decide | |
| self.requests: list[RuntimeGuardRequest] = [] | |
| async def authorize(self, request: RuntimeGuardRequest) -> RuntimeGuardDecision: | |
| self.requests.append(request) | |
| return self.decide(request) | |
| def event( | |
| tool: str, | |
| arguments: dict[str, Any], | |
| result: Any, | |
| *, | |
| side_effect: str = "read", | |
| ) -> TraceEvent: | |
| return TraceEvent( | |
| id=tool, | |
| tool=tool, | |
| arguments=arguments, | |
| result=result, | |
| side_effect=side_effect, | |
| idempotency="idempotent", | |
| ) | |
| def contract(tool: str, side_effect: str = "read") -> ToolContract: | |
| return ToolContract( | |
| tool, | |
| side_effect=side_effect, | |
| idempotency="idempotent", | |
| reviewed=True, | |
| ) | |
| def strict_registry( | |
| handlers: dict[str, Callable[..., Any]], | |
| contracts: dict[str, ToolContract], | |
| ) -> ToolRegistry: | |
| return ToolRegistry( | |
| handlers, | |
| contracts=contracts, | |
| require_contracts=True, | |
| require_reviewed_contracts=True, | |
| ) | |
| def test_guard_denied_mutating_call_has_zero_side_effect() -> None: | |
| state = {"handler_calls": 0, "writes": 0} | |
| async def mutate(command: str) -> str: | |
| del command | |
| state["handler_calls"] += 1 | |
| state["writes"] += 1 | |
| return "changed" | |
| workflow = compile_trace( | |
| Trace( | |
| inputs={}, | |
| events=[ | |
| event( | |
| "bash", | |
| {"command": "rm -rf ./data"}, | |
| "changed", | |
| side_effect="external_side_effect", | |
| ) | |
| ], | |
| ), | |
| allow_side_effects=True, | |
| ) | |
| guard = RecordingGuard( | |
| lambda request: RuntimeGuardDecision( | |
| RuntimeGuardAction.DENY, | |
| "mutation denied", | |
| source="test_guard", | |
| ) | |
| ) | |
| registry = strict_registry( | |
| {"bash": mutate}, | |
| {"bash": contract("bash", "external_side_effect")}, | |
| ) | |
| with pytest.raises(WorkflowExecutionError, match="mutation denied"): | |
| asyncio.run( | |
| execute_workflow( | |
| workflow, | |
| registry, | |
| approvals={"bash"}, | |
| runtime_guard=guard, | |
| ) | |
| ) | |
| assert len(guard.requests) == 1 | |
| assert state == {"handler_calls": 0, "writes": 0} | |
| def test_guard_allow_executes_exact_handler_once() -> None: | |
| calls = 0 | |
| async def lookup(value: int) -> int: | |
| nonlocal calls | |
| calls += 1 | |
| return value + 1 | |
| workflow = compile_trace(Trace(inputs={}, events=[event("lookup", {"value": 1}, 2)])) | |
| guard = RecordingGuard( | |
| lambda request: RuntimeGuardDecision(RuntimeGuardAction.ALLOW, "allowed") | |
| ) | |
| result = asyncio.run( | |
| execute_workflow( | |
| workflow, | |
| strict_registry({"lookup": lookup}, {"lookup": contract("lookup")}), | |
| runtime_guard=guard, | |
| ) | |
| ) | |
| assert result.output == 2 | |
| assert calls == 1 | |
| assert guard.requests[0].tool_name == "lookup" | |
| def test_trace2mcp_denial_wins_before_external_guard_allow() -> None: | |
| calls = 0 | |
| async def destroy() -> None: | |
| nonlocal calls | |
| calls += 1 | |
| workflow = compile_trace( | |
| Trace(inputs={}, events=[event("destroy", {}, None, side_effect="delete")]), | |
| allow_side_effects=True, | |
| ) | |
| guard = RecordingGuard( | |
| lambda request: RuntimeGuardDecision(RuntimeGuardAction.ALLOW, "external allow") | |
| ) | |
| with pytest.raises(PolicyViolationError, match="denied"): | |
| asyncio.run( | |
| execute_workflow( | |
| workflow, | |
| strict_registry( | |
| {"destroy": destroy}, | |
| {"destroy": contract("destroy", "delete")}, | |
| ), | |
| approvals={"destroy"}, | |
| runtime_guard=guard, | |
| ) | |
| ) | |
| assert calls == 0 | |
| assert guard.requests == [] | |
| def test_reviewed_contract_classification_reaches_guard_not_trace_metadata() -> None: | |
| async def send(message: str) -> str: | |
| return message | |
| workflow = compile_trace( | |
| Trace(inputs={}, events=[event("send", {"message": "hello"}, "hello")]) | |
| ) | |
| guard = RecordingGuard( | |
| lambda request: RuntimeGuardDecision(RuntimeGuardAction.ALLOW, "allowed") | |
| ) | |
| registry = strict_registry( | |
| {"send": send}, | |
| {"send": contract("send", "external_side_effect")}, | |
| ) | |
| assert ( | |
| asyncio.run( | |
| execute_workflow( | |
| workflow, | |
| registry, | |
| approvals={"send"}, | |
| runtime_guard=guard, | |
| ) | |
| ).output | |
| == "hello" | |
| ) | |
| request = guard.requests[0] | |
| assert request.side_effect == "external_side_effect" | |
| assert request.operation_classification == "external_side_effect" | |
| assert request.approval_satisfied is True | |
| assert request.contract_reviewed is True | |
| def test_nested_managed_invocation_traverses_same_guard_and_is_denied() -> None: | |
| state = {"nested_calls": 0} | |
| async def inner(command: str) -> str: | |
| del command | |
| state["nested_calls"] += 1 | |
| return "changed" | |
| async def outer() -> str: | |
| return await invoke_capability("bash", command="rm -rf ./data") | |
| workflow = compile_trace(Trace(inputs={}, events=[event("outer", {}, "changed")])) | |
| guard = RecordingGuard( | |
| lambda request: RuntimeGuardDecision( | |
| RuntimeGuardAction.DENY if request.tool_name == "bash" else RuntimeGuardAction.ALLOW, | |
| "nested deny" if request.tool_name == "bash" else "outer allow", | |
| ) | |
| ) | |
| registry = strict_registry( | |
| {"outer": outer, "bash": inner}, | |
| {"outer": contract("outer"), "bash": contract("bash")}, | |
| ) | |
| with pytest.raises(WorkflowExecutionError, match="nested deny"): | |
| asyncio.run(execute_workflow(workflow, registry, runtime_guard=guard)) | |
| assert [request.tool_name for request in guard.requests] == ["outer", "bash"] | |
| assert guard.requests[1].nested is True | |
| assert state["nested_calls"] == 0 | |
| def test_registry_proxy_indirect_call_traverses_same_guard() -> None: | |
| state = {"inner_calls": 0} | |
| registry: ToolRegistry | |
| async def inner() -> str: | |
| state["inner_calls"] += 1 | |
| return "inner" | |
| async def outer() -> str: | |
| capability = registry.get("inner") | |
| return await capability() | |
| workflow = compile_trace(Trace(inputs={}, events=[event("outer", {}, "inner")])) | |
| guard = RecordingGuard( | |
| lambda request: RuntimeGuardDecision( | |
| RuntimeGuardAction.DENY if request.tool_name == "inner" else RuntimeGuardAction.ALLOW, | |
| "indirect denied" if request.tool_name == "inner" else "outer allowed", | |
| ) | |
| ) | |
| registry = strict_registry( | |
| {"outer": outer, "inner": inner}, | |
| {"outer": contract("outer"), "inner": contract("inner")}, | |
| ) | |
| with pytest.raises(WorkflowExecutionError, match="indirect denied"): | |
| asyncio.run(execute_workflow(workflow, registry, runtime_guard=guard)) | |
| assert [request.tool_name for request in guard.requests] == ["outer", "inner"] | |
| assert state["inner_calls"] == 0 | |
| def test_guard_approval_decision_pauses_before_handler() -> None: | |
| calls = 0 | |
| async def handler() -> str: | |
| nonlocal calls | |
| calls += 1 | |
| return "ok" | |
| workflow = compile_trace(Trace(inputs={}, events=[event("handler", {}, "ok")])) | |
| guard = RecordingGuard( | |
| lambda request: RuntimeGuardDecision( | |
| RuntimeGuardAction.REQUIRE_APPROVAL, | |
| "HOL Guard review required", | |
| ) | |
| ) | |
| registry = strict_registry({"handler": handler}, {"handler": contract("handler")}) | |
| with pytest.raises(WorkflowExecutionError, match="requires approval"): | |
| asyncio.run(execute_workflow(workflow, registry, runtime_guard=guard)) | |
| assert calls == 0 | |
| def test_guard_exception_fails_closed_before_handler() -> None: | |
| calls = 0 | |
| async def handler() -> str: | |
| nonlocal calls | |
| calls += 1 | |
| return "ok" | |
| class BrokenGuard: | |
| async def authorize(self, request: RuntimeGuardRequest) -> RuntimeGuardDecision: | |
| del request | |
| raise RuntimeError("guard crashed") | |
| workflow = compile_trace(Trace(inputs={}, events=[event("handler", {}, "ok")])) | |
| registry = strict_registry({"handler": handler}, {"handler": contract("handler")}) | |
| with pytest.raises(WorkflowExecutionError, match="managed execution denied"): | |
| asyncio.run(execute_workflow(workflow, registry, runtime_guard=BrokenGuard())) | |
| assert calls == 0 | |
| def test_guard_receives_immutable_argument_snapshot() -> None: | |
| async def handler(payload: dict[str, list[int]]) -> str: | |
| return "ok" | |
| class MutationGuard: | |
| async def authorize(self, request: RuntimeGuardRequest) -> RuntimeGuardDecision: | |
| with pytest.raises(TypeError): | |
| request.arguments["payload"] = {} # type: ignore[index] | |
| with pytest.raises(TypeError): | |
| request.arguments["payload"]["values"] = () # type: ignore[index] | |
| return RuntimeGuardDecision(RuntimeGuardAction.ALLOW, "immutable") | |
| workflow = compile_trace( | |
| Trace( | |
| inputs={}, | |
| events=[event("handler", {"payload": {"values": [1, 2]}}, "ok")], | |
| ) | |
| ) | |
| registry = strict_registry({"handler": handler}, {"handler": contract("handler")}) | |
| assert ( | |
| asyncio.run(execute_workflow(workflow, registry, runtime_guard=MutationGuard())).output | |
| == "ok" | |
| ) | |
| def test_handler_swap_during_guard_is_rejected_before_any_handler() -> None: | |
| state = {"original_calls": 0, "replacement_calls": 0} | |
| async def original() -> str: | |
| state["original_calls"] += 1 | |
| return "original" | |
| async def replacement() -> str: | |
| state["replacement_calls"] += 1 | |
| return "replacement" | |
| workflow = compile_trace(Trace(inputs={}, events=[event("handler", {}, "original")])) | |
| registry = strict_registry({"handler": original}, {"handler": contract("handler")}) | |
| class SwappingGuard: | |
| async def authorize(self, request: RuntimeGuardRequest) -> RuntimeGuardDecision: | |
| del request | |
| registry.register("handler", replacement) | |
| return RuntimeGuardDecision(RuntimeGuardAction.ALLOW, "attempted swap") | |
| with pytest.raises(WorkflowExecutionError, match="changed during guard authorization"): | |
| asyncio.run(execute_workflow(workflow, registry, runtime_guard=SwappingGuard())) | |
| assert state == {"original_calls": 0, "replacement_calls": 0} | |
| def test_hol_guard_actions_are_normalized_fail_closed() -> None: | |
| assert _normalize_decision({"policy_action": "allow"}).action is RuntimeGuardAction.ALLOW | |
| assert _normalize_decision({"policy_action": "warn"}).action is RuntimeGuardAction.ALLOW | |
| assert ( | |
| _normalize_decision({"policy_action": "review"}).action | |
| is RuntimeGuardAction.REQUIRE_APPROVAL | |
| ) | |
| assert _normalize_decision({"policy_action": "block"}).action is RuntimeGuardAction.DENY | |
| with pytest.raises(Exception, match="Unsupported HOL Guard"): | |
| _normalize_decision({"policy_action": "future-action"}) | |
| def test_missing_hol_guard_is_clear_only_when_explicitly_enabled() -> None: | |
| assert hol_guard_from_environment({}) is None | |
| with pytest.raises(HolGuardUnavailableError, match="pipx install hol-guard"): | |
| hol_guard_from_environment( | |
| {"TRACE2MCP_HOL_GUARD_EXECUTABLE": "definitely-missing-hol-guard-binary"} | |
| ) | |