From a1cbaaeabb4fc193aec4d3613baec646c47e3051 Mon Sep 17 00:00:00 2001 From: ciregenz Date: Tue, 23 Jun 2026 05:15:34 -0700 Subject: [PATCH] [eric] agents: @typechecked + typed signature on the loop, p_ the live-partial mirror (convention) --- backend/apps/agents/agent_manager.py | 25 +++++++++++-------- .../apps/agents/manager/AgentLaunchMixin.py | 4 +-- backend/apps/agents/manager/MessagingMixin.py | 6 ++--- .../apps/agents/manager/RunSupportMixin.py | 4 +-- .../manager/session/SessionLifecycleMixin.py | 2 +- backend/tests/test_session_cleanup.py | 4 +-- backend/tests/test_streaming_harness.py | 14 +++++------ 7 files changed, 31 insertions(+), 28 deletions(-) diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 32eb7d71..c3be525f 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -7,7 +7,8 @@ import sys import time from datetime import datetime from uuid import uuid4 -from typing import Dict, Optional +from typing import Dict, List, Optional +from typeguard import typechecked from backend.apps.agents.core.models import ( AgentConfig, AgentSession, Message, MessageBranch, ApprovalRequest, ToolGroupMeta, @@ -90,13 +91,14 @@ os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000") class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunSupportMixin): + @typechecked def __init__(self): - self.sessions: dict[str, AgentSession] = {} - self.tasks: dict[str, asyncio.Task] = {} + self.sessions: Dict[str, AgentSession] = {} + self.tasks: Dict[str, asyncio.Task] = {} # Live mirror of the in-flight streamed assistant text per session, so a # stop can persist the partial reply instantly instead of waiting out the # multi-second SDK teardown the cancel handler sits behind. - self._live_partial: Dict[str, LivePartial] = {} + self.p_live_partial: Dict[str, LivePartial] = {} @@ -120,7 +122,8 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunS - async def _run_agent_loop(self, session_id: str, prompt: str, images: list | None = None, context_paths: list | None = None, forced_tools: list[str] | None = None, attached_skills: list | None = None, fork_session: bool = False, selected_browser_ids: list[str] | None = None, selected_app_output_ids: list[str] | None = None, selected_setting_ids: list[str] | None = None): + @typechecked + async def p_run_agent_loop(self, session_id: str, prompt: str, images: Optional[List] = None, context_paths: Optional[List] = None, forced_tools: Optional[List[str]] = None, attached_skills: Optional[List] = None, fork_session: bool = False, selected_browser_ids: Optional[List[str]] = None, selected_app_output_ids: Optional[List[str]] = None, selected_setting_ids: Optional[List[str]] = None): """Run the Claude Agent SDK query loop for a session.""" session = self.sessions.get(session_id) if not session: @@ -512,7 +515,7 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunS logger.info(f"[MCP-DEBUG] effective_disallowed: {effective_disallowed}") # `p_router_model_id` and `p_api_type_for_session` were resolved - # at the top of _run_agent_loop (before any closures were + # at the top of p_run_agent_loop (before any closures were # defined) so analytics closures could tag events with them. # Reuse those values here and keep session.provider in sync. resolved_model = p_router_model_id @@ -1158,12 +1161,12 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunS if isinstance(message, StreamEvent): await stream_event.handle_stream_event( - message, session, session_id, turn, thinking, self._live_partial + message, session, session_id, turn, thinking, self.p_live_partial ) elif isinstance(message, AssistantMessage): await assistant_message.handle_assistant_message( - message, session, session_id, turn, thinking, self._live_partial, self.sessions + message, session, session_id, turn, thinking, self.p_live_partial, self.sessions ) elif isinstance(message, ResultMessage): await result_message.handle_result_message( @@ -1214,7 +1217,7 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunS }) turn.stream_text_msg_id = None turn.stream_text_accum = "" - self._live_partial.pop(session_id, None) + self.p_live_partial.pop(session_id, None) for p_tool_msg_id in turn.stream_tool_msg_ids_ordered: await ws_manager.send_to_session(session_id, "agent:stream_end", { "session_id": session_id, @@ -1237,7 +1240,7 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunS # analogous flow) flagged pending_continuation during this # turn, kick off a follow-up turn immediately with the # captured prompt. We dispatch as a fire-and-forget task so - # the current _run_agent_loop frame can unwind cleanly + # the current p_run_agent_loop frame can unwind cleanly # before the next turn's options + history rebuild kicks in. # The follow-up is `hidden=True` so it doesn't add a user # bubble to the visible chat; the model sees it as a @@ -1523,7 +1526,7 @@ class AgentManager(SessionLifecycleMixin, MessagingMixin, AgentLaunchMixin, RunS # snapshot the live turn is writing. p_is_live_task = self.tasks.get(session_id) is asyncio.current_task() if p_is_live_task: - self._live_partial.pop(session_id, None) + self.p_live_partial.pop(session_id, None) if session_id in self.sessions and p_is_live_task: # For canvas-launched App Builder sessions, the workspace # folder IS the session_id (see launch_agent), so meta.json diff --git a/backend/apps/agents/manager/AgentLaunchMixin.py b/backend/apps/agents/manager/AgentLaunchMixin.py index 8fa7deae..c602b2af 100644 --- a/backend/apps/agents/manager/AgentLaunchMixin.py +++ b/backend/apps/agents/manager/AgentLaunchMixin.py @@ -1,6 +1,6 @@ """Agent run entry points for AgentManager: launch a new top-level run, run the no-SDK mock fallback, and the staticmethod invoke_agent helper. Split into a mixin to keep the manager file -under the size ceiling; self._run_agent_loop / self.p_stream_text / self.sessions resolve across +under the size ceiling; self.p_run_agent_loop / self.p_stream_text / self.sessions resolve across the MRO exactly as before.""" import asyncio @@ -327,7 +327,7 @@ class AgentLaunchMixin: "message": user_msg.model_dump(mode="json"), }) - await self._run_agent_loop(fork.id, message, fork_session=True) + await self.p_run_agent_loop(fork.id, message, fork_session=True) last_assistant = None for msg in reversed(fork.messages): diff --git a/backend/apps/agents/manager/MessagingMixin.py b/backend/apps/agents/manager/MessagingMixin.py index c63544aa..69c4b4de 100644 --- a/backend/apps/agents/manager/MessagingMixin.py +++ b/backend/apps/agents/manager/MessagingMixin.py @@ -1,6 +1,6 @@ """User-facing message operations for AgentManager (send / stop / edit / branch / approve / update), split into a mixin to keep the manager file under the size ceiling. Pure relocation: -self._run_agent_loop / self.sessions / self.stop_agent all resolve across the MRO as before.""" +self.p_run_agent_loop / self.sessions / self.stop_agent all resolve across the MRO as before.""" import asyncio import logging @@ -169,7 +169,7 @@ class MessagingMixin: if fast_verdict != "no": task = asyncio.create_task(browser_dispatch.run_browser_fast_path(session, session_id, prompt, selected_browser_ids, fast_brief, fast_verdict)) else: - task = asyncio.create_task(self._run_agent_loop(session_id, prompt, images=images, context_paths=context_paths, forced_tools=forced_tools, attached_skills=attached_skills, selected_browser_ids=selected_browser_ids, selected_app_output_ids=selected_app_output_ids, selected_setting_ids=selected_setting_ids)) + task = asyncio.create_task(self.p_run_agent_loop(session_id, prompt, images=images, context_paths=context_paths, forced_tools=forced_tools, attached_skills=attached_skills, selected_browser_ids=selected_browser_ids, selected_app_output_ids=selected_app_output_ids, selected_setting_ids=selected_setting_ids)) self.tasks[session_id] = task @typechecked @@ -308,7 +308,7 @@ class MessagingMixin: "session": session.model_dump(mode="json"), }) - task = asyncio.create_task(self._run_agent_loop( + task = asyncio.create_task(self.p_run_agent_loop( session_id, new_content, images=target_msg.images, context_paths=target_msg.context_paths, diff --git a/backend/apps/agents/manager/RunSupportMixin.py b/backend/apps/agents/manager/RunSupportMixin.py index 4c11628f..9e533208 100644 --- a/backend/apps/agents/manager/RunSupportMixin.py +++ b/backend/apps/agents/manager/RunSupportMixin.py @@ -1,7 +1,7 @@ """Per-run support methods for AgentManager: build the gated MCP server set, warm the prompt cache, stream-emit helpers, commit/drain a stopped turn, context-update broadcast, and the aux metadata + prompt/attachment delegators. Split into a mixin to keep the manager file under the -size ceiling; self.sessions / self.tasks / self._live_partial resolve across the MRO as before.""" +size ceiling; self.sessions / self.tasks / self.p_live_partial resolve across the MRO as before.""" import asyncio import logging @@ -200,7 +200,7 @@ class RunSupportMixin: push it to the client, idempotently. Lets a stop show the partial instantly instead of waiting out the SDK teardown the cancel handler sits behind. Returns True if it committed something.""" - live = self._live_partial.pop(session.id, None) + live = self.p_live_partial.pop(session.id, None) if not live: return False text = live.text or "" diff --git a/backend/apps/agents/manager/session/SessionLifecycleMixin.py b/backend/apps/agents/manager/session/SessionLifecycleMixin.py index 490d0334..5e3fa279 100644 --- a/backend/apps/agents/manager/session/SessionLifecycleMixin.py +++ b/backend/apps/agents/manager/session/SessionLifecycleMixin.py @@ -103,7 +103,7 @@ class SessionLifecycleMixin: wires its eviction in HERE and both removal paths get it for free.""" self.sessions.pop(session_id, None) self.tasks.pop(session_id, None) - self._live_partial.pop(session_id, None) + self.p_live_partial.pop(session_id, None) view_builder_render_retry_counts.pop(session_id, None) view_builder_dirty_sessions.discard(session_id) diff --git a/backend/tests/test_session_cleanup.py b/backend/tests/test_session_cleanup.py index 1317be5c..2f8e0e32 100644 --- a/backend/tests/test_session_cleanup.py +++ b/backend/tests/test_session_cleanup.py @@ -18,7 +18,7 @@ def test_purge_session_memory_clears_every_structure(): mgr = am.AgentManager() mgr.sessions = {"dead": object(), "alive": object()} mgr.tasks = {"dead": object()} - mgr._live_partial = {"dead": {"text": "half a reply"}} + mgr.p_live_partial = {"dead": {"text": "half a reply"}} vbs.view_builder_render_retry_counts["dead"] = 4 vbs.view_builder_dirty_sessions.add("dead") @@ -26,7 +26,7 @@ def test_purge_session_memory_clears_every_structure(): assert "dead" not in mgr.sessions assert "dead" not in mgr.tasks - assert "dead" not in mgr._live_partial + assert "dead" not in mgr.p_live_partial assert "dead" not in vbs.view_builder_render_retry_counts assert "dead" not in vbs.view_builder_dirty_sessions # Only the target id is purged; an unrelated live session survives. diff --git a/backend/tests/test_streaming_harness.py b/backend/tests/test_streaming_harness.py index a1f304d7..81bf2a1a 100644 --- a/backend/tests/test_streaming_harness.py +++ b/backend/tests/test_streaming_harness.py @@ -1,4 +1,4 @@ -"""Streaming harness: drive the real _run_agent_loop with a MOCKED claude_agent_sdk.query +"""Streaming harness: drive the real p_run_agent_loop with a MOCKED claude_agent_sdk.query that yields a controlled SDK message sequence, and assert the session state + emitted WS events. This is the safety net for restructuring the streaming loop (it had no isolated coverage), so it pins the observable contract: streamed text lands as an assistant message, @@ -22,7 +22,7 @@ def _mock_query_yielding(*messages): def _drive(monkeypatch, messages, prompt="hi"): - """Run one _run_agent_loop turn against a mocked SDK message stream; return (session, ws_events).""" + """Run one p_run_agent_loop turn against a mocked SDK message stream; return (session, ws_events).""" events = [] async def fake_send(session_id, event, data): @@ -35,7 +35,7 @@ def _drive(monkeypatch, messages, prompt="hi"): from backend.apps.agents.core.models import AgentSession session = AgentSession(name="t", model="sonnet", dashboard_id="d") mgr.sessions[session.id] = session - asyncio.run(mgr._run_agent_loop(session.id, prompt)) + asyncio.run(mgr.p_run_agent_loop(session.id, prompt)) return session, events @@ -55,7 +55,7 @@ def _assistant(blocks, **kw): def test_loop_wires_all_four_hooks_to_a_live_hook_context(monkeypatch): # Integration coverage the unit tests can't give: capture the ClaudeAgentOptions the real - # loop hands to query(), then invoke the WIRED hooks. This proves _run_agent_loop builds a + # loop hands to query(), then invoke the WIRED hooks. This proves p_run_agent_loop builds a # HookContext (all required fields, incl. the live `sessions` registry) and the four thin # wrappers delegate to the extracted hook modules. The SDK never fires these under a mocked # query, so without this the wiring (not just the functions) would be untested. @@ -76,7 +76,7 @@ def test_loop_wires_all_four_hooks_to_a_live_hook_context(monkeypatch): from backend.apps.agents.core.models import AgentSession session = AgentSession(name="t", model="sonnet", dashboard_id="d") mgr.sessions[session.id] = session - asyncio.run(mgr._run_agent_loop(session.id, "hi")) + asyncio.run(mgr.p_run_agent_loop(session.id, "hi")) options = captured["options"] assert options is not None @@ -174,7 +174,7 @@ def test_transient_capacity_error_is_retried_then_succeeds(monkeypatch): from backend.apps.agents.core.models import AgentSession session = AgentSession(name="t", model="sonnet", dashboard_id="d") mgr.sessions[session.id] = session - asyncio.run(mgr._run_agent_loop(session.id, "hi")) + asyncio.run(mgr.p_run_agent_loop(session.id, "hi")) assert state["n"] == 2 # retried exactly once assert any(m.role == "assistant" and "Recovered" in str(m.content) for m in session.messages) @@ -207,7 +207,7 @@ def test_thinking_pill_shows_per_turn_delta_not_cumulative(monkeypatch): session = AgentSession(name="t", model="sonnet", dashboard_id="d") session.tokens = {"input_fresh": 1000, "output": 500} # prior-turn accumulation mgr.sessions[session.id] = session - asyncio.run(mgr._run_agent_loop(session.id, "hi")) + asyncio.run(mgr.p_run_agent_loop(session.id, "hi")) assert pills, "expected a consolidated thinking pill" assert pills[-1]["input_tokens"] == 150 # (1100-1000)+(550-500), not the cumulative 1650