fischer-agentkit/tests/unit/test_react_phase_enforcemen...

620 lines
26 KiB
Python

"""Unit tests for ReActEngine phase enforcement (G6 wiring, R24).
Per plan U3 Execution note: characterization-first — verify that
`ReActEngine(phase_policy=None)` behaves identically to pre-change (no
enforcement, no advance_phase tool, no _current_phase mutation). Then add
enforcement tests.
"""
from __future__ import annotations
from unittest.mock import AsyncMock, MagicMock
import pytest
from agentkit.core.phase import PhasePolicy, PhaseState, default_policy
from agentkit.core.react import ReActEngine
from agentkit.tools.advance_phase import AdvancePhaseTool
# ---------------------------------------------------------------------------
# Characterization — phase_policy=None preserves existing behavior
# ---------------------------------------------------------------------------
class TestCharacterizationNoPolicy:
"""When phase_policy=None, no enforcement happens and behavior matches
pre-Wave-3."""
def test_init_without_phase_policy(self):
# Minimal stub LLM gateway — we're only testing constructor.
gateway = MagicMock()
engine = ReActEngine(llm_gateway=gateway)
assert engine._phase_policy is None
assert engine._current_phase is None
assert engine._steps_in_phase == 0
assert engine.current_phase is None
@pytest.mark.asyncio
async def test_execute_tool_dispatches_without_phase_check(self):
"""Tool dispatch proceeds normally when no policy set."""
gateway = MagicMock()
engine = ReActEngine(llm_gateway=gateway)
# MagicMock.name is a special attribute used internally by Mock for
# repr — setting it post-construction does not make mock.name == "x"
# hold. Patch _find_tool directly to bypass the name lookup.
fake_tool = MagicMock()
fake_tool.safe_execute = AsyncMock(return_value={"output": "ok"})
fake_tool.input_schema = None
engine._find_tool = lambda name, tools: fake_tool
result = await engine._execute_tool("any_tool", {"x": 1}, [fake_tool])
assert result == {"output": "ok"}
fake_tool.safe_execute.assert_awaited_once_with(x=1)
@pytest.mark.asyncio
async def test_advance_phase_returns_none_without_policy(self):
gateway = MagicMock()
engine = ReActEngine(llm_gateway=gateway)
assert engine.advance_phase() is None
def test_reset_does_not_touch_phase_state_when_no_policy(self):
gateway = MagicMock()
engine = ReActEngine(llm_gateway=gateway)
engine.reset()
assert engine._current_phase is None
# ---------------------------------------------------------------------------
# Initialization with phase_policy
# ---------------------------------------------------------------------------
class TestPhasePolicyInitialization:
def test_phase_policy_set_initializes_current_phase(self):
gateway = MagicMock()
engine = ReActEngine(
llm_gateway=gateway,
phase_policy=default_policy(),
)
assert engine._phase_policy is not None
assert engine._current_phase == PhaseState.PLANNING
assert engine._steps_in_phase == 0
def test_reset_resets_phase_to_start(self):
gateway = MagicMock()
engine = ReActEngine(
llm_gateway=gateway,
phase_policy=default_policy(),
)
# Manually move phase forward (simulating execute progress).
engine.advance_phase() # PLANNING → BUILDING
assert engine._current_phase == PhaseState.BUILDING
engine._steps_in_phase = 5
engine.reset()
assert engine._current_phase == PhaseState.PLANNING
assert engine._steps_in_phase == 0
# ---------------------------------------------------------------------------
# advance_phase() transitions
# ---------------------------------------------------------------------------
class TestAdvancePhase:
@pytest.fixture
def engine(self):
return ReActEngine(llm_gateway=MagicMock(), phase_policy=default_policy())
def test_planning_to_building(self, engine):
new_phase = engine.advance_phase()
assert new_phase == PhaseState.BUILDING
assert engine.current_phase == PhaseState.BUILDING
assert engine._steps_in_phase == 0 # counter reset on transition
def test_building_to_verification(self, engine):
engine.advance_phase() # → BUILDING
new_phase = engine.advance_phase()
assert new_phase == PhaseState.VERIFICATION
assert engine.current_phase == PhaseState.VERIFICATION
def test_verification_to_delivery(self, engine):
engine.advance_phase() # → BUILDING
engine.advance_phase() # → VERIFICATION
new_phase = engine.advance_phase()
assert new_phase == PhaseState.DELIVERY
assert engine.current_phase == PhaseState.DELIVERY
def test_delivery_returns_none(self, engine):
engine.advance_phase() # → BUILDING
engine.advance_phase() # → VERIFICATION
engine.advance_phase() # → DELIVERY
result = engine.advance_phase()
assert result is None
assert engine.current_phase == PhaseState.DELIVERY
# ---------------------------------------------------------------------------
# _check_phase_permission — whitelist enforcement
# ---------------------------------------------------------------------------
class TestPhasePermission:
@pytest.fixture
def engine(self):
return ReActEngine(llm_gateway=MagicMock(), phase_policy=default_policy())
def test_search_allowed_in_planning(self, engine):
assert engine._check_phase_permission("search", {}) is None
def test_write_file_blocked_in_planning(self, engine):
result = engine._check_phase_permission("write_file", {})
assert result is not None
assert result["error"] == "phase_violation"
assert "write_file" in result["message"]
assert result["current_phase"] == "planning"
def test_write_file_allowed_in_building(self, engine):
engine.advance_phase() # → BUILDING
assert engine._check_phase_permission("write_file", {}) is None
def test_any_tool_allowed_in_delivery(self, engine):
engine.advance_phase() # → BUILDING
engine.advance_phase() # → VERIFICATION
engine.advance_phase() # → DELIVERY
assert engine._check_phase_permission("literally_anything", {}) is None
def test_bash_command_filter_blocks_rm_in_planning(self, engine):
result = engine._check_phase_permission("shell", {"command": "rm -rf /tmp"})
assert result is not None
assert result["error"] == "phase_violation"
assert "rm" in result["message"] or "Bash command" in result["message"]
def test_bash_command_filter_allows_safe_in_planning(self, engine):
# `ls` and `git status` are not blocked.
assert engine._check_phase_permission("shell", {"command": "ls -la"}) is None
assert engine._check_phase_permission("shell", {"command": "git status"}) is None
def test_bash_command_filter_no_restriction_in_building(self, engine):
engine.advance_phase() # → BUILDING
# `rm` is allowed in building phase.
assert engine._check_phase_permission("shell", {"command": "rm -rf build/"}) is None
# ---------------------------------------------------------------------------
# _execute_tool integration — phase enforcement actually blocks dispatch
# ---------------------------------------------------------------------------
class TestExecuteToolPhaseEnforcement:
@pytest.fixture
def engine_with_tools(self):
engine = ReActEngine(llm_gateway=MagicMock(), phase_policy=default_policy())
# Two fake tools: one allowed in PLANNING (search), one not (write_file).
# MagicMock.name can't be set post-construction (special attribute),
# so we patch _find_tool with a dict-based lookup.
search_tool = MagicMock()
search_tool.input_schema = None
search_tool.safe_execute = AsyncMock(return_value={"results": []})
write_tool = MagicMock()
write_tool.input_schema = None
write_tool.safe_execute = AsyncMock(return_value={"written": True})
tools_by_name = {"search": search_tool, "write_file": write_tool}
engine._find_tool = lambda name, tools: tools_by_name.get(name)
return engine, [search_tool, write_tool]
@pytest.mark.asyncio
async def test_blocked_tool_returns_phase_violation_and_skips_dispatch(self, engine_with_tools):
engine, tools = engine_with_tools
# write_file in PLANNING should be blocked — write_tool.safe_execute
# should NEVER be called.
result = await engine._execute_tool("write_file", {"path": "/x"}, tools)
assert result["error"] == "phase_violation"
assert result["current_phase"] == "planning"
write_tool = tools[1]
write_tool.safe_execute.assert_not_called()
@pytest.mark.asyncio
async def test_allowed_tool_dispatches_normally(self, engine_with_tools):
engine, tools = engine_with_tools
result = await engine._execute_tool("search", {"query": "foo"}, tools)
assert result == {"results": []}
search_tool = tools[0]
search_tool.safe_execute.assert_awaited_once_with(query="foo")
@pytest.mark.asyncio
async def test_after_advance_phase_blocked_tool_now_dispatches(self, engine_with_tools):
engine, tools = engine_with_tools
# First: write_file blocked in PLANNING.
result = await engine._execute_tool("write_file", {"path": "/x"}, tools)
assert result["error"] == "phase_violation"
# Advance to BUILDING.
engine.advance_phase()
# Now: write_file allowed.
result = await engine._execute_tool("write_file", {"path": "/x"}, tools)
assert result == {"written": True}
# ---------------------------------------------------------------------------
# Auto-advance safety net (KTD6)
# ---------------------------------------------------------------------------
class TestAutoAdvance:
def test_auto_advance_after_threshold(self):
# Custom policy with auto-advance after 2 steps.
policy = PhasePolicy(
whitelist={
PhaseState.PLANNING: frozenset({"search"}),
PhaseState.BUILDING: frozenset({"write_file"}),
PhaseState.VERIFICATION: frozenset({"shell"}),
PhaseState.DELIVERY: frozenset({"*"}),
},
auto_advance_after_steps=2,
)
engine = ReActEngine(llm_gateway=MagicMock(), phase_policy=policy)
assert engine.current_phase == PhaseState.PLANNING
# Step 1: counter goes to 1, no advance yet.
engine._steps_in_phase += 1
assert engine._maybe_auto_advance() is False
assert engine.current_phase == PhaseState.PLANNING
# Step 2: counter hits 2, advance triggered.
engine._steps_in_phase += 1
assert engine._maybe_auto_advance() is True
assert engine.current_phase == PhaseState.BUILDING
assert engine._steps_in_phase == 0 # reset on advance
def test_auto_advance_none_default(self):
# default_policy has auto_advance_after_steps=None — no auto-advance.
engine = ReActEngine(llm_gateway=MagicMock(), phase_policy=default_policy())
engine._steps_in_phase = 100
assert engine._maybe_auto_advance() is False
assert engine.current_phase == PhaseState.PLANNING
# ---------------------------------------------------------------------------
# AdvancePhaseTool integration
# ---------------------------------------------------------------------------
class TestAdvancePhaseTool:
@pytest.mark.asyncio
async def test_advance_phase_tool_transitions_engine(self):
engine = ReActEngine(llm_gateway=MagicMock(), phase_policy=default_policy())
tool = AdvancePhaseTool(engine=engine)
result = await tool.execute()
assert result["is_error"] is False
assert result["current_phase"] == "building"
assert engine.current_phase == PhaseState.BUILDING
@pytest.mark.asyncio
async def test_advance_phase_tool_at_delivery_returns_error(self):
engine = ReActEngine(llm_gateway=MagicMock(), phase_policy=default_policy())
# Walk through all phases.
engine.advance_phase() # PLANNING → BUILDING
engine.advance_phase() # BUILDING → VERIFICATION
engine.advance_phase() # VERIFICATION → DELIVERY
tool = AdvancePhaseTool(engine=engine)
result = await tool.execute()
assert result["is_error"] is True
assert result["error"] == "already_at_final_phase"
assert result["current_phase"] == "delivery"
@pytest.mark.asyncio
async def test_advance_phase_tool_without_policy_returns_error(self):
engine = ReActEngine(llm_gateway=MagicMock()) # no policy
tool = AdvancePhaseTool(engine=engine)
result = await tool.execute()
assert result["is_error"] is True
assert result["error"] == "no_phase_policy"
def test_tool_schema_accepts_no_arguments(self):
engine = ReActEngine(llm_gateway=MagicMock(), phase_policy=default_policy())
tool = AdvancePhaseTool(engine=engine)
# input_schema has empty properties + additionalProperties:false —
# no arguments expected.
assert tool.input_schema["properties"] == {}
assert tool.input_schema["additionalProperties"] is False
def test_tool_bypasses_phase_check(self):
"""`advance_phase` is the LLM's escape hatch — must never be blocked."""
# _check_phase_permission should NOT block advance_phase even in PLANNING.
# The bypass is implemented in _execute_tool by name check.
# We verify the bypass indirectly: tool dispatches normally even in
# PLANNING (where only search/read_file/bash/tool_search are allowed).
# advance_phase is not in the whitelist, but the name-based bypass
# in _execute_tool lets it through.
# (Direct unit test of the bypass would require mocking _find_tool.)
# Sanity: advance_phase is not in any whitelist.
for phase, allowed in default_policy().whitelist.items():
assert "advance_phase" not in allowed, (
f"advance_phase must not be in {phase.value} whitelist"
)
# ---------------------------------------------------------------------------
# Wave 4 U2 — phase_violation accumulator + drain
# ---------------------------------------------------------------------------
class TestPhaseViolationAccumulator:
"""_check_phase_permission records violations; _drain_phase_violations
yields them as ReActEvents and clears the accumulator."""
@pytest.fixture
def engine(self):
return ReActEngine(
llm_gateway=MagicMock(),
phase_policy=default_policy(),
)
def test_violation_appended_on_tool_block(self, engine):
# write_file is blocked in PLANNING.
engine._check_phase_permission("write_file", {})
assert len(engine._phase_violations) == 1
v = engine._phase_violations[0]
assert v["error"] == "phase_violation"
assert v["tool"] == "write_file"
assert v["current_phase"] == "planning"
assert v["violation_kind"] == "tool_not_allowed"
def test_violation_appended_on_bash_block(self, engine):
engine._check_phase_permission("shell", {"command": "rm -rf /tmp"})
assert len(engine._phase_violations) == 1
v = engine._phase_violations[0]
assert v["violation_kind"] == "bash_command_blocked"
assert v["tool"] == "shell"
assert v["command_preview"] == "rm -rf /tmp"
def test_no_violation_when_allowed(self, engine):
# search is allowed in PLANNING.
engine._check_phase_permission("search", {})
assert engine._phase_violations == []
def test_no_violation_without_policy(self):
engine = ReActEngine(llm_gateway=MagicMock()) # no policy
engine._check_phase_permission("anything", {})
assert engine._phase_violations == []
def test_drain_returns_events_and_clears(self, engine):
# Trigger two violations.
engine._check_phase_permission("write_file", {"path": "/a"})
engine._check_phase_permission("write_file", {"path": "/b"})
assert len(engine._phase_violations) == 2
events = engine._drain_phase_violations(step=3)
assert len(events) == 2
assert all(e.event_type == "phase_violation" for e in events)
assert all(e.step == 3 for e in events)
# Each event data is a copy (caller can't mutate the accumulator).
assert events[0].data["tool"] == "write_file"
# Accumulator cleared after drain.
assert engine._phase_violations == []
def test_drain_empty_returns_empty(self, engine):
assert engine._drain_phase_violations(step=1) == []
def test_drain_returns_shallow_copy(self, engine):
"""Drained event data must not alias the original violation dict —
mutating one must not mutate the other."""
engine._check_phase_permission("write_file", {})
events = engine._drain_phase_violations(step=1)
# Mutate the drained event data.
events[0].data["tool"] = "MUTATED"
# Original accumulator (now empty) is unaffected — but more importantly,
# a fresh violation recorded after drain is unaffected.
engine._check_phase_permission("write_file", {})
new_violations = engine._phase_violations
assert new_violations[0]["tool"] == "write_file" # not "MUTATED"
def test_reset_clears_violations(self, engine):
engine._check_phase_permission("write_file", {})
assert len(engine._phase_violations) == 1
engine.reset()
assert engine._phase_violations == []
# ---------------------------------------------------------------------------
# Wave 4 U2 — execute_stream yields phase_violation events
# ---------------------------------------------------------------------------
class TestExecuteStreamPhaseViolationEvents:
"""execute_stream must yield phase_violation ReActEvents when a tool is
blocked by _check_phase_permission. The events are drained after each
tool_result yield."""
@pytest.mark.asyncio
async def test_stream_yields_phase_violation_on_tool_block(self):
"""When the LLM calls a tool blocked by the phase policy, execute_stream
yields a tool_call event, a tool_result event (with the error dict),
and a phase_violation event."""
from agentkit.core.react import ReActEvent
engine = ReActEngine(
llm_gateway=llm_mock_gateway_with_response(
tool_calls=[{"name": "write_file", "arguments": {"path": "/x"}}],
content=None,
),
phase_policy=default_policy(),
max_steps=1,
)
# Patch _find_tool so we don't need real tools registered. write_file
# should be blocked by phase_policy before _find_tool is called.
engine._find_tool = lambda name, tools: None
events: list[ReActEvent] = []
async for ev in engine.execute_stream(
messages=[{"role": "user", "content": "test"}],
tools=[],
):
events.append(ev)
# Expect: thinking → tool_call → tool_result → phase_violation → final_answer
event_types = [e.event_type for e in events]
assert "tool_call" in event_types
assert "tool_result" in event_types
assert "phase_violation" in event_types
# The phase_violation event must come AFTER tool_result.
tool_result_idx = next(i for i, e in enumerate(events) if e.event_type == "tool_result")
violation_idx = next(i for i, e in enumerate(events) if e.event_type == "phase_violation")
assert violation_idx > tool_result_idx
# Verify event data.
violation = events[violation_idx]
assert violation.data["error"] == "phase_violation"
assert violation.data["tool"] == "write_file"
assert violation.data["current_phase"] == "planning"
assert violation.data["violation_kind"] == "tool_not_allowed"
@pytest.mark.asyncio
async def test_stream_yields_phase_violation_on_bash_block(self):
"""When the LLM calls shell with a dangerous command in PLANNING,
execute_stream yields a phase_violation event with violation_kind
= bash_command_blocked."""
from agentkit.core.react import ReActEvent
engine = ReActEngine(
llm_gateway=llm_mock_gateway_with_response(
tool_calls=[{"name": "shell", "arguments": {"command": "rm -rf /tmp"}}],
content=None,
),
phase_policy=default_policy(),
max_steps=1,
)
engine._find_tool = lambda name, tools: None
events: list[ReActEvent] = []
async for ev in engine.execute_stream(
messages=[{"role": "user", "content": "test"}],
tools=[],
):
events.append(ev)
violation_events = [e for e in events if e.event_type == "phase_violation"]
assert len(violation_events) == 1
v = violation_events[0].data
assert v["violation_kind"] == "bash_command_blocked"
assert v["tool"] == "shell"
assert "rm -rf /tmp" in v["command_preview"]
@pytest.mark.asyncio
async def test_stream_no_violation_when_tool_allowed(self):
"""When the LLM calls an allowed tool, no phase_violation event is yielded."""
from agentkit.core.react import ReActEvent
engine = ReActEngine(
llm_gateway=llm_mock_gateway_with_response(
tool_calls=[{"name": "search", "arguments": {"query": "foo"}}],
content=None,
),
phase_policy=default_policy(),
max_steps=1,
)
# search is allowed in PLANNING; we still need _find_tool to return a
# tool object so dispatch proceeds.
search_tool = MagicMock()
search_tool.input_schema = None
search_tool.safe_execute = AsyncMock(return_value={"results": []})
engine._find_tool = lambda name, tools: search_tool
events: list[ReActEvent] = []
async for ev in engine.execute_stream(
messages=[{"role": "user", "content": "test"}],
tools=[search_tool],
):
events.append(ev)
assert not any(e.event_type == "phase_violation" for e in events)
@pytest.mark.asyncio
async def test_stream_no_violation_without_policy(self):
"""Without a phase_policy, no phase_violation events are yielded —
characterization of the no-policy path."""
from agentkit.core.react import ReActEvent
engine = ReActEngine(
llm_gateway=llm_mock_gateway_with_response(
tool_calls=[{"name": "any_tool", "arguments": {}}],
content=None,
),
max_steps=1,
)
any_tool = MagicMock()
any_tool.input_schema = None
any_tool.safe_execute = AsyncMock(return_value={"output": "ok"})
engine._find_tool = lambda name, tools: any_tool
events: list[ReActEvent] = []
async for ev in engine.execute_stream(
messages=[{"role": "user", "content": "test"}],
tools=[any_tool],
):
events.append(ev)
assert not any(e.event_type == "phase_violation" for e in events)
# ---------------------------------------------------------------------------
# Helpers — minimal LLM gateway mocks for execute_stream tests
# ---------------------------------------------------------------------------
def llm_mock_gateway():
"""Return a MagicMock LLM gateway (sufficient for constructor tests)."""
return MagicMock()
def llm_mock_gateway_with_response(tool_calls: list[dict], content: str | None):
"""Return a MagicMock LLM gateway whose chat_stream yields a single chunk
containing the given tool_calls (or content for a final-answer response).
The mock is shaped to match what execute_stream expects from
LLMGateway.chat_stream — an async iterable of chunks with attributes
`content`, `tool_calls`, `usage`, `model`.
"""
gateway = MagicMock()
# Build a fake chunk. execute_stream reads chunk.content, chunk.tool_calls,
# chunk.usage, chunk.model. The first three are typically accessed via
# attribute access; we make a small dataclass-like object.
class _Chunk:
def __init__(self, content, tool_calls, usage=None, model="default"):
self.content = content
self.tool_calls = tool_calls
self.usage = usage
self.model = model
# If tool_calls provided, emit a chunk with tool_calls (non-streaming path).
# Otherwise, emit a chunk with content (final answer path).
if tool_calls:
# Convert raw dicts to objects with .name/.arguments/.id attributes
# (LLMGateway normally returns tool_call objects).
class _TC:
def __init__(self, d):
self.name = d.get("name", "")
self.arguments = d.get("arguments", {})
self.id = d.get("id", "tc_test")
chunks = [_Chunk(content=None, tool_calls=[_TC(tc) for tc in tool_calls])]
# Follow with a final-answer chunk so execute_stream's loop exits
# cleanly after the tool call.
chunks.append(_Chunk(content="done", tool_calls=None))
else:
chunks = [_Chunk(content=content or "final answer", tool_calls=None)]
async def _fake_chat_stream(*args, **kwargs):
for c in chunks:
yield c
gateway.chat_stream = _fake_chat_stream
return gateway