diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 8d2124fb..7266b6a4 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -99,6 +99,54 @@ class AgentManager(SessionLifecycle, SessionPersistence, Messaging, SessionContr pass yield + @typechecked + async def prewarm_client(self, session_id: str) -> None: + """Spawn the session's CLI in the seconds between create and the first message, so the first + turn's acquire is a pool hit instead of a 0.6-1.6s cold connect. Best-effort: any failure + just means the first turn pays the connect it always paid. Kill switch OSW_PREWARM_CLI=0.""" + if os.environ.get("OSW_PREWARM_CLI", "1") == "0": + return + session = self.sessions.get(session_id) + if not session or session.messages: + return + try: + import claude_agent_sdk # noqa: F401 + except ImportError: + return + try: + from backend.apps.agents.providers.registry import ( + resolve_model_id_for_sdk as p_resolve, + get_api_type as p_api_of, + ) + p_router_model_id = p_resolve(session.model, load_settings()) + p_api_type = p_api_of(session.model) + builtin_perms = load_builtin_permissions() + # Representative-LENGTH prompt: thinking derives from prompt length (<50 chars forces it + # off), so an empty prewarm prompt would boot a different thinking config than a typical + # first message and fingerprint-miss into a respawn. 50+ chars matches the common case. + p_representative = "prewarm placeholder prompt of representative length for boot" + (options, options_kwargs, _pc, _stderr, _gs) = await self.build_agent_options( + session, session_id, p_representative, "", builtin_perms, + None, None, None, False, p_router_model_id, p_api_type) + from claude_agent_sdk import ClaudeSDKClient + from backend.apps.agents.manager.run.client_pool import acquire_client, boot_fingerprint + + async def p_connect(): + p_client = ClaudeSDKClient(options=options) + logger.info(f"[SPAWN-PHASE] prewarm-connect start session={session_id[:8]} t={time.monotonic():.3f}") + await p_client.connect() + logger.info(f"[SPAWN-PHASE] prewarm-connect done session={session_id[:8]} t={time.monotonic():.3f}") + return p_client + + fp = boot_fingerprint(options_kwargs, session) + await acquire_client(self.client_pool, session_id, fp, p_connect) + # Deleted mid-connect: the late-arriving client just pooled into a dead session; nothing else will ever dispose it. + if session_id not in self.sessions: + from backend.apps.agents.manager.run.client_pool import dispose_client + await dispose_client(self.client_pool, session_id) + except Exception: + logger.info("[client-pool] prewarm skipped for %s", session_id[:8], exc_info=True) + @typechecked async def 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, context_valve_retry: bool = False): """Run the Claude Agent SDK query loop for a session.""" diff --git a/backend/apps/agents/agents.py b/backend/apps/agents/agents.py index d25a8da3..ed653edb 100644 --- a/backend/apps/agents/agents.py +++ b/backend/apps/agents/agents.py @@ -135,6 +135,9 @@ async def launch_agent(config: AgentConfig): # A launch that carries a prompt runs it as the first turn through the same path /message uses. if config.prompt: asyncio.create_task(agent_manager.send_message(session.id, config.prompt)) + else: + # No prompt yet: the user is typing. Spend that window spawning the CLI so the first turn is a pool hit. + asyncio.create_task(agent_manager.prewarm_client(session.id)) return {"session_id": session.id, "session": session.model_dump(mode="json")} @agents.router.post("/sessions/{session_id}/message") diff --git a/backend/apps/agents/manager/run/client_pool.py b/backend/apps/agents/manager/run/client_pool.py index 3cebc2cc..910fd6ab 100644 --- a/backend/apps/agents/manager/run/client_pool.py +++ b/backend/apps/agents/manager/run/client_pool.py @@ -125,6 +125,11 @@ async def trim_pool_to_cap(pool: Dict[str, "ClientHandle"]) -> None: await dispose_client(pool, sid) +# One connect per session at a time: a pre-warm and a racing first turn must SHARE a spawn, or the +# second spawn silently leaks the first (two CLI processes, one pooled). +p_inflight_connects: Dict[str, "asyncio.Task[ClientHandle]"] = {} + + @typechecked async def acquire_client( pool: Dict[str, ClientHandle], @@ -145,15 +150,32 @@ async def acquire_client( reason = "force_respawn" if force_respawn else "fingerprint_changed" logger.info(f"[client-pool] {session_id}: respawn ({reason})") await dispose_client(pool, session_id) - client = await connect_fn() - now = time.monotonic() - handle = ClientHandle( - fingerprint=fingerprint, client=client, lock=asyncio.Lock(), connected_at=now, last_used=now, - ) - pool[session_id] = handle - logger.info(f"[client-pool] {session_id}: connected fresh client") - await trim_pool_to_cap(pool) - return handle + + inflight = p_inflight_connects.get(session_id) + if inflight is not None and not inflight.done(): + # Shielded so a cancelled waiter (user stops the turn) never kills the shared spawn. + handle = await asyncio.shield(inflight) + if not force_respawn and handle.fingerprint == fingerprint: + handle.last_used = time.monotonic() + return handle + await dispose_client(pool, session_id) + + async def p_connect_and_pool() -> ClientHandle: + client = await connect_fn() + now = time.monotonic() + handle = ClientHandle( + fingerprint=fingerprint, client=client, lock=asyncio.Lock(), connected_at=now, last_used=now, + ) + pool[session_id] = handle + logger.info(f"[client-pool] {session_id}: connected fresh client") + await trim_pool_to_cap(pool) + return handle + + task = asyncio.ensure_future(p_connect_and_pool()) + p_inflight_connects[session_id] = task + # Popped when the SPAWN finishes, not when this caller returns: a cancelled owner must not strand the entry. + task.add_done_callback(lambda t: p_inflight_connects.pop(session_id, None) if p_inflight_connects.get(session_id) is t else None) + return await asyncio.shield(task) @typechecked diff --git a/backend/tests/test_client_pool.py b/backend/tests/test_client_pool.py index ea3db72c..78524a9d 100644 --- a/backend/tests/test_client_pool.py +++ b/backend/tests/test_client_pool.py @@ -244,3 +244,61 @@ def test_seeded_simulation_invariants(): assert not pool["sim"].client.disconnected asyncio.run(run()) + + +def test_concurrent_acquires_share_one_spawn() -> None: + """A pre-warm and a racing first turn must never double-spawn: the second spawn used to leak the first CLI.""" + import asyncio + + from backend.apps.agents.manager.run.client_pool import acquire_client + + async def run() -> None: + pool: dict = {} + spawns = 0 + + class FakeClient: + async def disconnect(self) -> None: + return None + + async def connect_fn(): + nonlocal spawns + spawns += 1 + await asyncio.sleep(0.05) + return FakeClient() + + a, b = await asyncio.gather( + acquire_client(pool, "sess-race", "fp1", connect_fn), + acquire_client(pool, "sess-race", "fp1", connect_fn), + ) + assert spawns == 1, f"double spawn: {spawns}" + assert a is b + assert pool["sess-race"].client is a.client + + asyncio.run(run()) + + +def test_cancelled_waiter_does_not_kill_the_shared_spawn() -> None: + import asyncio + + from backend.apps.agents.manager.run.client_pool import acquire_client + + async def run() -> None: + pool: dict = {} + + class FakeClient: + async def disconnect(self) -> None: + return None + + async def connect_fn(): + await asyncio.sleep(0.08) + return FakeClient() + + first = asyncio.ensure_future(acquire_client(pool, "sess-cancel", "fp1", connect_fn)) + await asyncio.sleep(0.01) + second = asyncio.ensure_future(acquire_client(pool, "sess-cancel", "fp1", connect_fn)) + await asyncio.sleep(0.01) + second.cancel() + handle = await first + assert pool["sess-cancel"].client is handle.client + + asyncio.run(run())