WEBDIENSTE/openswarm
1
0
Fork 0
mirror of https://github.com/openswarm-ai/openswarm.git synced 2026-09-11 20:27:44 +02:00
Code Issues Packages Projects Releases Wiki Activity

[eric] agents: lift permission/pre-tool SDK hooks into gate_hooks + HookContext

Browse Source
This commit is contained in:
ciregenz
2026-06-23 01:59:14 -07:00
parent 0678ce3f77
commit aa3ce99978
4 changed files with 280 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/agents/agent_manager.py
+15 -142
View File
@@ -57,8 +57,9 @@ from backend.apps.agents.manager.session import lifecycle
from backend.apps.agents.manager.permissions import path_gate
from backend.apps.agents.manager import context_budget
from backend.apps.agents.manager.streaming.state import ThinkingState, TurnState
from backend.apps.agents.manager.streaming.hook_context import HookContext
from backend.apps.agents.manager.streaming import thinking as thinking_mod
from backend.apps.agents.manager.permissions.decision import effective_policy, request_user_approval
from backend.apps.agents.manager.permissions import gate_hooks
from backend.apps.agents.manager.session.workspace_git import _detect_git_identity, _ensure_cwd_git_repo
from backend.apps.agents.manager.prompt.tool_catalog import (
FULL_TOOLS,
@@ -85,8 +86,6 @@ from backend.apps.agents.manager.prompt.prompt_context import (
_resolve_attached_skills,
_resolve_forced_tools,
_resolve_mode,
TOOLSEARCH_LOOP_THRESHOLD,
toolsearch_loop_redirect,
)
from backend.apps.agents.manager.prompt.attachments import (
_build_dir_tree,
@@ -378,7 +377,7 @@ class AgentManager:
query, ClaudeAgentOptions, AssistantMessage, ResultMessage,
)
from claude_agent_sdk.types import (
HookMatcher, PermissionResultAllow, PermissionResultDeny,
HookMatcher,
TextBlock, ToolUseBlock, ThinkingBlock, StreamEvent,
SystemMessage,
)
@@ -406,156 +405,30 @@ class AgentManager:
# Per-tool DEFAULT policy (overridden by anything the user has set
# explicitly in builtin_permissions.json). Bash defaults to
# always_allow like every other builtin, for a frictionless run.
# Three guards below STILL force a prompt even on always_allow:
# Three guards in path_gate STILL force a prompt even on always_allow:
# the catastrophic-pattern match (rm -rf and friends), OS-scheduling
# (cron/launchd persistence), and the sensitive-path gate. So the
# poisoned-email -> destructive-command case is still caught; what
# this trades away is the prompt on ordinary shell commands. Users
# who want a prompt on every command can flip Bash to "ask" in the UI.
policy_defaults: Dict[str, str] = {}
hook_ctx = HookContext(
session=session,
session_id=session_id,
prompt=prompt,
builtin_perms=builtin_perms,
policy_defaults={},
)
async def can_use_tool(tool_name, input_data, context):
sensitive_pattern: str | None = None
if tool_name != "AskUserQuestion":
policy, sensitive_pattern = path_gate.maybe_override_policy(
effective_policy(tool_name, builtin_perms, policy_defaults), tool_name, input_data
)
if policy == "always_allow":
return PermissionResultAllow(updated_input=input_data)
if policy == "deny":
return PermissionResultDeny(message="Tool denied by permission policy")
decision = await request_user_approval(session, session_id, tool_name, input_data, builtin_perms, sensitive_pattern=sensitive_pattern)
if decision.get("behavior") == "allow":
return PermissionResultAllow(
updated_input=decision.get("updated_input", input_data)
)
return PermissionResultDeny(
message=decision.get("message", "User denied this action")
)
tool_start_times: dict[str, float] = {}
# Counts ToolSearch calls in a row (no other tool between them). A run
# of these with empty results is the "looping on ToolSearch" wedge.
_ts_loop = {"n": 0}
# One mid-run connect offer per session: a stuck agent fires the loop-breaker repeatedly,
# but the user should see the "connect this MCP" card once, not on every retry.
_mcp_offer_sent = {"done": False}
return await gate_hooks.can_use_tool(hook_ctx, tool_name, input_data, context)
async def pre_tool_hook(input_data, tool_use_id, context):
tool_name = input_data.get("tool_name", "")
hook_event = input_data.get("hook_event_name", "PreToolUse")
# ToolSearch loop-breaker. Gated MCP servers are withheld from the
# SDK until MCPActivate, so the CLI's native ToolSearch can never
# find them; small models thrash (empty ToolSearch, retry) for
# minutes until the user pauses. Let the first couple through, then
# redirect to the gate. Any non-ToolSearch call is real progress, so
# the counter resets. Gated-server lookup is deferred behind the
# threshold so the common (non-looping) path stays free.
if tool_name == "ToolSearch":
_ts_loop["n"] += 1
if _ts_loop["n"] >= TOOLSEARCH_LOOP_THRESHOLD:
_gated = gated_mcp_server_names(session.allowed_tools, session.active_mcps)
_reason = toolsearch_loop_redirect(_ts_loop["n"], _gated)
if _reason:
logger.info(f"[MCP-DEBUG] ToolSearch loop-breaker fired for {session_id} (n={_ts_loop['n']})")
# 2B-MCP: also surface a one-click connect offer to the USER for the vetted
# gated servers the agent keeps reaching for. Suggest-only: this just shows a
# card on the same channel the preflight uses; activation still requires
# MCPActivate + the dispatch gate, so it opens no side channel. Once per run,
# fail-open (an offer hiccup must never block the agent).
if not _mcp_offer_sent["done"]:
try:
from backend.apps.agents.core.mcp_preflight import offer_for_gated_server
_s = load_settings()
_offers = [o for o in (offer_for_gated_server(n, _s) for n in _gated) if o]
if _offers:
_mcp_offer_sent["done"] = True
await ws_manager.send_to_session(session_id, "agent:mcp_suggestions", {
"session_id": session_id,
"suggestions": _offers,
"is_vague": False,
})
except Exception:
logger.debug("mid-run MCP connect offer skipped", exc_info=True)
return {
"hookSpecificOutput": {
"hookEventName": hook_event,
"permissionDecision": "deny",
"permissionDecisionReason": _reason,
}
}
else:
_ts_loop["n"] = 0
# MCPSearch is the agent saying "I need an integration I don't have" (e.g. "no email
# connected"). Don't make the user read a wall of options: fire the same curated connect
# card the launch preflight uses, keyed to their original request. Non-blocking (the search
# proceeds) and once per run; covers the common path the ToolSearch-loop branch misses
# because a capable model does one MCPSearch instead of thrashing. Suggest-only as ever.
if (tool_name.endswith("MCPSearch") or tool_name.endswith("MCPList")) and not _mcp_offer_sent["done"]:
_mcp_offer_sent["done"] = True
async def _offer_from_prompt():
try:
from backend.apps.agents.core.mcp_preflight import run_preflight
result = await run_preflight(prompt, task_id=session_id, require_vague=False)
offers = result.get("suggestions", [])
if offers:
await ws_manager.send_to_session(session_id, "agent:mcp_suggestions", {
"session_id": session_id,
"suggestions": offers,
"is_vague": False,
})
except Exception:
logger.debug("MCPSearch-triggered connect offer skipped", exc_info=True)
asyncio.create_task(_offer_from_prompt())
if tool_name and tool_name != "AskUserQuestion":
tool_input = input_data.get("tool_input", {})
policy, sensitive_pattern = path_gate.maybe_override_policy(
effective_policy(tool_name, builtin_perms, policy_defaults), tool_name, tool_input
)
if policy == "deny":
return {
"hookSpecificOutput": {
"hookEventName": hook_event,
"permissionDecision": "deny",
"permissionDecisionReason": "Tool denied by permission policy",
}
}
if policy == "ask":
decision = await request_user_approval(session, session_id, tool_name, tool_input, builtin_perms, sensitive_pattern=sensitive_pattern)
if decision.get("behavior") == "allow":
if tool_use_id:
tool_start_times[tool_use_id] = time.time()
return {
"hookSpecificOutput": {
"hookEventName": hook_event,
"permissionDecision": "allow",
}
}
return {
"hookSpecificOutput": {
"hookEventName": hook_event,
"permissionDecision": "deny",
"permissionDecisionReason": decision.get("message", "User denied this action"),
}
}
if tool_use_id:
tool_start_times[tool_use_id] = time.time()
return {}
return await gate_hooks.pre_tool_hook(hook_ctx, input_data, tool_use_id, context)
async def post_tool_hook(input_data, tool_use_id, context):
elapsed_ms = None
if tool_use_id and tool_use_id in tool_start_times:
elapsed_ms = int((time.time() - tool_start_times.pop(tool_use_id)) * 1000)
if tool_use_id and tool_use_id in hook_ctx.tool_start_times:
elapsed_ms = int((time.time() - hook_ctx.tool_start_times.pop(tool_use_id)) * 1000)
raw_response = input_data.get("tool_response", "")
backend/apps/agents/manager/permissions/gate_hooks.py
+166
View File
@@ -0,0 +1,166 @@
"""The SDK permission/pre-tool hooks, lifted out of the agent loop. can_use_tool answers the
SDK's permission callback; pre_tool_hook runs before each tool call and also carries the two
MCP loop-breakers (ToolSearch-thrash redirect + one-shot connect offer). Both operate on a
HookContext passed by reference, so the shared counters survive across calls. The dict returns
are the claude_agent_sdk hook protocol (hookSpecificOutput), not internal state."""
import asyncio
import logging
import time
from typing import Dict, Optional, Union
from typeguard import typechecked
from claude_agent_sdk import PermissionResultAllow, PermissionResultDeny
from backend.apps.agents.core.ws_manager import ws_manager
from backend.apps.settings.settings import load_settings
from backend.apps.agents.manager.permissions import path_gate
from backend.apps.agents.manager.permissions.decision import effective_policy, request_user_approval
from backend.apps.agents.manager.prompt.tool_catalog import gated_mcp_server_names
from backend.apps.agents.manager.prompt.prompt_context import (
TOOLSEARCH_LOOP_THRESHOLD,
toolsearch_loop_redirect,
)
from backend.apps.agents.manager.streaming.hook_context import HookContext
logger = logging.getLogger(__name__)
@typechecked
async def can_use_tool(
ctx: HookContext, tool_name: str, input_data: object, context: object
) -> Union[PermissionResultAllow, PermissionResultDeny]:
sensitive_pattern: Optional[str] = None
if tool_name != "AskUserQuestion":
policy, sensitive_pattern = path_gate.maybe_override_policy(
effective_policy(tool_name, ctx.builtin_perms, ctx.policy_defaults), tool_name, input_data
)
if policy == "always_allow":
return PermissionResultAllow(updated_input=input_data)
if policy == "deny":
return PermissionResultDeny(message="Tool denied by permission policy")
decision = await request_user_approval(
ctx.session, ctx.session_id, tool_name, input_data, ctx.builtin_perms, sensitive_pattern=sensitive_pattern
)
if decision.get("behavior") == "allow":
return PermissionResultAllow(
updated_input=decision.get("updated_input", input_data)
)
return PermissionResultDeny(
message=decision.get("message", "User denied this action")
)
@typechecked
async def pre_tool_hook(ctx: HookContext, input_data: dict, tool_use_id: Optional[str], context: object) -> Dict[str, object]:
tool_name = input_data.get("tool_name", "")
hook_event = input_data.get("hook_event_name", "PreToolUse")
# ToolSearch loop-breaker. Gated MCP servers are withheld from the
# SDK until MCPActivate, so the CLI's native ToolSearch can never
# find them; small models thrash (empty ToolSearch, retry) for
# minutes until the user pauses. Let the first couple through, then
# redirect to the gate. Any non-ToolSearch call is real progress, so
# the counter resets. Gated-server lookup is deferred behind the
# threshold so the common (non-looping) path stays free.
if tool_name == "ToolSearch":
ctx.ts_loop_count += 1
if ctx.ts_loop_count >= TOOLSEARCH_LOOP_THRESHOLD:
gated = gated_mcp_server_names(ctx.session.allowed_tools, ctx.session.active_mcps)
reason = toolsearch_loop_redirect(ctx.ts_loop_count, gated)
if reason:
logger.info(f"[MCP-DEBUG] ToolSearch loop-breaker fired for {ctx.session_id} (n={ctx.ts_loop_count})")
# 2B-MCP: also surface a one-click connect offer to the USER for the vetted
# gated servers the agent keeps reaching for. Suggest-only: this just shows a
# card on the same channel the preflight uses; activation still requires
# MCPActivate + the dispatch gate, so it opens no side channel. Once per run,
# fail-open (an offer hiccup must never block the agent).
if not ctx.mcp_offer_sent:
try:
from backend.apps.agents.core.mcp_preflight import offer_for_gated_server
settings = load_settings()
offers = [o for o in (offer_for_gated_server(n, settings) for n in gated) if o]
if offers:
ctx.mcp_offer_sent = True
await ws_manager.send_to_session(ctx.session_id, "agent:mcp_suggestions", {
"session_id": ctx.session_id,
"suggestions": offers,
"is_vague": False,
})
except Exception:
logger.debug("mid-run MCP connect offer skipped", exc_info=True)
return {
"hookSpecificOutput": {
"hookEventName": hook_event,
"permissionDecision": "deny",
"permissionDecisionReason": reason,
}
}
else:
ctx.ts_loop_count = 0
# MCPSearch is the agent saying "I need an integration I don't have" (e.g. "no email
# connected"). Don't make the user read a wall of options: fire the same curated connect
# card the launch preflight uses, keyed to their original request. Non-blocking (the search
# proceeds) and once per run; covers the common path the ToolSearch-loop branch misses
# because a capable model does one MCPSearch instead of thrashing. Suggest-only as ever.
if (tool_name.endswith("MCPSearch") or tool_name.endswith("MCPList")) and not ctx.mcp_offer_sent:
ctx.mcp_offer_sent = True
async def offer_from_prompt():
try:
from backend.apps.agents.core.mcp_preflight import run_preflight
result = await run_preflight(ctx.prompt, task_id=ctx.session_id, require_vague=False)
offers = result.get("suggestions", [])
if offers:
await ws_manager.send_to_session(ctx.session_id, "agent:mcp_suggestions", {
"session_id": ctx.session_id,
"suggestions": offers,
"is_vague": False,
})
except Exception:
logger.debug("MCPSearch-triggered connect offer skipped", exc_info=True)
asyncio.create_task(offer_from_prompt())
if tool_name and tool_name != "AskUserQuestion":
tool_input = input_data.get("tool_input", {})
policy, sensitive_pattern = path_gate.maybe_override_policy(
effective_policy(tool_name, ctx.builtin_perms, ctx.policy_defaults), tool_name, tool_input
)
if policy == "deny":
return {
"hookSpecificOutput": {
"hookEventName": hook_event,
"permissionDecision": "deny",
"permissionDecisionReason": "Tool denied by permission policy",
}
}
if policy == "ask":
decision = await request_user_approval(
ctx.session, ctx.session_id, tool_name, tool_input, ctx.builtin_perms, sensitive_pattern=sensitive_pattern
)
if decision.get("behavior") == "allow":
if tool_use_id:
ctx.tool_start_times[tool_use_id] = time.time()
return {
"hookSpecificOutput": {
"hookEventName": hook_event,
"permissionDecision": "allow",
}
}
return {
"hookSpecificOutput": {
"hookEventName": hook_event,
"permissionDecision": "deny",
"permissionDecisionReason": decision.get("message", "User denied this action"),
}
}
if tool_use_id:
ctx.tool_start_times[tool_use_id] = time.time()
return {}
backend/apps/agents/manager/streaming/hook_context.py
+27
View File
@@ -0,0 +1,27 @@
"""Captured per-run state for the SDK tool hooks (can_use_tool / pre / post). Passed by
reference into the extracted hook functions so they can mutate the shared counters without
living inside the agent loop's closure. The session reference is the SAME object the loop
holds (pydantic keeps the instance, doesn't copy it), so hook-side mutations to status /
pending_approvals are visible to the loop."""
from typing import Dict
from pydantic import BaseModel, ConfigDict
from backend.apps.agents.core.models import AgentSession
class HookContext(BaseModel):
model_config = ConfigDict(validate_assignment=True)
session: AgentSession
session_id: str
prompt: str
builtin_perms: Dict[str, str]
policy_defaults: Dict[str, str]
# tool_use_id -> wall-clock start (s); pre records it, post pops it for elapsed_ms.
tool_start_times: Dict[str, float] = {}
# Consecutive ToolSearch calls; a run of these is the "looping on ToolSearch" wedge.
ts_loop_count: int = 0
# One mid-run "connect this MCP" card per run; a stuck agent retries, the user sees it once.
mcp_offer_sent: bool = False
backend/tests/test_gate_hooks.py
+72
View File
@@ -0,0 +1,72 @@
"""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.hook_context import HookContext
from backend.apps.agents.manager.permissions import gate_hooks
from backend.apps.agents.manager.prompt.prompt_context import TOOLSEARCH_LOOP_THRESHOLD
def _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={},
)
@pytest.mark.asyncio
async def test_can_use_tool_always_allow_returns_allow():
ctx = _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 = _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 = _ctx()
with patch.object(gate_hooks.path_gate, "maybe_override_policy", return_value=("ask", None)), \
patch.object(gate_hooks, "request_user_approval", new=AsyncMock(return_value={"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 = _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 = _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