[Haik]: ckpt, connection manager and feature bridge seem to be done, still gotta integrate the browser bridge to browser related tools tho

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,

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()

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")