mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-17 18:25:42 +02:00
75 lines
3.2 KiB
Python
75 lines
3.2 KiB
Python
"""Unit coverage for the extracted SDK permission/pre-tool hooks (gate_hooks). These ran
|
|
inside the agent loop's closure before and had no isolated tests; pin the contract now that
|
|
they're a module: policy -> allow/deny, and the ToolSearch loop-breaker fire + reset."""
|
|
|
|
import pytest
|
|
from unittest.mock import patch, AsyncMock
|
|
|
|
from claude_agent_sdk import PermissionResultAllow, PermissionResultDeny
|
|
|
|
from backend.apps.agents.core.models import AgentSession
|
|
from backend.apps.agents.manager.streaming.HookContext import HookContext
|
|
from backend.apps.agents.manager.permissions import gate_hooks
|
|
from backend.apps.agents.manager.permissions.ApprovalDecision import ApprovalDecision
|
|
from backend.apps.agents.manager.prompt.prompt_context import TOOLSEARCH_LOOP_THRESHOLD
|
|
|
|
|
|
def p_ctx() -> HookContext:
|
|
session = AgentSession(name="t", model="sonnet", dashboard_id="d")
|
|
return HookContext(
|
|
session=session,
|
|
session_id=session.id,
|
|
prompt="hi",
|
|
builtin_perms={},
|
|
policy_defaults={},
|
|
sessions={},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_can_use_tool_always_allow_returns_allow():
|
|
ctx = p_ctx()
|
|
with patch.object(gate_hooks.path_gate, "maybe_override_policy", return_value=("always_allow", None)):
|
|
result = await gate_hooks.can_use_tool(ctx, "Read", {"file_path": "/x"}, None)
|
|
assert isinstance(result, PermissionResultAllow)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_can_use_tool_deny_returns_deny():
|
|
ctx = p_ctx()
|
|
with patch.object(gate_hooks.path_gate, "maybe_override_policy", return_value=("deny", None)):
|
|
result = await gate_hooks.can_use_tool(ctx, "Bash", {"command": "rm -rf /"}, None)
|
|
assert isinstance(result, PermissionResultDeny)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_can_use_tool_ask_routes_through_approval():
|
|
ctx = p_ctx()
|
|
with patch.object(gate_hooks.path_gate, "maybe_override_policy", return_value=("ask", None)), \
|
|
patch.object(gate_hooks, "resolve_ask", new=AsyncMock(return_value=ApprovalDecision(behavior="allow"))):
|
|
result = await gate_hooks.can_use_tool(ctx, "Write", {"file_path": "/x"}, None)
|
|
assert isinstance(result, PermissionResultAllow)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_tool_hook_toolsearch_loopbreaker_fires_at_threshold():
|
|
ctx = p_ctx()
|
|
ctx.ts_loop_count = TOOLSEARCH_LOOP_THRESHOLD - 1
|
|
with patch.object(gate_hooks, "gated_mcp_server_names", return_value=["gmail"]), \
|
|
patch.object(gate_hooks, "toolsearch_loop_redirect", return_value="Stop calling ToolSearch; use MCPActivate"):
|
|
out = await gate_hooks.pre_tool_hook(ctx, {"tool_name": "ToolSearch"}, "tu1", None)
|
|
deny = out["hookSpecificOutput"]
|
|
assert deny["permissionDecision"] == "deny"
|
|
assert "MCPActivate" in deny["permissionDecisionReason"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pre_tool_hook_counter_resets_on_non_toolsearch():
|
|
ctx = p_ctx()
|
|
ctx.ts_loop_count = 5
|
|
with patch.object(gate_hooks.path_gate, "maybe_override_policy", return_value=("always_allow", None)):
|
|
out = await gate_hooks.pre_tool_hook(ctx, {"tool_name": "Read", "tool_input": {}}, "tu1", None)
|
|
assert ctx.ts_loop_count == 0
|
|
assert out == {}
|
|
assert "tu1" in ctx.tool_start_times # an allowed tool records its start time
|