diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 0195e04a..48910c26 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -29,6 +29,8 @@ from backend.apps.agents.manager.MockAgent import MockAgent from backend.apps.agents.manager.RunSupport import RunSupport from backend.apps.agents.manager.run.handle_run_error import handle_run_error from backend.apps.agents.manager.run.TurnRunner import TurnRunner +from backend.apps.agents.manager.run.client_pool import ClientHandle +from backend.apps.agents.manager.streaming.HookContext import HookContext from backend.apps.agents.manager.run.RunOptions import RunOptions logger = logging.getLogger(__name__) @@ -49,9 +51,9 @@ class AgentManager(SessionLifecycle, SessionPersistence, Messaging, SessionContr # Per-session cancel signal: the loop stashes its asyncio.Event here so a stop/close can set it. Lives on the manager, not the AgentSession model, so it stays out of serialization (an Event can't be model_dump'd). self.cancel_events: Dict[str, asyncio.Event] = {} # Persistent-client pool (lever A, flag-gated): one live CLI per session, reused across turns. - self.client_pool: Dict[str, object] = {} + self.client_pool: Dict[str, ClientHandle] = {} # Per-SESSION hook context + stderr buffer, updated in place each turn: a persistent client's hooks/stderr callback were bound at connect, so they must read stable objects, not per-turn rebuilds. - self.hook_ctxs: Dict[str, object] = {} + self.hook_ctxs: Dict[str, HookContext] = {} self.stderr_buffers: Dict[str, List[str]] = {} # Admission gate: one shared semaphore caps concurrent ROOT turns (children bypass). (Re)created per running loop by get_turn_admission so it never binds to a dead loop across a uvicorn reload or a test's asyncio.run. self.p_turn_admission_sema: Optional[asyncio.Semaphore] = None diff --git a/backend/apps/agents/manager/AgentManagerProtocol.py b/backend/apps/agents/manager/AgentManagerProtocol.py index c2498883..da6987b5 100644 --- a/backend/apps/agents/manager/AgentManagerProtocol.py +++ b/backend/apps/agents/manager/AgentManagerProtocol.py @@ -13,10 +13,12 @@ just sits once in the MRO. Re-enables the linter's pyright reportAttributeAccess from __future__ import annotations import asyncio -from typing import TYPE_CHECKING, Any, Dict +from typing import TYPE_CHECKING, Any, Dict, List if TYPE_CHECKING: from backend.apps.agents.core.models import AgentSession + from backend.apps.agents.manager.run.client_pool import ClientHandle + from backend.apps.agents.manager.streaming.HookContext import HookContext from backend.apps.agents.manager.streaming.PartialReply import PartialReply @@ -26,6 +28,9 @@ class AgentManagerProtocol: tasks: Dict[str, asyncio.Task] live_partial: Dict[str, PartialReply] cancel_events: Dict[str, asyncio.Event] + client_pool: Dict[str, ClientHandle] + hook_ctxs: Dict[str, HookContext] + stderr_buffers: Dict[str, List[str]] if TYPE_CHECKING: # Methods implemented on sibling mixins / AgentManager itself and called cross-mixin. Loose signatures on purpose: typeCheckingMode is off, so this only has to assert the names exist, not pin their call shapes. diff --git a/backend/apps/agents/manager/run/TurnRunner.py b/backend/apps/agents/manager/run/TurnRunner.py index 1167f547..b749a32a 100644 --- a/backend/apps/agents/manager/run/TurnRunner.py +++ b/backend/apps/agents/manager/run/TurnRunner.py @@ -6,7 +6,7 @@ except-handlers can still read them after a mid-stream failure.""" import asyncio import logging import time -from typing import Dict, List, Union +from typing import Dict, List, Union, cast from typeguard import typechecked from backend.apps.agents.core.models import AgentSession @@ -17,6 +17,7 @@ from backend.apps.agents.manager.streaming.handle_stream_event import handle_str from backend.apps.agents.manager.streaming.handle_assistant_message import handle_assistant_message from backend.apps.agents.manager.streaming.handle_result_message import handle_result_message from backend.apps.agents.manager.run.client_pool import ( + SdkClientLike, acquire_client, boot_fingerprint, dispose_client, @@ -132,8 +133,9 @@ class TurnRunner(AgentManagerProtocol): async with handle.lock: handle.turns_served += 1 try: - await handle.client.query(prompt_stream()) - await p_run_streaming_turn(p_stream=handle.client.receive_response()) + sdk = cast(SdkClientLike, handle.client) + await sdk.query(prompt_stream()) + await p_run_streaming_turn(p_stream=sdk.receive_response()) # LRU by turn-END so a session mid-long-turn isn't first cap-evicted the instant it finishes. handle.last_used = time.monotonic() except BaseException: diff --git a/backend/apps/agents/manager/run/client_pool.py b/backend/apps/agents/manager/run/client_pool.py index 368a94d5..722cb30a 100644 --- a/backend/apps/agents/manager/run/client_pool.py +++ b/backend/apps/agents/manager/run/client_pool.py @@ -14,7 +14,7 @@ import json import logging import os import time -from typing import Awaitable, Callable, Dict, List, Optional +from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Optional, Protocol, cast from pydantic import BaseModel, ConfigDict, InstanceOf from typeguard import typechecked @@ -58,6 +58,15 @@ def boot_fingerprint(options_kwargs: Dict, session: AgentSession) -> str: return hashlib.sha256(blob.encode()).hexdigest() +class SdkClientLike(Protocol): + """The slice of claude_agent_sdk.ClaudeSDKClient the pool touches. The real class can't be + module-imported here (mock-mode must import the manager without the SDK), so callers cast.""" + + async def query(self, prompt: Any) -> None: ... + def receive_response(self) -> AsyncIterator[Any]: ... + async def disconnect(self) -> None: ... + + class ClientHandle(BaseModel): model_config = ConfigDict(validate_assignment=True) @@ -150,7 +159,7 @@ async def dispose_client(pool: Dict[str, ClientHandle], session_id: str) -> None if handle is None: return try: - await handle.client.disconnect() + await cast(SdkClientLike, handle.client).disconnect() except Exception: logger.exception(f"[client-pool] {session_id}: disconnect failed (subprocess may already be dead)") @@ -164,7 +173,7 @@ def dispose_client_soon(pool: Dict[str, ClientHandle], session_id: str) -> None: async def p_bg() -> None: try: - await handle.client.disconnect() + await cast(SdkClientLike, handle.client).disconnect() except Exception: logger.exception(f"[client-pool] {session_id}: background disconnect failed")