"""Owned agent loop — replaces claude_agent_sdk's query() function. Generalizes the pattern from browser_agent.py (lines 243-334) into a provider-agnostic, streaming, HITL-aware tool-use loop. """ from __future__ import annotations import json import logging import time from typing import Any, Callable, Awaitable from uuid import uuid4 from backend.apps.agents.providers.base import ( BaseProvider, ContentBlock, ModelResponse, ProviderMessage, StreamEvent, ToolCall, ToolSchema, ) logger = logging.getLogger(__name__) # Type aliases for callbacks ToolExecutor = Callable[[str, dict], Awaitable[list[dict]]] # hitl_handler(tool_name, tool_input) -> (approved, updated_input_or_None) HITLHandler = Callable[[str, dict], Awaitable[tuple[bool, dict | None]]] # ws_emitter(event_type, data) -> None WSEmitter = Callable[[str, dict], Awaitable[None]] class AgentLoop: """Provider-agnostic agent loop with streaming and HITL support. The loop: 1. Sends user message to the model 2. Streams the response (emitting WebSocket events) 3. If the model requests tool use: a. For each tool call: check HITL permission → execute → collect result b. Append tool results → go to step 2 4. If the model stops (end_turn/max_tokens): done """ def __init__( self, session_id: str, provider: BaseProvider, model: str, system_prompt: str | None, tools: list[ToolSchema], tool_executor: ToolExecutor, hitl_handler: HITLHandler, ws_emitter: WSEmitter, max_turns: int | None = None, cwd: str | None = None, ): self.session_id = session_id self.provider = provider self.model = model self.system_prompt = system_prompt self.tools = tools self.tool_executor = tool_executor self.hitl_handler = hitl_handler self.ws_emitter = ws_emitter self.max_turns = max_turns self.cwd = cwd # Conversation history in provider-agnostic format self.messages: list[ProviderMessage] = [] # Token tracking self.total_input_tokens = 0 self.total_output_tokens = 0 async def run(self, user_content: Any) -> None: """Run the agent loop for a single user turn.""" # Append user message user_msg = self.provider.format_user_message(user_content) self.messages.append(user_msg) turn = 0 while True: if self.max_turns and turn >= self.max_turns: logger.info(f"Agent {self.session_id}: max turns ({self.max_turns}) reached") break turn += 1 # Stream the model response and collect it response = await self._stream_and_collect() # Track usage self.total_input_tokens += response.usage.get("input_tokens", 0) self.total_output_tokens += response.usage.get("output_tokens", 0) # Append assistant message to conversation history assistant_msg = self.provider.format_assistant_message(response) self.messages.append(assistant_msg) # If no tool use, we're done if response.stop_reason != "tool_use": break # Execute tools tool_results = await self._execute_tools(response) if not tool_results: break # Append tool results self.messages.append(ProviderMessage(role="tool_result", content=tool_results)) async def _stream_and_collect(self) -> ModelResponse: """Stream model output, emit WebSocket events, collect full response.""" collected_content: list[ContentBlock] = [] collected_usage: dict[str, int] = {} stop_reason = "end_turn" # Track streaming state for WS emissions stream_text_msg_id: str | None = None stream_tool_msg_ids: dict[int, str] = {} # block index -> msg_id block_index_map: dict[int, str] = {} # block index -> msg_id # Buffers for collecting content text_buffers: dict[int, str] = {} json_buffers: dict[int, str] = {} tool_names: dict[int, str] = {} tool_ids: dict[int, str] = {} block_types: dict[int, str] = {} # Wall-clock start time per content block (server-side stamps). # Used to compute elapsed_ms for thinking blocks so the persisted # ThinkingBubble can show the duration after streaming ends. block_start_ts: dict[int, float] = {} thinking_total_ms: int = 0 thinking_total_chars: int = 0 async for event in self.provider.stream_message( model=self.model, system=self.system_prompt, messages=self.messages, tools=self.tools, ): if event.type == "content_block_start": if event.block_type == "text": if stream_text_msg_id is None: stream_text_msg_id = uuid4().hex await self.ws_emitter("agent:stream_start", { "message_id": stream_text_msg_id, "role": "assistant", }) block_index_map[event.index] = stream_text_msg_id block_types[event.index] = "text" text_buffers[event.index] = "" elif event.block_type == "tool_use": tool_msg_id = uuid4().hex stream_tool_msg_ids[event.index] = tool_msg_id block_index_map[event.index] = tool_msg_id block_types[event.index] = "tool_use" tool_names[event.index] = event.tool_name tool_ids[event.index] = event.tool_id json_buffers[event.index] = "" await self.ws_emitter("agent:stream_start", { "message_id": tool_msg_id, "role": "tool_call", "tool_name": event.tool_name, }) elif event.block_type == "thinking": # Extended-thinking content block. Emit a distinct # WS stream with role="thinking" so the frontend # renders the live ThinkingBubble pill (rising # token counter, auto-collapse on first text). Each # thinking block gets its own message id — multiple # interleaved thinking/text blocks remain # individually addressable. thinking_msg_id = uuid4().hex block_index_map[event.index] = thinking_msg_id block_types[event.index] = "thinking" text_buffers[event.index] = "" # Server-stamp the start so we can compute exact # elapsed_ms server-side at content_block_stop. Using # time.time() (not monotonic) is fine here — we only # subtract two values from the same clock. block_start_ts[event.index] = time.time() await self.ws_emitter("agent:stream_start", { "message_id": thinking_msg_id, "role": "thinking", }) elif event.type == "content_block_delta": msg_id = block_index_map.get(event.index) if not msg_id: continue if event.delta_type == "text_delta": text_buffers.setdefault(event.index, "") text_buffers[event.index] += event.text await self.ws_emitter("agent:stream_delta", { "message_id": msg_id, "delta": event.text, }) elif event.delta_type == "input_json_delta": json_buffers.setdefault(event.index, "") json_buffers[event.index] += event.text await self.ws_emitter("agent:stream_delta", { "message_id": msg_id, "delta": event.text, }) elif event.delta_type == "thinking_delta": # Reuse the text buffer for thinking — same shape # (accumulated str), different sink. text_buffers.setdefault(event.index, "") text_buffers[event.index] += event.text await self.ws_emitter("agent:stream_delta", { "message_id": msg_id, "delta": event.text, }) elif event.type == "content_block_stop": msg_id = block_index_map.get(event.index) bt = block_types.get(event.index, "") if bt == "text": collected_content.append( ContentBlock(type="text", text=text_buffers.get(event.index, "")) ) elif bt == "tool_use": try: tool_input = json.loads(json_buffers.get(event.index, "{}")) except json.JSONDecodeError: tool_input = {} collected_content.append(ContentBlock( type="tool_use", tool_call=ToolCall( id=tool_ids.get(event.index, uuid4().hex), name=tool_names.get(event.index, ""), input=tool_input, ), )) elif bt == "thinking": thinking_text = text_buffers.get(event.index, "") collected_content.append( ContentBlock(type="thinking", text=thinking_text) ) # Accumulate per-block duration + char count for the # eventual persisted Message. We sum across multiple # thinking blocks in the same turn so a complex # interleaved (think → tool → think → answer) turn # still reports total time spent reasoning. start_ts = block_start_ts.get(event.index) if start_ts is not None: thinking_total_ms += int((time.time() - start_ts) * 1000) thinking_total_chars += len(thinking_text) # Send stream_end for tool + thinking blocks (text block # ends at message_stop). Thinking ends here so the # frontend can transition the pill from "live" to # "Thought for Ns" the moment the model stops thinking, # even if it then keeps streaming text. if msg_id and (bt == "tool_use" or bt == "thinking"): payload: dict[str, Any] = {"message_id": msg_id} if bt == "thinking": # Server-stamped truth so the persisted bubble # doesn't fall back to "Thoughts" — and so the # live bubble freezes on the exact server-side # duration instead of the client's clock. block_start = block_start_ts.get(event.index) if block_start is not None: block_elapsed = int((time.time() - block_start) * 1000) payload["elapsed_ms"] = block_elapsed # Token estimate for THIS block (chars/3.6 ≈ # Anthropic BPE for English prose). Matches the # heuristic the live UI used so the freeze # value doesn't visually jump. block_text = text_buffers.get(event.index, "") if block_text: payload["tokens"] = max(1, round(len(block_text) / 3.6)) await self.ws_emitter("agent:stream_end", payload) elif event.type == "usage": # Accumulate token usage from provider stream for k, v in event.usage.items(): collected_usage[k] = collected_usage.get(k, 0) + v elif event.type == "message_stop": # Check if any tool calls means stop_reason is tool_use has_tool_use = any(b.type == "tool_use" for b in collected_content) if has_tool_use: stop_reason = "tool_use" # End text stream if stream_text_msg_id: await self.ws_emitter("agent:stream_end", { "message_id": stream_text_msg_id, }) # Build and emit the collected messages await self._emit_collected_messages( collected_content, stream_text_msg_id, stream_tool_msg_ids, thinking_elapsed_ms=thinking_total_ms, thinking_total_chars=thinking_total_chars, ) return ModelResponse( content=collected_content, stop_reason=stop_reason, usage=collected_usage, ) async def _emit_collected_messages( self, content: list[ContentBlock], text_msg_id: str | None, tool_msg_ids: dict[int, str], thinking_elapsed_ms: int = 0, thinking_total_chars: int = 0, ) -> None: """Emit finalized agent:message events for the collected response.""" from backend.apps.agents.models import Message # Emit thinking blocks (extended thinking). Persisted as their own # messages so a session reload still shows the reasoning trail. # Multiple thinking blocks per turn are concatenated into a single # persisted message — the streaming UI already showed each block # individually, this is just for the historical record. thinking_parts = [b.text for b in content if b.type == "thinking" and b.text] if thinking_parts: joined = "\n\n".join(thinking_parts) # Stamp duration + token estimate so the persisted bubble can # show "Thought for Ns · M tokens" on reload instead of the # generic "Thoughts" fallback. Use the server-side accumulated # times so multi-block turns aggregate correctly. msg = Message( role="thinking", content=joined, elapsed_ms=thinking_elapsed_ms or None, tokens=max(1, round(thinking_total_chars / 3.6)) if thinking_total_chars else None, ) await self.ws_emitter("agent:message", { "message": msg.model_dump(mode="json"), }) # Emit text message text_parts = [b.text for b in content if b.type == "text" and b.text] if text_parts: msg = Message( id=text_msg_id or uuid4().hex, role="assistant", content="\n".join(text_parts), ) await self.ws_emitter("agent:message", { "message": msg.model_dump(mode="json"), }) # Emit tool call messages tool_blocks = [b for b in content if b.type == "tool_use" and b.tool_call] tool_id_list = sorted(tool_msg_ids.items(), key=lambda x: x[0]) for i, block in enumerate(tool_blocks): tc = block.tool_call msg_id = tool_id_list[i][1] if i < len(tool_id_list) else uuid4().hex msg = Message( id=msg_id, role="tool_call", content={ "id": tc.id, "tool": tc.name, "input": tc.input, }, ) await self.ws_emitter("agent:message", { "message": msg.model_dump(mode="json"), }) async def _execute_tools(self, response: ModelResponse) -> list[dict]: """Execute all tool calls from a response, respecting HITL permissions. Returns a list of tool result dicts formatted for the provider. """ from backend.apps.agents.models import Message results = [] for block in response.content: if block.type != "tool_use" or not block.tool_call: continue tc = block.tool_call start_time = time.time() # HITL permission check approved, updated_input = await self.hitl_handler(tc.name, tc.input) if not approved: result_content = [{"type": "text", "text": "Tool use was denied by the user."}] else: tool_input = updated_input if updated_input else tc.input try: result_content = await self.tool_executor(tc.name, tool_input) except Exception as e: logger.warning(f"Tool execution error: {tc.name}: {e}") result_content = [{"type": "text", "text": f"Error executing {tc.name}: {e}"}] elapsed_ms = int((time.time() - start_time) * 1000) # Emit tool result to frontend result_text = "" for block_item in result_content: if isinstance(block_item, dict) and block_item.get("type") == "text": result_text = block_item.get("text", "") break result_msg = Message( role="tool_result", content={ "text": result_text[:15000] if result_text else "Done.", "tool_name": tc.name, "elapsed_ms": elapsed_ms, }, ) await self.ws_emitter("agent:message", { "message": result_msg.model_dump(mode="json"), }) # Format for provider results.append( self.provider.format_tool_result(tc.id, result_content) ) return results