diff --git a/backend/apps/agents/agents.py b/backend/apps/agents/agents.py index 87c9cc24..bbd9c241 100644 --- a/backend/apps/agents/agents.py +++ b/backend/apps/agents/agents.py @@ -3,8 +3,8 @@ Endpoints operate directly on a module-level sessions dict and the Agent class. No manager layer — Agent already encapsulates its own runtime state. -ws_manager is used ONLY in this file — the Agent class and its internals -communicate via the on_event callback, never importing ws_manager directly. +The ws module is used ONLY in this file — the Agent class and its internals +communicate via the on_event callback, never importing ws directly. """ from contextlib import asynccontextmanager @@ -33,7 +33,7 @@ from backend.apps.agents.session_store import ( reconcile_on_startup, load, ) -from backend.OLDapps.agents.manager.ws_manager import ws_manager +from backend.apps.agents import ws from claude_agent_sdk import ClaudeAgentOptions SESSIONS: dict[str, Agent] = {} @@ -41,9 +41,9 @@ SESSIONS: dict[str, Agent] = {} @typechecked def p_make_session_emitter(session_id: str) -> EventCallback: - """Create an event callback that routes typed events to ws_manager for a session.""" + """Create an event callback that routes typed events to the WS connection pool.""" async def emit(event: AnyEvent) -> None: - await ws_manager.send_to_session(session_id, event.event, event.model_dump(mode="json")) + await ws.send_to_session(session_id, event.event, event.model_dump(mode="json")) return emit def get_agent(session_id: str) -> Agent: @@ -211,7 +211,7 @@ class ApprovalBody(BaseModel): @agents.router.post("/approval") async def handle_approval(body: ApprovalBody) -> dict: - ws_manager.resolve_approval(body.request_id, { + ws.APPROVAL_BRIDGE.resolve(body.request_id, { "behavior": body.behavior, "message": body.message, "updated_input": body.updated_input, diff --git a/backend/apps/agents/ws.py b/backend/apps/agents/ws.py new file mode 100644 index 00000000..e35c315b --- /dev/null +++ b/backend/apps/agents/ws.py @@ -0,0 +1,117 @@ +"""WebSocket connection pool and async future bridges. + +Three concerns, one module: + +1. Connection pool — holds WebSocket objects, delivers JSON payloads. +2. FutureBridge — generic async request/response over WebSocket. +3. Bridge instances — approval_bridge and browser_bridge. +""" + +import asyncio +import json +import logging +from typing import Callable, Awaitable + +from fastapi import WebSocket + +from pydantic import BaseModel, Field + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# 1. Connection pool +# --------------------------------------------------------------------------- + +P_SESSION_CONNECTIONS: dict[str, list[WebSocket]] = {} +P_GLOBAL_CONNECTIONS: list[WebSocket] = [] + + +async def connect_session(session_id: str, ws: WebSocket) -> None: + await ws.accept() + P_SESSION_CONNECTIONS.setdefault(session_id, []).append(ws) + + +async def connect_global(ws: WebSocket) -> None: + await ws.accept() + P_GLOBAL_CONNECTIONS.append(ws) + + +def disconnect_session(session_id: str, ws: WebSocket) -> None: + conns = P_SESSION_CONNECTIONS.get(session_id) + if not conns: + return + conns[:] = [c for c in conns if c is not ws] + if not conns: + del P_SESSION_CONNECTIONS[session_id] + + +def disconnect_global(ws: WebSocket) -> None: + P_GLOBAL_CONNECTIONS[:] = [c for c in P_GLOBAL_CONNECTIONS if c is not ws] + + +async def send_to_session(session_id: str, event: str, data: dict) -> None: + payload = json.dumps({"event": event, "session_id": session_id, "data": data}) + for ws in P_SESSION_CONNECTIONS.get(session_id, []): + try: + await ws.send_text(payload) + except Exception: + pass + for ws in P_GLOBAL_CONNECTIONS: + try: + await ws.send_text(payload) + except Exception: + pass + + +async def broadcast_global(event: str, data: dict) -> None: + payload = json.dumps({"event": event, "data": data}) + for ws in P_GLOBAL_CONNECTIONS: + try: + await ws.send_text(payload) + except Exception: + pass + + +# --------------------------------------------------------------------------- +# 2. FutureBridge +# --------------------------------------------------------------------------- + +class FutureBridge(BaseModel): + """Async request/response bridge over WebSocket. + + Pattern: create a Future, send a question to the frontend, + block until the frontend responds (or timeout). + """ + + p_pending: dict[str, asyncio.Future] = Field(default_factory=dict) + + async def request( + self, + request_id: str, + send_fn: Callable[[], Awaitable[None]], + timeout: float, + ) -> dict: + future: asyncio.Future = asyncio.get_event_loop().create_future() + self.p_pending[request_id] = future + await send_fn() + try: + return await asyncio.wait_for(future, timeout=timeout) + except asyncio.TimeoutError: + logger.warning("FutureBridge request %s timed out after %ss", request_id, timeout) + return {"error": "Timed out"} + finally: + self.p_pending.pop(request_id, None) + + def resolve(self, request_id: str, result: dict) -> None: + future = self.p_pending.get(request_id) + if future and not future.done(): + future.set_result(result) + + +# --------------------------------------------------------------------------- +# 3. Bridge instances +# --------------------------------------------------------------------------- + +APPROVAL_BRIDGE = FutureBridge() +BROWSER_BRIDGE = FutureBridge() diff --git a/backend/core/Agent/Agent.py b/backend/core/Agent/Agent.py index cec32627..b91aec83 100644 --- a/backend/core/Agent/Agent.py +++ b/backend/core/Agent/Agent.py @@ -14,7 +14,7 @@ from backend.core.Agent.shared_structs.ApprovalRequest import ApprovalRequest from backend.core.Agent.shared_structs.MessageLog import MessageLog from backend.core.events.events import ( AgentSnapshot, AgentStatusEvent, AgentMessageEvent, - EventCallback, + EventCallback, AnyEvent, ) os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000")