step 3.3 complete: device WS endpoint + DeviceConnectionManager

app/core/device_manager.py

@@ -0,0 +1,183 @@
+"""Device connection manager.
+
+Maintains in-memory state for all active Electron → backend WebSocket
+connections.  One connection per user (latest replaces previous).
+
+The manager participates in two interaction patterns:
+
+1. **Tool-call round-trip** (bidirectional CRUD):
+   - Backend sends ``tool_call`` frame → Electron executes CRUD → returns
+     ``tool_result`` frame.
+   - ``create_pending_call`` registers a Future keyed by ``call_id``.
+   - ``resolve_pending_call`` fulfils the Future; callers awaiting it
+     receive the result dict from Electron.
+
+2. **Agent-data streaming** (local directory agent runs):
+   - Backend sends ``agent_run`` frame → Electron reads files and sends
+     back a stream of ``agent_data`` frames followed by ``agent_complete``.
+   - ``get_agent_data_queue`` returns (or creates) an asyncio.Queue for
+     a specific ``run_id`` so the agent runner can iterate frames.
+
+The ``device_manager`` module-level singleton is imported by both the
+device WS route and the agent runner.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from dataclasses import dataclass, field
+
+from fastapi import WebSocket
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DeviceConnection:
+    """State for a single connected Electron device."""
+
+    ws: WebSocket
+    device_id: str
+    # Futures indexed by tool_call id — resolved when tool_result arrives.
+    pending_calls: dict[str, asyncio.Future[dict]] = field(default_factory=dict)
+    # Per-run queues for agent_data / agent_complete frames.
+    agent_data_queues: dict[str, asyncio.Queue[dict | None]] = field(default_factory=dict)
+
+
+class DeviceConnectionManager:
+    """Singleton registry of active Electron WebSocket connections.
+
+    Thread/task safety note: asyncio is single-threaded by design.  All
+    mutations happen inside await-points on the main event loop, so no
+    locking is required for the in-memory dicts.
+    """
+
+    def __init__(self) -> None:
+        self._connections: dict[str, DeviceConnection] = {}
+
+    # ── Registration ──────────────────────────────────────────────────
+
+    def register(self, user_id: str, device_id: str, ws: WebSocket) -> None:
+        """Store the active connection for *user_id*, replacing any previous one."""
+        if user_id in self._connections:
+            old = self._connections[user_id]
+            logger.info(
+                "device_manager: replacing existing connection for user=%s device=%s",
+                user_id,
+                old.device_id,
+            )
+            # Cancel any futures that were waiting on the old connection.
+            for fut in old.pending_calls.values():
+                if not fut.done():
+                    fut.cancel()
+        self._connections[user_id] = DeviceConnection(ws=ws, device_id=device_id)
+        logger.info(
+            "device_manager: registered user=%s device=%s", user_id, device_id
+        )
+
+    def unregister(self, user_id: str) -> None:
+        """Remove the connection for *user_id* and cancel any pending futures."""
+        conn = self._connections.pop(user_id, None)
+        if conn is None:
+            return
+        for fut in conn.pending_calls.values():
+            if not fut.done():
+                fut.cancel()
+        logger.info("device_manager: unregistered user=%s", user_id)
+
+    # ── Presence queries ──────────────────────────────────────────────
+
+    def get_ws(self, user_id: str) -> WebSocket | None:
+        """Return the active WebSocket for *user_id*, or ``None`` if offline."""
+        conn = self._connections.get(user_id)
+        return conn.ws if conn else None
+
+    def is_online(self, user_id: str, device_id: str | None = None) -> bool:
+        """Return ``True`` if the user has an active connection.
+
+        If *device_id* is provided also checks that it matches the connected device.
+        """
+        conn = self._connections.get(user_id)
+        if conn is None:
+            return False
+        if device_id is not None:
+            return conn.device_id == device_id
+        return True
+
+    # ── Frame sending ─────────────────────────────────────────────────
+
+    async def send_frame(self, user_id: str, frame: dict) -> None:
+        """Send *frame* as a JSON text message to the device.
+
+        Raises ``RuntimeError`` if the user is not connected.
+        """
+        conn = self._connections.get(user_id)
+        if conn is None:
+            raise RuntimeError(
+                f"send_frame: user {user_id!r} is not connected"
+            )
+        await conn.ws.send_text(json.dumps(frame))
+
+    # ── Tool-call round-trip ──────────────────────────────────────────
+
+    def create_pending_call(
+        self, user_id: str, call_id: str
+    ) -> asyncio.Future[dict]:
+        """Register a Future that will be resolved when the tool_result arrives.
+
+        Raises ``RuntimeError`` if the user is not connected.
+        """
+        conn = self._connections.get(user_id)
+        if conn is None:
+            raise RuntimeError(
+                f"create_pending_call: user {user_id!r} is not connected"
+            )
+        loop = asyncio.get_event_loop()
+        fut: asyncio.Future[dict] = loop.create_future()
+        conn.pending_calls[call_id] = fut
+        return fut
+
+    def resolve_pending_call(
+        self, user_id: str, call_id: str, result: dict
+    ) -> None:
+        """Fulfil the Future registered under *call_id* with the Electron result.
+
+        No-ops if the call_id is unknown (already timed out or cancelled).
+        """
+        conn = self._connections.get(user_id)
+        if conn is None:
+            return
+        fut = conn.pending_calls.pop(call_id, None)
+        if fut is not None and not fut.done():
+            fut.set_result(result)
+
+    # ── Agent-data queue ──────────────────────────────────────────────
+
+    def get_agent_data_queue(
+        self, user_id: str, run_id: str
+    ) -> asyncio.Queue[dict | None]:
+        """Return (creating if absent) the queue for *run_id* agent frames.
+
+        The agent runner reads from this queue.  The device WS handler writes
+        to it.  ``None`` is the sentinel that signals the stream is finished.
+        """
+        conn = self._connections.get(user_id)
+        if conn is None:
+            raise RuntimeError(
+                f"get_agent_data_queue: user {user_id!r} is not connected"
+            )
+        if run_id not in conn.agent_data_queues:
+            conn.agent_data_queues[run_id] = asyncio.Queue()
+        return conn.agent_data_queues[run_id]
+
+    def cleanup_agent_data_queue(self, user_id: str, run_id: str) -> None:
+        """Remove the queue for *run_id* once a run has completed."""
+        conn = self._connections.get(user_id)
+        if conn:
+            conn.agent_data_queues.pop(run_id, None)
+
+
+# Module-level singleton — import this everywhere.
+device_manager = DeviceConnectionManager()