feat(runtime): route host_browser_request to connected chrome-extension clients

assistant/src/__tests__/conversation-confirmation-signals.test.ts

@@ -633,4 +633,83 @@ describe("restoreBrowserProxyAvailability", () => {
     // browser proxy.
     expect(bashProxy.isAvailable()).toBe(false);
   });
+
+  test("uses hostBrowserSenderOverride when set so drain-queue restores preserve the registry-routed sender", () => {
+    // Regression (PR #24129 cycle 2): the queue-drain path calls
+    // `restoreBrowserProxyAvailability()` on dequeue, which used to pass
+    // `this.sendToClient` (the SSE hub emitter) to the proxy, clobbering the
+    // chrome-extension registry-routed sender established by the POST
+    // /messages handler. The override field lets the HTTP handler pin the
+    // registry-routed sender so the drain path preserves it.
+    const sseHub: ServerMessage[] = [];
+    const registry: ServerMessage[] = [];
+    const conversation = makeConversation((msg) => sseHub.push(msg));
+    const browserProxy = new HostBrowserProxy(() => {});
+    conversation.setHostBrowserProxy(browserProxy);
+
+    // Simulate updateClient setting sendToClient to the SSE hub and
+    // marking the conversation as client-less (chrome-extension is
+    // non-interactive).
+    conversation.updateClient((msg) => sseHub.push(msg), true);
+    expect(browserProxy.isAvailable()).toBe(false);
+
+    // The HTTP handler stashes the registry-routed sender as the override.
+    const registrySender = (msg: ServerMessage) => registry.push(msg);
+    conversation.hostBrowserSenderOverride = registrySender;
+
+    // Drain-queue path calls restoreBrowserProxyAvailability — it must now
+    // prefer the override over sendToClient.
+    conversation.restoreBrowserProxyAvailability();
+    expect(browserProxy.isAvailable()).toBe(true);
+
+    // Send a frame through the proxy and verify it flows through the
+    // registry sender, not the SSE hub.
+    const internalSend = (
+      browserProxy as unknown as {
+        sendToClient: (msg: ServerMessage) => void;
+      }
+    ).sendToClient;
+    const probe: ServerMessage = {
+      type: "host_browser_cancel",
+      requestId: "probe-1",
+    } as ServerMessage;
+    internalSend(probe);
+    expect(registry).toHaveLength(1);
+    expect(sseHub.some((m) => m === probe)).toBe(false);
+  });
+
+  test("falls back to sendToClient when hostBrowserSenderOverride is cleared", () => {
+    // When a non-chrome-extension turn takes over, the HTTP handler clears
+    // the override and restoreBrowserProxyAvailability must fall back to
+    // sendToClient (the SSE hub), otherwise macOS turns would route their
+    // host_browser frames through the stale chrome-extension registry.
+    const sseHub: ServerMessage[] = [];
+    const conversation = makeConversation((msg) => sseHub.push(msg));
+    const browserProxy = new HostBrowserProxy(() => {});
+    conversation.setHostBrowserProxy(browserProxy);
+
+    // First the chrome-extension path pins the override.
+    const registry: ServerMessage[] = [];
+    conversation.hostBrowserSenderOverride = (msg) => registry.push(msg);
+    conversation.updateClient((msg) => sseHub.push(msg), true);
+    conversation.restoreBrowserProxyAvailability();
+
+    // Then a macOS handoff clears the override.
+    conversation.hostBrowserSenderOverride = undefined;
+    conversation.updateClient((msg) => sseHub.push(msg), false);
+    conversation.restoreBrowserProxyAvailability();
+
+    const internalSend = (
+      browserProxy as unknown as {
+        sendToClient: (msg: ServerMessage) => void;
+      }
+    ).sendToClient;
+    const probe: ServerMessage = {
+      type: "host_browser_cancel",
+      requestId: "probe-2",
+    } as ServerMessage;
+    internalSend(probe);
+    expect(sseHub).toContain(probe);
+    expect(registry).not.toContain(probe);
+  });
 });

assistant/src/browser-extension-relay/server.ts

