gdmcode 0.1.0__py3-none-any.whl

src/server/bridge.py

@@ -0,0 +1,427 @@
+"""asyncio WebSocket bridge between gdm CLI and Chrome extension.
+
+The bridge is a separate process with its own event loop so it survives
+AgentLoop restarts and keeps the Chrome extension persistently connected.
+
+Security model:
+  - Binds to 127.0.0.1 only — never 0.0.0.0
+  - Bearer token required on first message (auth handshake)
+  - Origin validation for browser connections
+  - Token written to disk with chmod 600
+
+Usage::
+
+    config = BridgeConfig(port=9321, session_id="abc")
+    server = BridgeServer(config)
+    asyncio.run(server.run())
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import platform
+import re
+import secrets
+import time
+from collections import deque
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+__all__ = ["BridgeConfig", "BridgeServer", "TokenManager", "NONCE_TTL_SECS"]
+
+log = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+BRIDGE_DEFAULT_PORT: int = 9321
+BRIDGE_DEFAULT_HOST: str = "127.0.0.1"
+COMMAND_QUEUE_TTL_SECS: float = 30.0
+MAX_MESSAGE_SIZE_BYTES: int = 10 * 1024 * 1024  # 10 MB (screenshots)
+HEARTBEAT_INTERVAL_SECS: float = 20.0
+TOKEN_FILE_NAME: str = "bridge.token"
+HEALTH_CHECK_PATH: str = "/health"
+NONCE_TTL_SECS: float = 60.0
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+def _default_config_dir() -> Path:
+    if platform.system() == "Windows":
+        base = Path.home() / "AppData" / "Roaming" / "gdm"
+    else:
+        base = Path.home() / ".config" / "gdm"
+    base.mkdir(parents=True, exist_ok=True)
+    return base
+
+
+def _resolve_token_path(project_root: Path | None = None) -> Path:
+    """Resolve the canonical bridge token path.
+
+    Checks ``.context-memory/bridge.token`` relative to *project_root* (if
+    provided and the file exists), then falls back to the user-scoped config
+    directory (``~/.config/gdm/bridge.token`` on Linux/macOS).
+    """
+    if project_root is not None:
+        override = project_root / ".context-memory" / TOKEN_FILE_NAME
+        if override.exists():
+            return override
+    return _default_config_dir() / TOKEN_FILE_NAME
+
+
+@dataclass
+class BridgeConfig:
+    host: str = BRIDGE_DEFAULT_HOST
+    port: int = BRIDGE_DEFAULT_PORT
+    token_dir: Path = field(default_factory=_default_config_dir)
+    session_id: str = ""
+    db_path: Path | None = None
+
+
+# ---------------------------------------------------------------------------
+# Token management
+# ---------------------------------------------------------------------------
+
+@dataclass
+class TokenManager:
+    """Manages bearer token generation, storage, and validation."""
+    token_dir: Path
+
+    def generate(self) -> str:
+        token = secrets.token_urlsafe(32)
+        token_path = self.token_dir / TOKEN_FILE_NAME
+        token_path.write_text(token, encoding="utf-8")
+        try:
+            token_path.chmod(0o600)
+        except NotImplementedError:
+            pass  # Windows — chmod is a no-op
+        log.info("Bridge token written to %s", token_path)
+        return token
+
+    def load(self) -> str | None:
+        token_path = self.token_dir / TOKEN_FILE_NAME
+        if not token_path.exists():
+            return None
+        return token_path.read_text(encoding="utf-8").strip()
+
+    def validate(self, candidate: str) -> bool:
+        stored = self.load()
+        if stored is None:
+            return False
+        return secrets.compare_digest(candidate, stored)
+
+    def revoke(self) -> None:
+        (self.token_dir / TOKEN_FILE_NAME).unlink(missing_ok=True)
+
+
+# ---------------------------------------------------------------------------
+# Client registry
+# ---------------------------------------------------------------------------
+
+@dataclass
+class ConnectedClient:
+    ws: Any                                          # WebSocketServerProtocol
+    client_type: str = "unknown"
+    authenticated: bool = False
+    connected_at: float = field(default_factory=time.monotonic)
+    last_heartbeat: float = field(default_factory=time.monotonic)
+    tab_id: int | None = None
+    origin: str = ""
+
+
+class ClientRegistry:
+    def __init__(self) -> None:
+        self._clients: dict[int, ConnectedClient] = {}
+
+    def add(self, ws: Any) -> ConnectedClient:
+        client = ConnectedClient(ws=ws)
+        self._clients[id(ws)] = client
+        return client
+
+    def remove(self, ws: Any) -> None:
+        self._clients.pop(id(ws), None)
+
+    def get(self, ws: Any) -> ConnectedClient | None:
+        return self._clients.get(id(ws))
+
+    def get_cli_clients(self) -> list[ConnectedClient]:
+        return [c for c in self._clients.values()
+                if c.client_type == "cli" and c.authenticated]
+
+    def get_extension_clients(self) -> list[ConnectedClient]:
+        return [c for c in self._clients.values()
+                if c.client_type == "extension" and c.authenticated]
+
+    @property
+    def count(self) -> int:
+        return len(self._clients)
+
+    def summary(self) -> dict[str, int]:
+        cli = len(self.get_cli_clients())
+        ext = len(self.get_extension_clients())
+        return {"cli": cli, "extension": ext, "total": self.count}
+
+
+# ---------------------------------------------------------------------------
+# Command queue
+# ---------------------------------------------------------------------------
+
+@dataclass
+class PendingCommand:
+    id: str
+    payload: dict[str, Any]
+    queued_at: float = field(default_factory=time.monotonic)
+    source_ws: Any | None = None
+
+
+class CommandQueue:
+    """Buffers CLI→extension commands when no extension is connected."""
+
+    def __init__(self, ttl: float = COMMAND_QUEUE_TTL_SECS) -> None:
+        self._queue: deque[PendingCommand] = deque()
+        self._ttl = ttl
+
+    def enqueue(self, cmd: PendingCommand) -> None:
+        self._expire_stale()
+        self._queue.append(cmd)
+        log.debug("Command %s queued (%d pending)", cmd.id, len(self._queue))
+
+    def drain(self) -> list[PendingCommand]:
+        self._expire_stale()
+        commands = list(self._queue)
+        self._queue.clear()
+        return commands
+
+    def _expire_stale(self) -> None:
+        now = time.monotonic()
+        while self._queue and (now - self._queue[0].queued_at) > self._ttl:
+            expired = self._queue.popleft()
+            log.warning("Command %s expired after %.1fs", expired.id, self._ttl)
+
+    @property
+    def size(self) -> int:
+        return len(self._queue)
+
+
+# ---------------------------------------------------------------------------
+# Main server
+# ---------------------------------------------------------------------------
+
+_LOCALHOST_ORIGIN_RE = re.compile(
+    r"https?://(localhost|127\.0\.0\.1)(:\d+)?$"
+)
+
+
+class BridgeServer:
+    """asyncio WebSocket bridge between gdm CLI and Chrome extension.
+
+    Handles:
+    - Bearer-token authentication on first message
+    - Command forwarding: CLI → extension (with queue when disconnected)
+    - Response forwarding: extension → CLI
+    - Status broadcast: extension → all CLI clients
+    - /health HTTP endpoint (no auth required)
+    - Origin validation for browser-side connections
+    """
+
+    def __init__(self, config: BridgeConfig) -> None:
+        self._config = config
+        self._token_mgr = TokenManager(token_dir=config.token_dir)
+        self._clients = ClientRegistry()
+        self._command_queue = CommandQueue()
+        self._pending_responses: dict[str, PendingCommand] = {}
+        self._audit_log: list[dict[str, Any]] = []
+        self._nonce_cache: dict[str, float] = {}
+
+    async def run(self) -> None:
+        """Start the server. Blocks until shutdown."""
+        try:
+            import websockets
+            from websockets.server import serve
+        except ImportError as exc:
+            raise RuntimeError(
+                "websockets package required: pip install websockets"
+            ) from exc
+
+        self._token_mgr.generate()
+        log.info(
+            "Bridge server starting on %s:%d",
+            self._config.host, self._config.port,
+        )
+        async with serve(
+            self._handle_client,
+            self._config.host,
+            self._config.port,
+            max_size=MAX_MESSAGE_SIZE_BYTES,
+            process_request=self._handle_http,
+        ):
+            log.info("Bridge ready — waiting for connections")
+            await asyncio.Future()  # run forever
+
+    async def _handle_http(
+        self, path: str, headers: Any
+    ) -> tuple[int, list[tuple[str, str]], bytes] | None:
+        if path == HEALTH_CHECK_PATH:
+            body = json.dumps({
+                "status": "ok",
+                "clients": self._clients.summary(),
+                "queue_depth": self._command_queue.size,
+            }).encode()
+            return (200, [("Content-Type", "application/json")], body)
+        return None
+
+    async def _handle_client(self, ws: Any) -> None:
+        from src.server.protocol_version import (
+            is_compatible,
+            make_hello_message,
+            make_version_mismatch_error,
+        )
+
+        client = self._clients.add(ws)
+        origin = getattr(ws, "request_headers", {}).get("Origin", "")
+        if origin and not self._is_allowed_origin(origin):
+            await ws.close(4003, f"Origin {origin} not allowed")
+            self._clients.remove(ws)
+            return
+        client.origin = origin
+
+        # Send protocol hello (protocol-001) — must arrive before auth handshake
+        hello = make_hello_message(server_name="gdm-bridge")
+        await ws.send(json.dumps(hello))
+
+        # Wait for hello_ack (5 second timeout)
+        _first_msg: dict[str, Any] | None = None
+        try:
+            raw = await asyncio.wait_for(ws.recv(), timeout=5.0)
+            ack = json.loads(raw)
+            if ack.get("type") == "hello_ack":
+                client_ver = ack.get("protocol_version", "0.0")
+                compatible, reason = is_compatible(client_ver)
+                if not compatible:
+                    await ws.send(json.dumps(make_version_mismatch_error(client_ver)))
+                    await ws.close(code=4000, reason=reason)
+                    self._clients.remove(ws)
+                    return
+                elif reason:
+                    log.warning("Protocol minor version warning: %s", reason)
+            else:
+                # Not hello_ack — backward compat: treat as first regular message
+                _first_msg = ack
+        except asyncio.TimeoutError:
+            log.warning("Client did not send hello_ack within 5s — proceeding without version check")
+
+        try:
+            if _first_msg is not None:
+                await self._route_message(client, _first_msg)
+            async for raw_msg in ws:
+                try:
+                    msg = json.loads(raw_msg)
+                except json.JSONDecodeError:
+                    log.warning("Malformed JSON from %s — ignored", client.client_type)
+                    continue
+                await self._route_message(client, msg)
+        except Exception:  # noqa: BLE001 — websockets.ConnectionClosed + others
+            log.debug("Client %s disconnected", client.client_type)
+        finally:
+            self._clients.remove(ws)
+
+    async def _route_message(self, client: ConnectedClient, msg: dict[str, Any]) -> None:
+        msg_type = msg.get("type", "")
+
+        if not client.authenticated:
+            if msg_type != "auth":
+                await client.ws.close(4001, "First message must be auth")
+                return
+            if not self._token_mgr.validate(msg.get("token", "")):
+                await client.ws.close(4002, "Invalid token")
+                return
+            client.authenticated = True
+            client.client_type = msg.get("client_type", "cli")
+            await client.ws.send(json.dumps({"type": "auth_ok"}))
+            # Flush queued commands to newly connected extension
+            if client.client_type == "extension":
+                for cmd in self._command_queue.drain():
+                    await client.ws.send(json.dumps(cmd.payload))
+            return
+
+        self._audit(msg_type, msg)
+        match msg_type:
+            case "heartbeat":
+                client.last_heartbeat = time.monotonic()
+
+            case "command":
+                nonce = msg.get("nonce", "")
+                if nonce and not self._check_and_record_nonce(nonce):
+                    self._audit("replay_rejected", msg)
+                    await client.ws.send(json.dumps({
+                        "type": "error",
+                        "code": "replay_detected",
+                        "id": msg.get("id", ""),
+                        "message": "Replayed nonce rejected",
+                    }))
+                    return
+                extensions = self._clients.get_extension_clients()
+                cmd = PendingCommand(id=msg.get("id", ""), payload=msg, source_ws=client.ws)
+                if extensions:
+                    await extensions[0].ws.send(json.dumps(msg))
+                    self._pending_responses[msg.get("id", "")] = cmd
+                else:
+                    self._command_queue.enqueue(cmd)
+
+            case "response":
+                cmd_id = msg.get("id", "")
+                pending = self._pending_responses.pop(cmd_id, None)
+                if pending and pending.source_ws:
+                    await pending.source_ws.send(json.dumps(msg))
+
+            case "status":
+                for cli in self._clients.get_cli_clients():
+                    await cli.ws.send(json.dumps(msg))
+
+    def _check_and_record_nonce(self, nonce: str) -> bool:
+        """Record nonce; return True if fresh, False if replayed within NONCE_TTL_SECS."""
+        now = time.monotonic()
+        stale = [n for n, ts in self._nonce_cache.items() if (now - ts) > NONCE_TTL_SECS]
+        for n in stale:
+            del self._nonce_cache[n]
+        if nonce in self._nonce_cache:
+            return False
+        self._nonce_cache[nonce] = now
+        return True
+
+    def _is_allowed_origin(self, origin: str) -> bool:
+        if origin.startswith("chrome-extension://"):
+            return True
+        if origin in ("http://localhost", "http://127.0.0.1"):
+            return True
+        return bool(_LOCALHOST_ORIGIN_RE.match(origin))
+
+    def _audit(self, msg_type: str, msg: dict[str, Any]) -> None:
+        self._audit_log.append({
+            "timestamp": time.time(),
+            "type": msg_type,
+            "msg_id": msg.get("id", ""),
+            "summary": json.dumps(msg)[:500],
+        })
+
+    async def shutdown(self) -> None:
+        """Graceful shutdown: notify pending commands, close all clients."""
+        for cmd in self._command_queue.drain():
+            if cmd.source_ws:
+                try:
+                    await cmd.source_ws.send(json.dumps({
+                        "type": "response", "id": cmd.id,
+                        "ok": False, "error": "Bridge shutting down",
+                    }))
+                except Exception:  # noqa: BLE001
+                    pass
+        for client in list(self._clients._clients.values()):
+            try:
+                await client.ws.close(1001, "Bridge shutting down")
+            except Exception:  # noqa: BLE001
+                pass

src/server/bridge_cli.py

@@ -0,0 +1,103 @@
+"""Entry point for launching the bridge server as a subprocess.
+
+Usage::
+
+    python -m src.server.bridge_cli --port 9321
+    gdm browser start
+"""
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+from src.server.bridge import BRIDGE_DEFAULT_PORT, BridgeConfig, BridgeServer
+
+__all__ = ["main", "start_bridge", "stop_bridge", "bridge_status"]
+
+_bridge_process: subprocess.Popen | None = None
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="gdm-bridge",
+        description="Start the gdm WebSocket bridge server",
+    )
+    p.add_argument("--port", type=int, default=9321, help="Listen port (default: 9321)")
+    p.add_argument("--host", default="127.0.0.1", help="Bind address (default: 127.0.0.1)")
+    p.add_argument("--session-id", default="", help="Associate with a gdm session ID")
+    p.add_argument("--token-dir", type=Path, default=None, help="Directory to store auth token")
+    p.add_argument("--log-level", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"])
+    return p
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(
+        level=getattr(logging, args.log_level),
+        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    )
+
+    cfg = BridgeConfig(
+        host=args.host,
+        port=args.port,
+        session_id=args.session_id,
+    )
+    if args.token_dir:
+        cfg.token_dir = args.token_dir
+
+    server = BridgeServer(cfg)
+    try:
+        asyncio.run(server.run())
+    except KeyboardInterrupt:
+        pass
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()
+
+
+def start_bridge(port: int = BRIDGE_DEFAULT_PORT, token_dir: Path | None = None) -> None:
+    """Start the bridge server in a background subprocess (no-op if already running)."""
+    global _bridge_process
+    if _bridge_process is not None and _bridge_process.poll() is None:
+        return
+    cmd = [sys.executable, "-m", "src.server.bridge_cli", "--port", str(port)]
+    if token_dir is not None:
+        cmd += ["--token-dir", str(token_dir)]
+    _bridge_process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+
+
+def stop_bridge() -> None:
+    """Gracefully stop the bridge server subprocess."""
+    global _bridge_process
+    if _bridge_process is None:
+        return
+    _bridge_process.terminate()
+    try:
+        _bridge_process.wait(timeout=5.0)
+    except subprocess.TimeoutExpired:
+        _bridge_process.kill()
+    _bridge_process = None
+
+
+def bridge_status() -> dict[str, Any]:
+    """Return running status and port of the bridge server."""
+    global _bridge_process
+    running = _bridge_process is not None and _bridge_process.poll() is None
+    return {
+        "running": running,
+        "port": BRIDGE_DEFAULT_PORT,
+        "pid": _bridge_process.pid if running else None,
+    }

src/server/bridge_client.py

@@ -0,0 +1,170 @@
+"""Synchronous CLI-side WebSocket client for communicating with the bridge server.
+
+Runs an asyncio event loop in a background thread, exposing a blocking
+send_command() for use from the synchronous AgentLoop.
+
+Usage::
+
+    client = BridgeClient(url="ws://127.0.0.1:9321", token="abc123")
+    client.connect()
+    response = client.send_command(action="click", params={"selector": "#btn"})
+    client.disconnect()
+"""
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import json
+import logging
+import threading
+import uuid
+from typing import Any
+
+__all__ = ["BridgeClient"]
+
+log = logging.getLogger(__name__)
+
+
+class BridgeClient:
+    """Synchronous WebSocket client wrapping an asyncio WS in a background thread."""
+
+    def __init__(self, url: str, token: str) -> None:
+        self._url = url
+        self._token = token
+        self._ws: Any = None
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._thread: threading.Thread | None = None
+        self._connected = threading.Event()
+        self._disconnect_flag = threading.Event()
+        self._pending: dict[str, concurrent.futures.Future[dict[str, Any]]] = {}
+        self._lock = threading.Lock()
+
+    # ── Public interface ───────────────────────────────────────────────────
+
+    @property
+    def is_connected(self) -> bool:
+        return self._connected.is_set() and not self._disconnect_flag.is_set()
+
+    def connect(self, timeout: float = 10.0) -> None:
+        """Start the background event loop and connect to the bridge."""
+        self._loop = asyncio.new_event_loop()
+        self._disconnect_flag.clear()
+        self._thread = threading.Thread(
+            target=self._run_loop,
+            daemon=True,
+            name="BridgeClient-io",
+        )
+        self._thread.start()
+        if not self._connected.wait(timeout=timeout):
+            raise ConnectionError(
+                f"Failed to connect to bridge at {self._url} within {timeout}s"
+            )
+        log.debug("BridgeClient connected to %s", self._url)
+
+    def disconnect(self) -> None:
+        """Close the WebSocket and stop the background event loop."""
+        self._disconnect_flag.set()
+        self._connected.clear()
+        if self._loop and self._ws:
+            asyncio.run_coroutine_threadsafe(self._ws.close(), self._loop)
+        if self._thread and self._thread.is_alive():
+            self._thread.join(timeout=5)
+
+    def send_command(
+        self,
+        action: str,
+        params: dict[str, Any],
+        timeout: float = 10.0,
+    ) -> dict[str, Any]:
+        """Send a command and block until response is received.
+
+        Raises:
+            TimeoutError: No response within `timeout` seconds.
+            ConnectionError: Bridge is not connected.
+        """
+        if not self.is_connected:
+            raise ConnectionError("Not connected to bridge")
+
+        cmd_id = uuid.uuid4().hex[:12]
+        msg = {"type": "command", "id": cmd_id, "action": action, **params}
+
+        future: concurrent.futures.Future[dict[str, Any]] = concurrent.futures.Future()
+        with self._lock:
+            self._pending[cmd_id] = future
+
+        asyncio.run_coroutine_threadsafe(
+            self._ws.send(json.dumps(msg)), self._loop  # type: ignore[union-attr]
+        )
+
+        try:
+            return future.result(timeout=timeout)
+        except concurrent.futures.TimeoutError as exc:
+            with self._lock:
+                self._pending.pop(cmd_id, None)
+            raise TimeoutError(f"Command {cmd_id} timed out after {timeout}s") from exc
+        except Exception as exc:
+            with self._lock:
+                self._pending.pop(cmd_id, None)
+            raise ConnectionError(f"Command failed: {exc}") from exc
+
+    # ── Background loop ────────────────────────────────────────────────────
+
+    def _run_loop(self) -> None:
+        asyncio.set_event_loop(self._loop)
+        try:
+            self._loop.run_until_complete(self._main())  # type: ignore[union-attr]
+        except Exception:
+            log.debug("BridgeClient loop exited")
+        finally:
+            self._connected.clear()
+
+    async def _main(self) -> None:
+        try:
+            import websockets
+        except ImportError as exc:
+            raise RuntimeError("websockets package required: pip install websockets") from exc
+
+        try:
+            async with websockets.connect(self._url) as ws:
+                self._ws = ws
+                await ws.send(json.dumps({
+                    "type": "auth",
+                    "token": self._token,
+                    "client_type": "cli",
+                }))
+                auth_reply = json.loads(await ws.recv())
+                if auth_reply.get("type") != "auth_ok":
+                    raise ConnectionError(f"Auth failed: {auth_reply}")
+                self._connected.set()
+                log.debug("BridgeClient authenticated")
+
+                async for raw in ws:
+                    if self._disconnect_flag.is_set():
+                        break
+                    try:
+                        msg = json.loads(raw)
+                    except json.JSONDecodeError:
+                        continue
+                    self._dispatch_response(msg)
+        except Exception as exc:
+            log.debug("BridgeClient WS error: %s", exc)
+            self._connected.clear()
+            self._fail_all_pending(str(exc))
+
+    def _dispatch_response(self, msg: dict[str, Any]) -> None:
+        msg_id = msg.get("id")
+        if not msg_id:
+            return
+        with self._lock:
+            future = self._pending.pop(msg_id, None)
+        if future and not future.done():
+            future.set_result(msg)
+
+    def _fail_all_pending(self, error: str) -> None:
+        with self._lock:
+            pending = list(self._pending.items())
+            self._pending.clear()
+        exc = ConnectionError(error)
+        for _cmd_id, future in pending:
+            if not future.done():
+                future.set_exception(exc)