From c2d0c4d28e19902450de57713028a602e4102def Mon Sep 17 00:00:00 2001 From: haikdc Date: Thu, 2 Apr 2026 22:46:54 -0700 Subject: [PATCH] [Haik]: ckpt, sdk hooks & approval flow don --- backend/apps/agents/agents.py | 54 ++++++- backend/apps/agents/create_sdk_hooks.py | 139 +++++++++++++++++++ backend/core/Agent/Agent.py | 56 +++++++- backend/core/events/events.py | 19 ++- backend/core/tools/shared_structs/Toolkit.py | 20 ++- 5 files changed, 277 insertions(+), 11 deletions(-) create mode 100644 backend/apps/agents/create_sdk_hooks.py diff --git a/backend/apps/agents/agents.py b/backend/apps/agents/agents.py index 4ac47aa7..21a2b170 100644 --- a/backend/apps/agents/agents.py +++ b/backend/apps/agents/agents.py @@ -15,14 +15,14 @@ import asyncio from fastapi import HTTPException from pydantic import BaseModel -from typing import Optional, List +from typing import Optional, List, Dict from backend.config.Apps import SubApp from backend.core.Agent.Agent import Agent from backend.core.shared_structs.agent.Message.Message import UserMessage from backend.core.events.events import ( AnyEvent, AgentStatusEvent, AgentClosedEvent, BranchSwitchedEvent, - EventCallback, + ApprovalRequestEvent, EventCallback, ) from backend.apps.agents.session_store import ( load_all, @@ -36,15 +36,42 @@ from backend.apps.agents.session_store import ( from backend.apps.agents import ws from backend.apps.agents.compose_system_prompt import compose_system_prompt from backend.core.tools.make_builtin_toolkit.make_builtin_toolkit import make_builtin_toolkit +from backend.apps.agents.create_sdk_hooks import create_sdk_hooks from claude_agent_sdk import ClaudeAgentOptions +from claude_agent_sdk.types import HookMatcher +from backend.core.tools.shared_structs.Toolkit import Toolkit +from claude_agent_sdk.types import McpServerConfig SESSIONS: dict[str, Agent] = {} @typechecked def p_make_session_emitter(session_id: str) -> EventCallback: - """Create an event callback that routes typed events to the WS connection pool.""" + """Create an event callback that routes typed events to the WS connection pool. + + ApprovalRequestEvents are special-cased: instead of just broadcasting, + the emitter routes through the APPROVAL_BRIDGE and resolves the + embedded future with the user's decision. + """ async def emit(event: AnyEvent) -> None: + if isinstance(event, ApprovalRequestEvent): + if not ws.has_global_connections(): + if not event.future.done(): + event.future.set_result({"behavior": "deny", "message": "No dashboard connected for approval."}) + return + result = await ws.APPROVAL_BRIDGE.request( + request_id=event.request_id, + send_fn=lambda: ws.send_to_session(session_id, event.event, { + "request_id": event.request_id, + "session_id": event.session_id, + "tool_name": event.tool_name, + "tool_input": event.tool_input, + }), + timeout=600.0, + ) + if not event.future.done(): + event.future.set_result(result) + return await ws.send_to_session(session_id, event.event, event.model_dump(mode="json")) return emit @@ -68,6 +95,7 @@ async def p_send_browser_command( timeout=30.0, ) + def get_agent(session_id: str) -> Agent: agent: Optional[Agent] = SESSIONS.get(session_id) if not agent: @@ -91,6 +119,8 @@ async def agents_lifespan(): agent: Agent = Agent(**data) agent.status = "stopped" agent.on_event = p_make_session_emitter(agent.session_id) + toolkit: Toolkit = make_builtin_toolkit(agent, SESSIONS, p_send_browser_command) + agent.toolkit = toolkit SESSIONS[agent.session_id] = agent delete(sid) except Exception as e: @@ -135,7 +165,6 @@ async def launch(body: LaunchBody) -> dict: system_prompt = compose_system_prompt( session_prompt=body.system_prompt or None, ) - # Placeholder config — replaced below once the toolkit is built agent: Agent = Agent( model=body.model, mode=body.mode, @@ -145,16 +174,25 @@ async def launch(body: LaunchBody) -> dict: agent.on_event = p_make_session_emitter(agent.session_id) SESSIONS[agent.session_id] = agent - toolkit = make_builtin_toolkit(agent, SESSIONS, p_send_browser_command) - mcp_servers = toolkit.collect_mcp_servers() + toolkit: Toolkit = make_builtin_toolkit(agent, SESSIONS, p_send_browser_command) + agent.toolkit = toolkit + mcp_servers: Dict[str, McpServerConfig] = toolkit.collect_mcp_servers() allowed_tools, disallowed_tools = toolkit.collect_tool_permissions() + can_use_tool, pre_tool_hook, post_tool_hook = create_sdk_hooks(agent) + agent.config = ClaudeAgentOptions( system_prompt=system_prompt, max_turns=body.max_turns, mcp_servers=mcp_servers if mcp_servers else None, allowed_tools=allowed_tools, disallowed_tools=disallowed_tools, + permission_mode="default", + can_use_tool=can_use_tool, + hooks={ + "PreToolUse": [HookMatcher(matcher=None, hooks=[pre_tool_hook])], + "PostToolUse": [HookMatcher(matcher=None, hooks=[post_tool_hook])], + }, ) await agent.emit(AgentStatusEvent( @@ -317,6 +355,8 @@ async def resume_session(session_id: str) -> dict: agent: Agent = Agent(**data) agent.status = "stopped" agent.on_event = p_make_session_emitter(agent.session_id) + toolkit: Toolkit = make_builtin_toolkit(agent, SESSIONS, p_send_browser_command) + agent.toolkit = toolkit SESSIONS[agent.session_id] = agent delete(session_id) await agent.emit(AgentStatusEvent( @@ -345,6 +385,8 @@ async def duplicate_session(session_id: str, body: dict = {}) -> dict: clone.pending_approvals = [] clone.sub_agents = [] clone.on_event = p_make_session_emitter(clone.session_id) + toolkit: Toolkit = make_builtin_toolkit(clone, SESSIONS, p_send_browser_command) + clone.toolkit = toolkit SESSIONS[clone.session_id] = clone await clone.emit(AgentStatusEvent( session_id=clone.session_id, status=clone.status, diff --git a/backend/apps/agents/create_sdk_hooks.py b/backend/apps/agents/create_sdk_hooks.py new file mode 100644 index 00000000..a05362a7 --- /dev/null +++ b/backend/apps/agents/create_sdk_hooks.py @@ -0,0 +1,139 @@ +"""SDK hook factories for the agent loop. + +create_sdk_hooks(agent) returns the three callables that ClaudeAgentOptions +expects: can_use_tool, pre_tool_hook, post_tool_hook. + +All transport-specific behavior (WebSocket, persistence) is accessed +through Agent.emit (on_event) so this module stays framework-agnostic. +""" + +import json +import time +from typing import Any, Callable, Dict, Tuple, Optional, List + +from claude_agent_sdk.types import PermissionResultAllow, PermissionResultDeny, PermissionResult + +from backend.core.shared_structs.agent.Message.Message import ToolResultMessage +from backend.core.shared_structs.agent.Message.agent_outputs import ToolResultContent +from backend.core.events.events import AgentMessageEvent + +from backend.core.Agent.Agent import Agent +from typeguard import typechecked +from backend.core.tools.shared_structs.TOOL_PERMISSIONS import TOOL_PERMISSIONS + + +@typechecked +def create_sdk_hooks( + agent: "Agent", +) -> Tuple[Callable, Callable, Callable]: + """Build (can_use_tool, pre_tool_hook, post_tool_hook) closures for an Agent.""" + + tool_start_times: Dict[str, float] = {} + + @typechecked + async def can_use_tool(tool_name: str, input_data: Any) -> PermissionResult: + permission: Optional[TOOL_PERMISSIONS] = ( + agent.toolkit.resolve_permission(tool_name) if agent.toolkit else None + ) + if permission == "allow": + return PermissionResultAllow(updated_input=input_data) + if permission == "deny": + return PermissionResultDeny(message="Tool denied by permission policy") + + # TODO: better type spec for decision + decision: Dict[str, Any] = await agent.request_approval( + tool_name, input_data if isinstance(input_data, dict) else {}, + ) + 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"), + ) + + # TODO: better type spec for input_data and return value + @typechecked + async def pre_tool_hook(input_data: dict, tool_use_id: str) -> Dict[str, Any]: + tool_name: str = input_data.get("tool_name", "") + hook_event: str = input_data.get("hook_event_name", "PreToolUse") + + if tool_name: + permission: Optional[TOOL_PERMISSIONS] = ( + agent.toolkit.resolve_permission(tool_name) if agent.toolkit else None + ) + if permission == "deny": + return { + "hookSpecificOutput": { + "hookEventName": hook_event, + "permissionDecision": "deny", + "permissionDecisionReason": "Tool denied by permission policy", + }, + } + if permission == "ask": + tool_input: Dict[str, Any] = input_data.get("tool_input", {}) + decision: Dict[str, Any] = await agent.request_approval(tool_name, tool_input) + 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 {} + + @typechecked + async def post_tool_hook(input_data: dict, tool_use_id: str) -> Dict[str, Any]: + elapsed_ms: Optional[int] = 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) + + raw_response = input_data.get("tool_response", "") + + if isinstance(raw_response, list) and raw_response: + text_parts: List[str] = [ + b.get("text", "") + for b in raw_response + if isinstance(b, dict) and b.get("type") == "text" + ] + if text_parts: + raw_response: str = "\n".join(text_parts) if len(text_parts) > 1 else text_parts[0] + + if isinstance(raw_response, str): + content: str = raw_response + else: + try: + content: str = json.dumps(raw_response, indent=2, default=str) + except Exception: + content: str = str(raw_response) + + result_msg: ToolResultMessage = ToolResultMessage( + content=ToolResultContent( + tool_use_id=tool_use_id or "", + text=content, + is_error=isinstance(raw_response, str) and raw_response.startswith("Error"), + ), + branch_id=agent.branch_id, + ) + agent.messages.append(result_msg) + await agent.emit(AgentMessageEvent( + session_id=agent.session_id, message=result_msg, + )) + + return {"continue_": True} + + return can_use_tool, pre_tool_hook, post_tool_hook diff --git a/backend/core/Agent/Agent.py b/backend/core/Agent/Agent.py index 1bd6a4a6..bc0fec8e 100644 --- a/backend/core/Agent/Agent.py +++ b/backend/core/Agent/Agent.py @@ -2,11 +2,11 @@ import asyncio import logging import os from copy import deepcopy +from typing import Any, Dict, List, Literal, Optional from uuid import uuid4 from claude_agent_sdk import ClaudeAgentOptions from pydantic import BaseModel, Field, InstanceOf -from typing import List, Literal, Optional from typeguard import typechecked from backend.core.Agent.run_agent_loop.run_agent_loop import run_agent_loop @@ -15,8 +15,9 @@ from backend.core.shared_structs.agent.ApprovalRequest import ApprovalRequest from backend.core.shared_structs.agent.MessageLog import MessageLog from backend.core.events.events import ( AgentSnapshot, AgentStatusEvent, AgentMessageEvent, - EventCallback, AnyEvent, + ApprovalRequestEvent, EventCallback, AnyEvent, ) +from backend.core.tools.shared_structs.Toolkit import Toolkit os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000") @@ -39,6 +40,7 @@ class Agent(BaseModel): sub_branches: List["Agent"] = Field(default_factory=list) parent_id: Optional[str] = None + toolkit: Optional[Toolkit] = Field(default=None, exclude=True) on_event: Optional[EventCallback] = Field(default=None, exclude=True) task: Optional[asyncio.Task] = None @@ -70,6 +72,55 @@ class Agent(BaseModel): if self.on_event: await self.on_event(event) + @typechecked + async def request_approval( + self, tool_name: str, tool_input: Dict[str, Any], + ) -> Dict[str, Any]: + """HITL approval flow: pause the agent, ask the user, resume. + + Emits an ApprovalRequestEvent through on_event. The transport layer + resolves the embedded future with the user's decision. + Returns {"behavior": "allow"|"deny", ...}. + """ + if not self.on_event: + return {"behavior": "allow"} + + request: ApprovalRequest = ApprovalRequest( + session_id=self.session_id, + tool_name=tool_name, + tool_input=tool_input, + ) + self.pending_approvals.append(request) + self.status = "waiting_approval" + await self.emit(AgentStatusEvent( + session_id=self.session_id, status="waiting_approval", + )) + + future: asyncio.Future = asyncio.get_event_loop().create_future() + try: + await self.emit(ApprovalRequestEvent( + session_id=self.session_id, + request_id=request.id, + tool_name=tool_name, + tool_input=tool_input, + future=future, + )) + decision: Dict[str, Any] = await future + except asyncio.TimeoutError: + decision = {"behavior": "deny", "message": "Approval timed out"} + except asyncio.CancelledError: + decision = {"behavior": "deny", "message": "Agent stopped"} + raise + + self.pending_approvals = [ + a for a in self.pending_approvals if a.id != request.id + ] + self.status = "running" + await self.emit(AgentStatusEvent( + session_id=self.session_id, status="running", + )) + return decision + @typechecked async def send_message(self, msg: Message) -> None: async with self.lock: @@ -126,6 +177,7 @@ class Agent(BaseModel): "messages": MessageLog(messages=branched_messages), "sub_agents": [], "pending_approvals": [], + "toolkit": self.toolkit, "on_event": self.on_event, "task": None, "lock": asyncio.Lock(), diff --git a/backend/core/events/events.py b/backend/core/events/events.py index 371e5ac0..42ac648c 100644 --- a/backend/core/events/events.py +++ b/backend/core/events/events.py @@ -1,5 +1,7 @@ +import asyncio + from pydantic import BaseModel, Field -from typing import Annotated, List, Literal, Optional, Union, Callable, Awaitable +from typing import Annotated, Any, Dict, Literal, Optional, Union, Callable, Awaitable from backend.core.shared_structs.agent.Message.Message import AnyMessage from backend.core.shared_structs.agent.AgentSnapshot import AgentSnapshot @@ -62,12 +64,25 @@ class BrowserCardAddedEvent(BaseModel): browser_card: BrowserCardPosition +class ApprovalRequestEvent(BaseModel): + """Emitted when the agent needs human approval for a tool call. + + The transport layer resolves `future` with the user's decision dict. + """ + event: Literal["agent:approval_request"] = "agent:approval_request" + session_id: str + request_id: str + tool_name: str + tool_input: Dict[str, Any] + future: asyncio.Future = Field(exclude=True) + + AnyEvent = Annotated[ Union[ AgentStatusEvent, AgentMessageEvent, StreamStartEvent, StreamDeltaEvent, StreamEndEvent, BranchSwitchedEvent, AgentClosedEvent, - BrowserCardAddedEvent, + BrowserCardAddedEvent, ApprovalRequestEvent, ], Field(discriminator="event"), ] diff --git a/backend/core/tools/shared_structs/Toolkit.py b/backend/core/tools/shared_structs/Toolkit.py index 2f22824a..5b5a6180 100644 --- a/backend/core/tools/shared_structs/Toolkit.py +++ b/backend/core/tools/shared_structs/Toolkit.py @@ -88,4 +88,22 @@ class Toolkit(BaseModel): a, d = toolkit.collect_tool_permissions() allowed.extend(a) disallowed.extend(d) - return allowed, disallowed \ No newline at end of file + return allowed, disallowed + + @typechecked + def resolve_permission(self, sdk_name: str) -> Optional[TOOL_PERMISSIONS]: + """Look up the permission for a single tool by its SDK-format name. + + Returns the tool's permission if found, or None if the tool + doesn't exist in this toolkit tree. + """ + if self.tools is not None: + for tool in self.tools: + if tool.to_sdk_args() == sdk_name: + return tool.permission + if self.nested_toolkits is not None: + for toolkit in self.nested_toolkits: + found: Optional[TOOL_PERMISSIONS] = toolkit.resolve_permission(sdk_name) + if found is not None: + return found + return None \ No newline at end of file