@@ -28,6 +28,14 @@ interface PendingCommand {
 export interface BrowserRelayWebSocketData {
   wsType: "browser-relay";
   connectionId: string;
+  /**
+   * Guardian identity derived from the JWT claims at WebSocket upgrade
+   * time. Used by the ChromeExtensionRegistry (runtime/) to route
+   * host_browser_request frames to the correct extension. Undefined when
+   * HTTP auth is disabled (dev bypass) or when the token's sub cannot be
+   * parsed into an actor principal.
+   */
+  guardianId?: string;
 }
 
 export interface ExtensionRelayStatus {

assistant/src/daemon/conversation.ts

@@ -184,6 +184,19 @@ export class Conversation {
   /** @internal */ hostBrowserProxy?: HostBrowserProxy;
   /** @internal */ hostCuProxy?: HostCuProxy;
   /** @internal */ hostFileProxy?: HostFileProxy;
+  /**
+   * Optional override sender used by `restoreBrowserProxyAvailability` so
+   * non-SSE transports (e.g. chrome-extension, whose host_browser_request
+   * frames flow through the ChromeExtensionRegistry WebSocket rather than
+   * the SSE hub) can preserve their registry-routed sender across drain
+   * queue restores. When set, `restoreBrowserProxyAvailability()` uses this
+   * function instead of `sendToClient` so the drain-queue path doesn't
+   * clobber the chrome-extension sender with the SSE hub emitter.
+   *
+   * Populated by the POST /messages handler for chrome-extension turns and
+   * cleared when an unrelated interface takes over (see `updateClient`).
+   */
+  /** @internal */ hostBrowserSenderOverride?: (msg: ServerMessage) => void;
   /** @internal */ cesClient?: CesClient;
   /** @internal */ readonly queue = new MessageQueue();
   /** @internal */ currentActiveSurfaceId?: string;
@@ -559,11 +572,21 @@ export class Conversation {
    * it available would be to flip `hasNoClient` false, which would
    * incorrectly enable host_bash/host_file/host_cu tool gating downstream.
    *
+   * When `hostBrowserSenderOverride` is set, that function is used as the
+   * sender instead of `sendToClient`. This is required for the
+   * chrome-extension interface whose host_browser frames route through the
+   * ChromeExtensionRegistry WebSocket rather than the SSE hub: if the
+   * queue-drain path called this helper with `sendToClient`, the
+   * registry-routed sender established at turn-start would be clobbered by
+   * the SSE hub emitter and host_browser_request frames would stop
+   * reaching the extension.
+   *
    * Callers must only invoke this when they know the current interface
    * supports host_browser (see `supportsHostProxy(id, "host_browser")`).
    */
   restoreBrowserProxyAvailability(): void {
-    this.hostBrowserProxy?.updateSender(this.sendToClient, true);
+    const sender = this.hostBrowserSenderOverride ?? this.sendToClient;
+    this.hostBrowserProxy?.updateSender(sender, true);
   }
 
   setSubagentAllowedTools(tools: Set<string> | undefined): void {

assistant/src/runtime/__tests__/chrome-extension-registry.test.ts

@@ -0,0 +1,162 @@
+import { beforeEach, describe, expect, test } from "bun:test";
+
+import type { ServerMessage } from "../../daemon/message-protocol.js";
+import {
+  __resetChromeExtensionRegistryForTests,
+  type ChromeExtensionConnection,
+  ChromeExtensionRegistry,
+  getChromeExtensionRegistry,
+} from "../chrome-extension-registry.js";
+
+// Minimal structural stand-in for Bun's ServerWebSocket. Only the methods
+// the registry touches (`send`, `close`) are modeled; the rest of the Bun
+// ServerWebSocket API is out of scope for these unit tests.
+interface FakeWs {
+  send: (data: string) => number;
+  close: (code?: number, reason?: string) => void;
+  sent: string[];
+  closed: { code?: number; reason?: string }[];
+  sendShouldThrow?: boolean;
+}
+
+function makeFakeWs(): FakeWs {
+  const sent: string[] = [];
+  const closed: { code?: number; reason?: string }[] = [];
+  const ws: FakeWs = {
+    sent,
+    closed,
+    send(data: string) {
+      if (ws.sendShouldThrow) {
+        throw new Error("simulated ws.send failure");
+      }
+      sent.push(data);
+      return data.length;
+    },
+    close(code?: number, reason?: string) {
+      closed.push({ code, reason });
+    },
+  };
+  return ws;
+}
+
+function makeConnection(
+  guardianId: string,
+  id?: string,
+): { conn: ChromeExtensionConnection; fakeWs: FakeWs } {
+  const fakeWs = makeFakeWs();
+  const conn: ChromeExtensionConnection = {
+    id: id ?? crypto.randomUUID(),
+    guardianId,
+    ws: fakeWs as unknown as ChromeExtensionConnection["ws"],
+    connectedAt: Date.now(),
+  };
+  return { conn, fakeWs };
+}
+
+describe("ChromeExtensionRegistry", () => {
+  beforeEach(() => {
+    __resetChromeExtensionRegistryForTests();
+  });
+
+  test("register stores the connection under the guardianId", () => {
+    const registry = new ChromeExtensionRegistry();
+    const { conn } = makeConnection("guardian-alpha");
+    registry.register(conn);
+    expect(registry.get("guardian-alpha")).toBe(conn);
+  });
+
+  test("unregister removes the connection", () => {
+    const registry = new ChromeExtensionRegistry();
+    const { conn } = makeConnection("guardian-alpha");
+    registry.register(conn);
+    registry.unregister(conn.id);
+    expect(registry.get("guardian-alpha")).toBeUndefined();
+  });
+
+  test("unregister is a no-op when the connectionId is unknown", () => {
+    const registry = new ChromeExtensionRegistry();
+    // Should not throw even though nothing is registered.
+    expect(() => registry.unregister("unknown-connection")).not.toThrow();
+  });
+
+  test("registering a second connection for the same guardianId closes the prior one", () => {
+    const registry = new ChromeExtensionRegistry();
+    const { conn: conn1, fakeWs: fakeWs1 } = makeConnection(
+      "guardian-alpha",
+      "conn-1",
+    );
+    const { conn: conn2 } = makeConnection("guardian-alpha", "conn-2");
+    registry.register(conn1);
+    registry.register(conn2);
+    // Prior connection should have been closed with code 1000.
+    expect(fakeWs1.closed).toHaveLength(1);
+    expect(fakeWs1.closed[0].code).toBe(1000);
+    // Registry should hold the new connection.
+    expect(registry.get("guardian-alpha")).toBe(conn2);
+  });
+
+  test("registering the same connection id twice is idempotent and does not close itself", () => {
+    const registry = new ChromeExtensionRegistry();
+    const { conn, fakeWs } = makeConnection("guardian-alpha", "conn-1");
+    registry.register(conn);
+    registry.register(conn);
+    expect(fakeWs.closed).toHaveLength(0);
+    expect(registry.get("guardian-alpha")).toBe(conn);
+  });
+
+  test("send returns false when no connection exists for the guardian", () => {
+    const registry = new ChromeExtensionRegistry();
+    const msg: ServerMessage = {
+      type: "host_browser_cancel",
+      requestId: "req-1",
+    } as ServerMessage;
+    expect(registry.send("missing-guardian", msg)).toBe(false);
+  });
+
+  test("send returns true and forwards the JSON-serialized message when a connection exists", () => {
+    const registry = new ChromeExtensionRegistry();
+    const { conn, fakeWs } = makeConnection("guardian-alpha");
+    registry.register(conn);
+    const msg: ServerMessage = {
+      type: "host_browser_cancel",
+      requestId: "req-1",
+    } as ServerMessage;
+    const ok = registry.send("guardian-alpha", msg);
+    expect(ok).toBe(true);
+    expect(fakeWs.sent).toHaveLength(1);
+    const parsed = JSON.parse(fakeWs.sent[0]);
+    expect(parsed.type).toBe("host_browser_cancel");
+    expect(parsed.requestId).toBe("req-1");
+  });
+
+  test("send returns false when ws.send throws (best-effort delivery)", () => {
+    const registry = new ChromeExtensionRegistry();
+    const { conn, fakeWs } = makeConnection("guardian-alpha");
+    fakeWs.sendShouldThrow = true;
+    registry.register(conn);
+    const msg: ServerMessage = {
+      type: "host_browser_cancel",
+      requestId: "req-1",
+    } as ServerMessage;
+    expect(registry.send("guardian-alpha", msg)).toBe(false);
+  });
+
+  test("getChromeExtensionRegistry returns a module-level singleton", () => {
+    const first = getChromeExtensionRegistry();
+    const second = getChromeExtensionRegistry();
+    expect(first).toBe(second);
+  });
+
+  test("unregister after supersession does not remove the new connection", () => {
+    // When a new connection supersedes an older one, the close handler for
+    // the older socket will fire later and call unregister with the OLD id.
+    // That must not clobber the newer registration.
+    const registry = new ChromeExtensionRegistry();
+    const { conn: old } = makeConnection("guardian-alpha", "old-id");
+    const { conn: fresh } = makeConnection("guardian-alpha", "fresh-id");
+    registry.register(old);
+    registry.register(fresh);
+    registry.unregister("old-id");
+    expect(registry.get("guardian-alpha")).toBe(fresh);
+  });
+});