189 lines
7.2 KiB
Python
189 lines
7.2 KiB
Python
"""Regression test: legacy single-agent runs are unaffected by WriteGuard.
|
|
|
|
The WriteGuard middleware is registered unconditionally (when configured) but
|
|
must be a pure pass-through when the run context has no ``maf_run_id`` flag.
|
|
This test proves a write tool call in a legacy run executes normally (0
|
|
interceptions, real handler invoked), so the new middleware introduces zero
|
|
behaviour change for non-MAF runs.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from deerflow.agents.middlewares.write_guard_middleware import WriteGuardMiddleware
|
|
|
|
|
|
class _StubPolicy:
|
|
def should_intercept(self, tool_name: str, args: Any, run_context: Any = None) -> bool: # noqa: ANN401
|
|
return True # would intercept — but only if active
|
|
|
|
|
|
class _StubInterceptor:
|
|
async def intercept(self, *, tool_name: str, args: Any, tool_call_id: str, run_context: dict) -> dict: # noqa: ANN401
|
|
raise AssertionError("interceptor must NOT be called in a legacy run")
|
|
|
|
async def validate_approved_action(self, *, tool_name: str, run_context: dict) -> bool:
|
|
# Legacy stub: no real store. The approved-passthrough tests supply
|
|
# their own validating interceptor; this default fails closed.
|
|
return False
|
|
|
|
|
|
class _ValidatingStubInterceptor:
|
|
"""Approved-passthrough interceptor whose action validation is configurable."""
|
|
|
|
def __init__(self, *, valid: bool) -> None:
|
|
self._valid = valid
|
|
self.validated_with: dict | None = None
|
|
|
|
async def intercept(self, *, tool_name: str, args: Any, tool_call_id: str, run_context: dict) -> dict: # noqa: ANN401
|
|
return {"status": "approval_required", "action_id": "a?", "tool_name": tool_name}
|
|
|
|
async def validate_approved_action(self, *, tool_name: str, run_context: dict) -> bool:
|
|
self.validated_with = {"tool_name": tool_name, "run_context": run_context}
|
|
return self._valid
|
|
|
|
|
|
def _request(run_context: dict | None = None):
|
|
runtime = MagicMock()
|
|
runtime.context = run_context or {}
|
|
req = MagicMock()
|
|
req.tool_call = {"name": "write_anything", "args": {"x": 1}, "id": "c1"}
|
|
req.runtime = runtime
|
|
return req
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_legacy_run_write_tool_not_intercepted():
|
|
"""No ``maf_run_id`` → middleware passes through, real handler runs."""
|
|
mw = WriteGuardMiddleware(policy=_StubPolicy(), interceptor=_StubInterceptor())
|
|
called = {"n": 0}
|
|
|
|
async def handler(request):
|
|
called["n"] += 1
|
|
return MagicMock(content="ok")
|
|
|
|
# No maf_run_id in context = legacy run.
|
|
result = await mw.awrap_tool_call(_request({}), handler)
|
|
assert called["n"] == 1
|
|
assert result.content == "ok"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_maf_run_write_tool_intercepted():
|
|
"""With ``maf_run_id`` set, the same policy DOES intercept (sanity check)."""
|
|
called = {"n": 0}
|
|
|
|
async def handler(request):
|
|
called["n"] += 1
|
|
return MagicMock(content="ok")
|
|
|
|
# Active MAF run → interceptor would be called, but ours raises to prove it.
|
|
# Use a real interceptor here to avoid the assertion.
|
|
class _Real:
|
|
async def intercept(self, *, tool_name, args, tool_call_id, run_context):
|
|
return {"status": "approval_required", "action_id": "a1", "approval_id": "p1"}
|
|
|
|
mw2 = WriteGuardMiddleware(policy=_StubPolicy(), interceptor=_Real())
|
|
result = await mw2.awrap_tool_call(_request({"maf_run_id": "r1"}), handler)
|
|
assert called["n"] == 0 # handler NOT called
|
|
envelope = json.loads(result.content)
|
|
assert envelope["status"] == "approval_required"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_approved_passthrough_executes():
|
|
"""An approved-tool run (maf_approved_tool + maf_approved_action_id bound and
|
|
the action validates) lets that one tool through (P0-4 action-id security)."""
|
|
interceptor = _ValidatingStubInterceptor(valid=True)
|
|
mw = WriteGuardMiddleware(policy=_StubPolicy(), interceptor=interceptor)
|
|
called = {"n": 0}
|
|
|
|
async def handler(request):
|
|
called["n"] += 1
|
|
return MagicMock(content="approved-exec")
|
|
|
|
ctx = {
|
|
"maf_run_id": "r1",
|
|
"maf_session_id": "s1",
|
|
"agent_id": "seat-1",
|
|
"maf_approved_tool": "write_anything",
|
|
"maf_approved_action_id": "action-1",
|
|
}
|
|
result = await mw.awrap_tool_call(_request(ctx), handler)
|
|
assert called["n"] == 1
|
|
assert result.content == "approved-exec"
|
|
# The async path actually validated the bound action (not just field presence).
|
|
assert interceptor.validated_with is not None
|
|
assert interceptor.validated_with["tool_name"] == "write_anything"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_approved_passthrough_fails_closed_when_validation_false():
|
|
"""P0-4: if the bound action does NOT validate (status/ownership/tool mismatch),
|
|
the approved passthrough fails closed — the tool does not run."""
|
|
interceptor = _ValidatingStubInterceptor(valid=False)
|
|
mw = WriteGuardMiddleware(policy=_StubPolicy(), interceptor=interceptor)
|
|
called = {"n": 0}
|
|
|
|
async def handler(request):
|
|
called["n"] += 1
|
|
return MagicMock(content="should-not-run")
|
|
|
|
ctx = {
|
|
"maf_run_id": "r1",
|
|
"maf_approved_tool": "write_anything",
|
|
"maf_approved_action_id": "action-bad",
|
|
}
|
|
result = await mw.awrap_tool_call(_request(ctx), handler)
|
|
assert called["n"] == 0, "invalid approved action must fail closed"
|
|
envelope = json.loads(result.content)
|
|
assert envelope["status"] == "approval_invalid"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_approved_passthrough_requires_action_id():
|
|
"""P0-4: maf_approved_tool WITHOUT maf_approved_action_id no longer passes
|
|
through — action-id binding is mandatory."""
|
|
mw = WriteGuardMiddleware(policy=_StubPolicy(), interceptor=_ValidatingStubInterceptor(valid=True))
|
|
called = {"n": 0}
|
|
|
|
async def handler(request):
|
|
called["n"] += 1
|
|
return MagicMock(content="ok")
|
|
|
|
ctx = {"maf_run_id": "r1", "maf_approved_tool": "write_anything"} # no action id
|
|
result = await mw.awrap_tool_call(_request(ctx), handler)
|
|
assert called["n"] == 0
|
|
envelope = json.loads(result.content)
|
|
assert envelope["status"] in ("approval_required", "intercept_failed")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_approved_passthrough_only_for_named_tool():
|
|
"""When maf_approved_tool is set but the call is a DIFFERENT tool, it is
|
|
still intercepted (only the approved tool passes through)."""
|
|
class _Real:
|
|
async def intercept(self, *, tool_name, args, tool_call_id, run_context):
|
|
return {"status": "approval_required", "action_id": "a2"}
|
|
|
|
mw = WriteGuardMiddleware(policy=_StubPolicy(), interceptor=_Real())
|
|
called = {"n": 0}
|
|
|
|
async def handler(request):
|
|
called["n"] += 1
|
|
return MagicMock(content="ok")
|
|
|
|
ctx = {"maf_run_id": "r1", "maf_approved_tool": "update_task"}
|
|
# The request is for a different tool → not allowed through.
|
|
req = _request(ctx)
|
|
req.tool_call = {"name": "delete_thing", "args": {}, "id": "c2"}
|
|
result = await mw.awrap_tool_call(req, handler)
|
|
assert called["n"] == 0 # intercepted, not executed
|
|
envelope = json.loads(result.content)
|
|
assert envelope["status"] == "approval_required"
|