diff --git a/clients/chrome-extension/background/__tests__/relay-connection.test.ts b/clients/chrome-extension/background/__tests__/relay-connection.test.ts new file mode 100644 index 00000000000..d7d774a827a --- /dev/null +++ b/clients/chrome-extension/background/__tests__/relay-connection.test.ts @@ -0,0 +1,567 @@ +/** + * Tests for the RelayConnection helper. + * + * Drives the class against a fake global WebSocket so we can exercise + * the open/message/close/reconnect lifecycle without touching a real + * socket. Covers both self-hosted and cloud modes and the caller-close + * vs unexpected-close branches. + */ + +import { describe, test, expect, beforeEach, afterEach } from 'bun:test'; + +import { RelayConnection, type RelayMode } from '../relay-connection.js'; + +// ── Fake WebSocket ────────────────────────────────────────────────── + +type WsListener = (ev: { data?: unknown; code?: number; reason?: string }) => void; + +interface FakeWebSocket { + url: string; + readyState: number; + listeners: Map>; + sent: string[]; + close: (code?: number, reason?: string) => void; + send: (data: string) => void; + addEventListener: (type: string, listener: WsListener) => void; + removeEventListener: (type: string, listener: WsListener) => void; + dispatch: (type: string, ev: { data?: unknown; code?: number; reason?: string }) => void; + /** Track whether close() was called by the helper (caller-side) */ + closeCallsByCaller: Array<{ code?: number; reason?: string }>; +} + +let instances: FakeWebSocket[] = []; + +function makeFakeWebSocket(url: string): FakeWebSocket { + const listeners = new Map>(); + const sent: string[] = []; + const closeCallsByCaller: Array<{ code?: number; reason?: string }> = []; + const ws: FakeWebSocket = { + url, + readyState: 0, // CONNECTING + listeners, + sent, + closeCallsByCaller, + close(code, reason) { + closeCallsByCaller.push({ code, reason }); + ws.readyState = 3; // CLOSED + }, + send(data) { + sent.push(data); + }, + addEventListener(type, listener) { + if (!listeners.has(type)) listeners.set(type, new Set()); + listeners.get(type)!.add(listener); + }, + removeEventListener(type, listener) { + listeners.get(type)?.delete(listener); + }, + dispatch(type, ev) { + const set = listeners.get(type); + if (!set) return; + for (const l of set) l(ev); + }, + }; + return ws; +} + +// Mimic the WebSocket.OPEN etc. static constants used by the class. +function installFakeWebSocket(): void { + instances = []; + const FakeCtor = function (this: unknown, url: string) { + const instance = makeFakeWebSocket(url); + instances.push(instance); + return instance as unknown as WebSocket; + } as unknown as typeof WebSocket; + (FakeCtor as unknown as { CONNECTING: number }).CONNECTING = 0; + (FakeCtor as unknown as { OPEN: number }).OPEN = 1; + (FakeCtor as unknown as { CLOSING: number }).CLOSING = 2; + (FakeCtor as unknown as { CLOSED: number }).CLOSED = 3; + (globalThis as unknown as { WebSocket: typeof WebSocket }).WebSocket = FakeCtor; +} + +const originalWebSocket = (globalThis as unknown as { WebSocket?: typeof WebSocket }).WebSocket; + +beforeEach(() => { + installFakeWebSocket(); +}); + +afterEach(() => { + (globalThis as unknown as { WebSocket?: typeof WebSocket }).WebSocket = originalWebSocket; +}); + +/** Walk the fake-ws instance into the OPEN state and fire the open event. */ +function openSocket(ws: FakeWebSocket): void { + ws.readyState = 1; + ws.dispatch('open', {}); +} + +/** Fire a close event as if the server kicked us. */ +function closeSocket(ws: FakeWebSocket, code = 1006, reason = 'abnormal'): void { + ws.readyState = 3; + ws.dispatch('close', { code, reason }); +} + +// ── Harness ───────────────────────────────────────────────────────── + +interface Callbacks { + openCalls: number; + closeCalls: Array<{ code: number; reason: string }>; + messages: string[]; +} + +function makeCallbacks(): Callbacks { + return { openCalls: 0, closeCalls: [], messages: [] }; +} + +function makeConn(mode: RelayMode, callbacks: Callbacks, onReconnect?: () => Promise): RelayConnection { + return new RelayConnection({ + mode, + onOpen: () => { + callbacks.openCalls += 1; + }, + onMessage: (data) => { + callbacks.messages.push(data); + }, + onClose: (code, reason) => { + callbacks.closeCalls.push({ code, reason }); + }, + onReconnect, + }); +} + +// ── Tests ─────────────────────────────────────────────────────────── + +describe('RelayConnection', () => { + describe('start', () => { + test('opens a self-hosted WebSocket to the expected URL', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { + kind: 'self-hosted', + baseUrl: 'http://127.0.0.1:7830', + token: 'local-token-abc', + }, + cbs, + ); + + conn.start(); + + expect(instances.length).toBe(1); + expect(instances[0].url).toBe( + 'ws://127.0.0.1:7830/v1/browser-relay?token=local-token-abc', + ); + + openSocket(instances[0]); + expect(cbs.openCalls).toBe(1); + expect(conn.isOpen()).toBe(true); + }); + + test('opens a cloud WebSocket to the expected wss URL', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { + kind: 'cloud', + baseUrl: 'https://api.vellum.ai', + token: 'cloud-jwt-xyz', + }, + cbs, + ); + + conn.start(); + + expect(instances.length).toBe(1); + expect(instances[0].url).toBe( + 'wss://api.vellum.ai/v1/browser-relay?token=cloud-jwt-xyz', + ); + }); + + test('URL-encodes special characters in the token', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { + kind: 'cloud', + baseUrl: 'https://api.vellum.ai/', + token: 'a b+c/d=', + }, + cbs, + ); + + conn.start(); + + expect(instances.length).toBe(1); + expect(instances[0].url).toBe( + 'wss://api.vellum.ai/v1/browser-relay?token=a%20b%2Bc%2Fd%3D', + ); + }); + + test('strips a trailing slash on the base URL', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { + kind: 'cloud', + baseUrl: 'https://api.vellum.ai/', + token: 'tok', + }, + cbs, + ); + + conn.start(); + + expect(instances[0].url).not.toContain('ai//'); + expect(instances[0].url).toBe('wss://api.vellum.ai/v1/browser-relay?token=tok'); + }); + + test('omits the token query param when the caller passes null', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { + kind: 'self-hosted', + baseUrl: 'http://127.0.0.1:7830', + token: null, + }, + cbs, + ); + + conn.start(); + + expect(instances[0].url).toBe('ws://127.0.0.1:7830/v1/browser-relay'); + }); + }); + + describe('onMessage', () => { + test('forwards incoming messages to the caller', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + instances[0].dispatch('message', { data: 'hello-from-daemon' }); + instances[0].dispatch('message', { data: 'second' }); + + expect(cbs.messages).toEqual(['hello-from-daemon', 'second']); + }); + + test('stringifies non-string event data (belt-and-suspenders)', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + // Simulate a binary frame that arrived as a Blob-like object; the + // class uses String(ev.data) to keep the callback signature simple. + instances[0].dispatch('message', { data: 42 }); + + expect(cbs.messages).toEqual(['42']); + }); + }); + + describe('close', () => { + test('caller-close prevents reconnect on a subsequent unexpected close', async () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + expect(instances.length).toBe(1); + + conn.close(); + // close() synchronously calls the underlying ws.close — the fake + // dispatches the close event manually only on explicit dispatch. + closeSocket(instances[0], 1000, 'caller closed'); + + // Wait a tick to be sure any stray setTimeout didn't enqueue a + // reconnect. + await new Promise((r) => setTimeout(r, 5)); + expect(instances.length).toBe(1); + }); + + test('marks closedByCaller so isOpen returns false', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + expect(conn.isOpen()).toBe(true); + conn.close(); + expect(conn.isOpen()).toBe(false); + }); + + test('close with code 1000 on the helper forwards to the socket', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + conn.close(1000, 'bye'); + + expect(instances[0].closeCallsByCaller.length).toBe(1); + expect(instances[0].closeCallsByCaller[0].code).toBe(1000); + expect(instances[0].closeCallsByCaller[0].reason).toBe('bye'); + }); + }); + + describe('reconnect', () => { + test('unexpected close triggers reconnect after a delay', async () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + expect(instances.length).toBe(1); + + // Server-side abnormal close (e.g. the daemon restarted). + closeSocket(instances[0], 1006, 'abnormal'); + + expect(cbs.closeCalls.length).toBe(1); + expect(cbs.closeCalls[0].code).toBe(1006); + + // The reconnect is scheduled behind a real setTimeout — wait long + // enough for it to fire. The base delay is 1000ms; we tolerate + // some scheduling jitter. + await new Promise((r) => setTimeout(r, 1100)); + + expect(instances.length).toBe(2); + expect(instances[1].url).toBe( + 'ws://127.0.0.1:7830/v1/browser-relay?token=t', + ); + + // Clean up. + conn.close(); + }); + + test('normal close (code 1000) does NOT call onReconnect', async () => { + const cbs = makeCallbacks(); + let reconnectCalls = 0; + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + async () => { + reconnectCalls += 1; + return 'new-token'; + }, + ); + + conn.start(); + openSocket(instances[0]); + + // Normal close — should still schedule a reconnect but without + // calling the refresh hook. + closeSocket(instances[0], 1000, 'normal'); + await new Promise((r) => setTimeout(r, 1100)); + + expect(reconnectCalls).toBe(0); + expect(instances.length).toBe(2); + + conn.close(); + }); + + test('onReconnect replaces the token used for the next URL', async () => { + const cbs = makeCallbacks(); + let refreshCalls = 0; + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 'old' }, + cbs, + async () => { + refreshCalls += 1; + return 'fresh-token'; + }, + ); + + conn.start(); + openSocket(instances[0]); + expect(instances[0].url).toContain('token=old'); + + closeSocket(instances[0], 4001, 'auth rotated'); + await new Promise((r) => setTimeout(r, 1100)); + + expect(refreshCalls).toBe(1); + expect(instances.length).toBe(2); + expect(instances[1].url).toContain('token=fresh-token'); + + conn.close(); + }); + + test('onReconnect returning void leaves the existing token in place', async () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 'keep' }, + cbs, + async () => { + // no return → void + }, + ); + + conn.start(); + openSocket(instances[0]); + closeSocket(instances[0], 1006, 'abnormal'); + await new Promise((r) => setTimeout(r, 1100)); + + expect(instances.length).toBe(2); + expect(instances[1].url).toContain('token=keep'); + + conn.close(); + }); + + test('close called before scheduled reconnect fires prevents reconnection', async () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + closeSocket(instances[0], 1006, 'abnormal'); + + // Cancel before the reconnect timer fires. + conn.close(); + + await new Promise((r) => setTimeout(r, 1100)); + + // No second socket should have been constructed. + expect(instances.length).toBe(1); + }); + }); + + describe('setMode', () => { + test('closes the current socket and opens a new one for the new mode', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + expect(instances.length).toBe(1); + expect(instances[0].url).toContain('ws://127.0.0.1'); + + conn.setMode({ kind: 'cloud', baseUrl: 'https://api.vellum.ai', token: 'cloud-jwt' }); + + // Old socket was closed by the caller; new one was constructed + // against the cloud URL. + expect(instances[0].closeCallsByCaller.length).toBe(1); + expect(instances.length).toBe(2); + expect(instances[1].url).toBe( + 'wss://api.vellum.ai/v1/browser-relay?token=cloud-jwt', + ); + + conn.close(); + }); + + test('updates the mode getter', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + conn.start(); + expect(conn.mode.kind).toBe('self-hosted'); + + conn.setMode({ kind: 'cloud', baseUrl: 'https://api.vellum.ai', token: 'c' }); + expect(conn.mode.kind).toBe('cloud'); + + conn.close(); + }); + + test('stale close event from a superseded socket does not clear the new ws or schedule reconnect', async () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + expect(instances.length).toBe(1); + const oldSocket = instances[0]; + + // Switch modes mid-flight: the helper closes socket A (oldSocket) + // and constructs socket B (newSocket) for the cloud gateway. We + // keep newSocket in CONNECTING so we can observe the state that + // would be disturbed by a stale close event. + conn.setMode({ + kind: 'cloud', + baseUrl: 'https://api.vellum.ai', + token: 'cloud-jwt', + }); + expect(instances.length).toBe(2); + const newSocket = instances[1]; + expect(newSocket.url).toBe( + 'wss://api.vellum.ai/v1/browser-relay?token=cloud-jwt', + ); + expect(conn.mode.kind).toBe('cloud'); + + // Now simulate the asynchronous close event that socket A fires + // after setMode already re-pointed this.ws at socket B. The + // helper should ignore it entirely: this.ws stays pinned to + // newSocket, no reconnect is queued, and onClose is NOT invoked + // (we already told the caller we switched modes). + closeSocket(oldSocket, 1006, 'stale'); + + // No onClose call — the close event came from a superseded socket. + expect(cbs.closeCalls.length).toBe(0); + + // Open the new socket to confirm the helper still holds a valid + // reference to it. If the stale close had nulled out this.ws we'd + // see isOpen() stay false here. + openSocket(newSocket); + expect(conn.isOpen()).toBe(true); + + // Wait long enough that any reconnect timer would have fired. + await new Promise((r) => setTimeout(r, 1100)); + + // Still only the original two sockets — no spurious reconnect. + expect(instances.length).toBe(2); + + conn.close(); + }); + }); + + describe('send', () => { + test('writes to the underlying socket when OPEN', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + openSocket(instances[0]); + conn.send('hello-daemon'); + + expect(instances[0].sent).toEqual(['hello-daemon']); + }); + + test('is a no-op before the socket is OPEN', () => { + const cbs = makeCallbacks(); + const conn = makeConn( + { kind: 'self-hosted', baseUrl: 'http://127.0.0.1:7830', token: 't' }, + cbs, + ); + + conn.start(); + // ws.readyState is still CONNECTING (0). + conn.send('too-early'); + expect(instances[0].sent).toEqual([]); + }); + }); +}); diff --git a/clients/chrome-extension/background/relay-connection.ts b/clients/chrome-extension/background/relay-connection.ts new file mode 100644 index 00000000000..3b5b7c45546 --- /dev/null +++ b/clients/chrome-extension/background/relay-connection.ts @@ -0,0 +1,248 @@ +/** + * Relay WebSocket connection helper. + * + * Extracted from worker.ts so we can share the open/close/reconnect + * lifecycle between the two relay transports: + * + * - `self-hosted` — ws://127.0.0.1:/v1/browser-relay, token minted + * by the local daemon (legacy path; default for back-compat). + * - `cloud` — wss:///v1/browser-relay, token from + * the cloud OAuth flow (see cloud-auth.ts). + * + * The class only knows how to open the socket, forward incoming messages + * to the caller, and reconnect after unexpected closes. It does NOT parse + * relay messages — worker.ts owns the envelope dispatch (ExtensionCommand + * + host_browser_request via the PR 9 dispatcher) via the `onMessage` + * callback. + */ + +/** Reconnect backoff bounds mirror the legacy inline worker.ts values. */ +const RECONNECT_BASE_MS = 1_000; +const RECONNECT_MAX_MS = 30_000; + +/** WebSocket close codes that represent intentional, non-error closures. */ +const NORMAL_CLOSE_CODES = new Set([1000, 1001]); + +/** + * Connection mode with the corresponding base URL + bearer token. The + * base URL is normalised by {@link RelayConnection.buildUrl}: any + * `http(s)://` scheme is rewritten to `ws(s)://` and a trailing slash is + * stripped. Pass the daemon's HTTP origin for self-hosted mode and the + * cloud gateway's HTTPS origin for cloud mode — the class figures out + * the WebSocket scheme. + */ +export type RelayMode = + | { kind: 'self-hosted'; baseUrl: string; token: string | null } + | { kind: 'cloud'; baseUrl: string; token: string | null }; + +export interface RelayConnectionDeps { + /** + * Mode + token. The token is pre-fetched by the caller (so the caller + * can decide whether to skip the connection entirely when there's no + * token yet, e.g. before cloud sign-in or before self-hosted pairing). + */ + mode: RelayMode; + /** Invoked with the raw string payload for every incoming message. */ + onMessage: (data: string) => void; + /** Invoked when the socket transitions to OPEN. */ + onOpen: () => void; + /** Invoked when the socket closes (user-initiated or unexpected). */ + onClose: (code: number, reason: string) => void; + /** + * Optional: invoked right before a reconnect attempt is scheduled for + * an unexpected close. Callers use this to refresh stale tokens before + * the next `start()` attempt. If this returns a string, the new token + * replaces `mode.token` for the next URL build. + */ + onReconnect?: () => Promise; +} + +/** + * Long-lived WebSocket helper. One instance per live relay session — + * switching modes closes the current socket and constructs a new one. + */ +export class RelayConnection { + private ws: WebSocket | null = null; + private deps: RelayConnectionDeps; + private reconnectTimer: ReturnType | null = null; + private reconnectDelay = RECONNECT_BASE_MS; + private closedByCaller = false; + + constructor(deps: RelayConnectionDeps) { + this.deps = deps; + } + + /** Current connection mode (read-only; use `setMode` to switch). */ + get mode(): RelayMode { + return this.deps.mode; + } + + /** Is the underlying socket currently in the OPEN readyState? */ + isOpen(): boolean { + return this.ws !== null && this.ws.readyState === WebSocket.OPEN; + } + + /** Begin (or resume) connecting. Idempotent while already connected. */ + start(): void { + this.closedByCaller = false; + this.reconnectDelay = RECONNECT_BASE_MS; + this.connect(); + } + + /** + * Swap the connection mode / token without destroying the class + * instance. The current socket is closed cleanly and a fresh one is + * opened for the new mode. Used by the popup's mode switcher. + */ + setMode(mode: RelayMode): void { + this.deps = { ...this.deps, mode }; + // Tear down the current socket without marking the caller as having + // closed us permanently — `start()` below re-arms shouldConnect. + if (this.reconnectTimer !== null) { + clearTimeout(this.reconnectTimer); + this.reconnectTimer = null; + } + if (this.ws) { + try { + this.ws.close(1000, 'mode switched'); + } catch { + /* ignore */ + } + this.ws = null; + } + this.start(); + } + + /** + * Send a raw string payload. No-op if the socket is not currently OPEN + * — matches the existing worker.ts semantics where heartbeats and + * responses silently drop when the socket is mid-reconnect. + */ + send(data: string): void { + if (this.ws && this.ws.readyState === WebSocket.OPEN) { + this.ws.send(data); + } + } + + /** + * Close the socket permanently. After this the connection will not + * reconnect on its own; call `start()` again to resume. + */ + close(code = 1000, reason = 'closed by caller'): void { + this.closedByCaller = true; + if (this.reconnectTimer !== null) { + clearTimeout(this.reconnectTimer); + this.reconnectTimer = null; + } + if (this.ws) { + try { + this.ws.close(code, reason); + } catch { + /* ignore */ + } + this.ws = null; + } + } + + // ── Internals ───────────────────────────────────────────────────── + + private connect(): void { + if (this.ws && (this.ws.readyState === WebSocket.OPEN || this.ws.readyState === WebSocket.CONNECTING)) { + return; + } + + const url = this.buildUrl(); + // Capture a local reference to the socket so that every listener + // can verify it is still the active one before mutating shared + // state. Without this, a `setMode()` that closes socket A and + // immediately opens socket B can get A's asynchronous close event + // delivered afterward — that stale event would otherwise clear the + // reference to B and schedule a spurious reconnect. + const ws = new WebSocket(url); + this.ws = ws; + + ws.addEventListener('open', () => { + if (this.ws !== ws) return; // stale event from a superseded socket + this.reconnectDelay = RECONNECT_BASE_MS; + this.deps.onOpen(); + }); + + ws.addEventListener('message', (event: MessageEvent) => { + if (this.ws !== ws) return; // stale event from a superseded socket + this.deps.onMessage(String(event.data)); + }); + + ws.addEventListener('close', (event: CloseEvent) => { + if (this.ws !== ws) return; // stale event from a superseded socket + const code = event.code; + const reason = event.reason; + this.ws = null; + this.deps.onClose(code, reason); + if (!this.closedByCaller) { + if (!NORMAL_CLOSE_CODES.has(code)) { + this.scheduleReconnectWithRefresh(); + } else { + this.scheduleReconnect(); + } + } + }); + + ws.addEventListener('error', () => { + if (this.ws !== ws) return; // stale event from a superseded socket + // A close event will follow — nothing to do here beyond letting + // the socket transition into CLOSING/CLOSED so we can reconnect. + }); + } + + /** Build the WebSocket URL from the current mode. */ + private buildUrl(): string { + const { mode } = this.deps; + const base = mode.baseUrl.replace(/\/$/, ''); + const wsBase = base.replace(/^http/, 'ws'); + const url = `${wsBase}/v1/browser-relay`; + if (mode.token) { + return `${url}?token=${encodeURIComponent(mode.token)}`; + } + return url; + } + + private scheduleReconnect(): void { + if (this.reconnectTimer !== null) return; + const delay = this.reconnectDelay; + this.reconnectTimer = setTimeout(() => { + this.reconnectTimer = null; + if (!this.closedByCaller) this.connect(); + }, delay); + this.reconnectDelay = Math.min(this.reconnectDelay * 2, RECONNECT_MAX_MS); + } + + /** + * Unexpected close path: give the caller a chance to refresh the + * token (e.g. the self-hosted daemon rotated its edge JWT, or the + * cloud OAuth flow expired) before the next connect attempt. + */ + private scheduleReconnectWithRefresh(): void { + if (this.reconnectTimer !== null) return; + const delay = this.reconnectDelay; + this.reconnectTimer = setTimeout(async () => { + this.reconnectTimer = null; + if (this.closedByCaller) return; + if (this.deps.onReconnect) { + try { + const newToken = await this.deps.onReconnect(); + if (typeof newToken === 'string') { + this.deps = { + ...this.deps, + mode: { ...this.deps.mode, token: newToken }, + }; + } + } catch { + // Refresh failures fall through to a bare reconnect attempt — + // the server will reject the handshake and we'll loop. + } + } + if (!this.closedByCaller) this.connect(); + }, delay); + this.reconnectDelay = Math.min(this.reconnectDelay * 2, RECONNECT_MAX_MS); + } +} diff --git a/clients/chrome-extension/background/worker.ts b/clients/chrome-extension/background/worker.ts index a421839d634..07d660ff2a5 100644 --- a/clients/chrome-extension/background/worker.ts +++ b/clients/chrome-extension/background/worker.ts @@ -1,13 +1,32 @@ /** * Chrome MV3 service worker — browser-relay bridge. * - * Connects to ws://127.0.0.1:/v1/browser-relay and dispatches - * ExtensionCommands from the server to browser APIs, sending back - * ExtensionResponses. + * Connects to either + * - the local daemon's browser-relay endpoint + * (`ws://127.0.0.1:/v1/browser-relay`), or + * - the cloud gateway's browser-relay endpoint + * (`wss:///v1/browser-relay`) + * + * depending on the `vellum.relayMode` key in chrome.storage.local + * (default `"self-hosted"` for back-compat). Both transports share the + * same envelope vocabulary — the choice is strictly about where the + * socket points and which token is presented on the handshake. + * + * Once connected, the worker dispatches incoming server messages: + * - `host_browser_request` / `host_browser_cancel` envelopes are + * routed to the CDP proxy dispatcher (Phase 2 PR 9, gated behind + * the `vellum.cdpProxyEnabled` feature flag). + * - Every other payload is treated as a legacy `ExtensionCommand` + * and dispatched to the existing browser-API handlers. */ import type { ExtensionCommand, ExtensionResponse, ExtensionHeartbeat } from '../../../assistant/src/browser-extension-relay/protocol.js'; -import { signInCloud, type CloudAuthConfig, type StoredCloudToken } from './cloud-auth.js'; +import { + signInCloud, + getStoredToken as getStoredCloudToken, + type CloudAuthConfig, + type StoredCloudToken, +} from './cloud-auth.js'; import { bootstrapLocalToken, type StoredLocalToken, @@ -19,31 +38,37 @@ import { type HostBrowserCancelEnvelope, type HostBrowserResultEnvelope, } from './host-browser-dispatcher.js'; +import { RelayConnection, type RelayMode } from './relay-connection.js'; // Cloud OAuth defaults — kept here so the popup can stay a thin client and the // service worker is the single owner of the launchWebAuthFlow lifecycle. This // avoids the MV3 popup teardown race where closing the popup mid-auth kills // the awaited promise before the token is persisted. -// -// PR 14 will plumb these through config; hard-coded for the Phase 2 skeleton. const CLOUD_GATEWAY_BASE_URL = 'https://api.vellum.ai'; const CLOUD_OAUTH_CLIENT_ID = 'vellum-chrome-extension'; const DEFAULT_RELAY_PORT = 7830; const HEARTBEAT_INTERVAL_MS = 30_000; -const RECONNECT_BASE_MS = 1_000; -const RECONNECT_MAX_MS = 30_000; const EXTENSION_VERSION = chrome.runtime.getManifest().version; -let ws: WebSocket | null = null; -let reconnectDelay = RECONNECT_BASE_MS; +// ── Mode selection (Phase 2 PR 14) ───────────────────────────────── +// +// Existing installs have no `vellum.relayMode` key and must keep using +// the local daemon transport. New installs can flip to cloud via the +// popup radio group. +const RELAY_MODE_KEY = 'vellum.relayMode'; +type RelayModeKind = 'self-hosted' | 'cloud'; + +function isRelayModeKind(v: unknown): v is RelayModeKind { + return v === 'self-hosted' || v === 'cloud'; +} + +let relayMode: RelayModeKind = 'self-hosted'; +let relayConnection: RelayConnection | null = null; let heartbeatTimer: ReturnType | null = null; let shouldConnect = false; -/** WebSocket close codes that represent intentional, non-error closures. */ -const NORMAL_CLOSE_CODES = new Set([1000, 1001]); - // ── Host browser dispatcher (Phase 2 PR 9) ────────────────────────── // // Feature-flagged behind `vellum.cdpProxyEnabled` in chrome.storage.local. @@ -132,68 +157,151 @@ async function refreshToken(): Promise { } } -// ── WebSocket lifecycle ───────────────────────────────────────────── +// ── Relay connection lifecycle ────────────────────────────────────── -async function connect(): Promise { - if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) { - return; - } +async function loadRelayMode(): Promise { + const result = await chrome.storage.local.get(RELAY_MODE_KEY); + const stored = result[RELAY_MODE_KEY]; + return isRelayModeKind(stored) ? stored : 'self-hosted'; +} +async function buildRelayModeConfig(kind: RelayModeKind): Promise { + if (kind === 'cloud') { + const stored = await getStoredCloudToken(); + return { + kind: 'cloud', + baseUrl: CLOUD_GATEWAY_BASE_URL, + token: stored?.token ?? null, + }; + } + // Self-hosted: re-use the existing local-token flow. The plan explicitly + // defers the switch to PR 13's getStoredLocalToken() to a follow-up. const [token, port] = await Promise.all([getBearerToken(), getRelayPort()]); - const relayUrlBase = `ws://127.0.0.1:${port}/v1/browser-relay`; - const url = token ? `${relayUrlBase}?token=${encodeURIComponent(token)}` : relayUrlBase; - - ws = new WebSocket(url); + return { + kind: 'self-hosted', + baseUrl: `http://127.0.0.1:${port}`, + token, + }; +} - ws.addEventListener('open', () => { - console.log('[vellum-relay] Connected to relay server'); - reconnectDelay = RECONNECT_BASE_MS; - startHeartbeat(); +/** + * Wire a RelayConnection up with the worker's message/open/close + * callbacks. Does NOT start it. + */ +function createRelayConnection(mode: RelayMode): RelayConnection { + return new RelayConnection({ + mode, + onOpen: () => { + console.log(`[vellum-relay] Connected (${mode.kind})`); + startHeartbeat(); + }, + onMessage: (data) => { + handleServerMessage(data); + }, + onClose: (code, reason) => { + console.log(`[vellum-relay] Disconnected (code=${code}, reason=${reason || 'n/a'})`); + stopHeartbeat(); + }, + onReconnect: async () => { + // Self-hosted: attempt to mint a fresh gateway token. Cloud: no-op + // for now — the cloud token is stored independently via OAuth and + // we'd rather surface the failure to the user than silently loop. + if (mode.kind === 'self-hosted') { + const ok = await refreshToken(); + if (ok) { + const refreshed = await getBearerToken(); + return refreshed; + } + } + }, }); +} - ws.addEventListener('message', (event) => { - handleServerMessage(event.data as string); - }); +/** + * Thrown by `connect()` when the selected relay mode has no usable + * token yet. Callers (e.g. the popup connect handler) surface the + * message verbatim to the user so they can take action — signing in + * to cloud or re-pairing the local daemon — instead of seeing a + * silent no-op after pressing "Connect". + */ +class MissingTokenError extends Error { + constructor(message: string) { + super(message); + this.name = 'MissingTokenError'; + } +} - ws.addEventListener('close', (event) => { - console.log(`[vellum-relay] Disconnected (code=${event.code}). Reconnecting in ${reconnectDelay}ms…`); - stopHeartbeat(); - ws = null; - if (shouldConnect) { - if (!NORMAL_CLOSE_CODES.has(event.code)) { - // Any unexpected close (including 1006 from failed HTTP 401 handshakes, - // 1008, 4001, etc.) — attempt a token refresh before reconnecting. - refreshToken().then(() => scheduleReconnect()); - } else { - scheduleReconnect(); - } - } - }); +function missingTokenMessage(kind: RelayModeKind): string { + if (kind === 'cloud') { + return 'Sign in with Vellum (cloud) before connecting'; + } + return 'Pair the Vellum daemon (self-hosted) before connecting'; +} - ws.addEventListener('error', () => { - // close event will follow; just log - console.warn('[vellum-relay] WebSocket error'); - }); +async function connect(): Promise { + if (relayConnection && relayConnection.isOpen()) return; + const mode = await buildRelayModeConfig(relayMode); + if (!mode.token) { + const msg = missingTokenMessage(mode.kind); + console.warn(`[vellum-relay] ${msg}`); + throw new MissingTokenError(msg); + } + // Tear down any stale instance before constructing a new one. This + // keeps the close/reconnect lifecycle simple — one RelayConnection + // per live socket, no hidden state carried across mode switches. + if (relayConnection) { + relayConnection.close(1000, 'reconfigured'); + } + relayConnection = createRelayConnection(mode); + relayConnection.start(); } -function scheduleReconnect(): void { - setTimeout(() => { - if (shouldConnect) connect(); - }, reconnectDelay); - reconnectDelay = Math.min(reconnectDelay * 2, RECONNECT_MAX_MS); +function disconnect(): void { + stopHeartbeat(); + if (relayConnection) { + relayConnection.close(1000, 'User disconnected'); + relayConnection = null; + } +} + +/** + * Handle a runtime switch of `vellum.relayMode` (e.g. the popup radio + * group flipped). Closes any current socket and opens a new one in the + * new mode — see plan PR 14 step 2. + */ +async function applyModeChange(newKind: RelayModeKind): Promise { + if (newKind === relayMode) return; + relayMode = newKind; + if (!shouldConnect) return; + disconnect(); + try { + await connect(); + } catch (err) { + // The user switched modes before signing in / pairing. Leave the + // extension disconnected and let the next user-initiated connect + // bubble the error up through the popup message handler. + if (err instanceof MissingTokenError) { + shouldConnect = false; + console.warn( + `[vellum-relay] Mode switch to ${newKind} left disconnected: ${err.message}`, + ); + return; + } + throw err; + } } function startHeartbeat(): void { stopHeartbeat(); heartbeatTimer = setInterval(async () => { - if (!ws || ws.readyState !== WebSocket.OPEN) return; + if (!relayConnection || !relayConnection.isOpen()) return; const tabs = await chrome.tabs.query({}); const heartbeat: ExtensionHeartbeat = { type: 'heartbeat', extensionVersion: EXTENSION_VERSION, connectedTabs: tabs.length, }; - ws.send(JSON.stringify(heartbeat)); + relayConnection.send(JSON.stringify(heartbeat)); }, HEARTBEAT_INTERVAL_MS); } @@ -205,8 +313,8 @@ function stopHeartbeat(): void { } function sendResponse(response: ExtensionResponse): void { - if (ws && ws.readyState === WebSocket.OPEN) { - ws.send(JSON.stringify(response)); + if (relayConnection && relayConnection.isOpen()) { + relayConnection.send(JSON.stringify(response)); } } @@ -422,21 +530,31 @@ async function handleScreenshot(cmd: ExtensionCommand): Promise // ── Extension message listener (from popup) ───────────────────────── -chrome.runtime.onMessage.addListener((message, _sender, sendResponse) => { +chrome.runtime.onMessage.addListener((message, _sender, sendResponseFn) => { if (message.type === 'connect') { shouldConnect = true; - connect().then(() => sendResponse({ ok: true })).catch((err) => sendResponse({ ok: false, error: String(err) })); + connect() + .then(() => sendResponseFn({ ok: true })) + .catch((err) => { + // Reset shouldConnect so a subsequent storage change or + // bootstrap doesn't silently retry a doomed connect. The user + // will press Connect again after signing in / pairing. + shouldConnect = false; + const errorMessage = err instanceof Error ? err.message : String(err); + sendResponseFn({ ok: false, error: errorMessage }); + }); return true; // async } if (message.type === 'disconnect') { shouldConnect = false; - ws?.close(1000, 'User disconnected'); - sendResponse({ ok: true }); + disconnect(); + sendResponseFn({ ok: true }); return false; } if (message.type === 'get_status') { - sendResponse({ - connected: ws !== null && ws.readyState === WebSocket.OPEN, + sendResponseFn({ + connected: relayConnection !== null && relayConnection.isOpen(), + mode: relayMode, }); return false; } @@ -451,8 +569,8 @@ chrome.runtime.onMessage.addListener((message, _sender, sendResponse) => { typeof message.clientId === 'string' ? message.clientId : CLOUD_OAUTH_CLIENT_ID, }; signInCloud(config) - .then((stored: StoredCloudToken) => sendResponse({ ok: true, token: stored })) - .catch((err) => sendResponse({ ok: false, error: err instanceof Error ? err.message : String(err) })); + .then((stored: StoredCloudToken) => sendResponseFn({ ok: true, token: stored })) + .catch((err) => sendResponseFn({ ok: false, error: err instanceof Error ? err.message : String(err) })); return true; // async } if (message.type === 'self-hosted-pair') { @@ -469,14 +587,33 @@ chrome.runtime.onMessage.addListener((message, _sender, sendResponse) => { }); // Auto-connect on service worker start if previously connected. -// Refresh the token first so we don't reconnect with stale credentials. -chrome.storage.local.get('autoConnect').then(async (result) => { - if (result.autoConnect === true) { - shouldConnect = true; +// Refresh the self-hosted token first so we don't reconnect with stale +// credentials — cloud-mode auto-connect just reads the stored OAuth +// token and trusts the caller to re-sign in if it's expired. +async function bootstrap(): Promise { + relayMode = await loadRelayMode(); + const { autoConnect } = await chrome.storage.local.get('autoConnect'); + if (autoConnect !== true) return; + shouldConnect = true; + if (relayMode === 'self-hosted') { await refreshToken(); - connect(); } -}); + try { + await connect(); + } catch (err) { + // A missing token at auto-connect time is not a hard failure — + // the user will see the disconnected state in the popup and can + // sign in / pair to try again. Log and move on. + if (err instanceof MissingTokenError) { + shouldConnect = false; + console.warn(`[vellum-relay] Skipping auto-connect: ${err.message}`); + return; + } + throw err; + } +} + +bootstrap(); // Load the CDP proxy feature flag at startup. Missing / non-boolean values // are treated as false so existing deployments exhibit no behavior change. @@ -488,9 +625,8 @@ chrome.storage.local.get(CDP_PROXY_ENABLED_KEY).then((result) => { } }); -// Keep the flag live-updatable from the popup without requiring the service -// worker to restart. `chrome.storage.onChanged` fires in the same service -// worker context the value is set from, which is perfect here. +// Keep feature flag + relay mode live-updatable from the popup without +// requiring the service worker to restart. chrome.storage.onChanged.addListener((changes, areaName) => { if (areaName !== 'local') return; if (CDP_PROXY_ENABLED_KEY in changes) { @@ -500,4 +636,11 @@ chrome.storage.onChanged.addListener((changes, areaName) => { `[vellum-relay] CDP proxy feature flag updated: ${cdpProxyEnabled}`, ); } + if (RELAY_MODE_KEY in changes) { + const newValue = changes[RELAY_MODE_KEY]?.newValue; + if (isRelayModeKind(newValue)) { + console.log(`[vellum-relay] Relay mode updated: ${newValue}`); + void applyModeChange(newValue); + } + } }); diff --git a/clients/chrome-extension/popup/popup.html b/clients/chrome-extension/popup/popup.html index 54890e86b47..46e3217c10b 100644 --- a/clients/chrome-extension/popup/popup.html +++ b/clients/chrome-extension/popup/popup.html @@ -158,6 +158,33 @@ cursor: pointer; } + .mode-group { + margin-bottom: 14px; + } + + .mode-radio-row { + display: flex; + gap: 14px; + margin-top: 4px; + } + + .mode-radio-row label { + display: flex; + align-items: center; + gap: 6px; + font-size: 12px; + font-weight: 500; + color: #374151; + cursor: pointer; + user-select: none; + margin-bottom: 0; + } + + .mode-radio-row input[type="radio"] { + margin: 0; + cursor: pointer; + } + .hint { font-size: 11px; color: #9ca3af; @@ -198,6 +225,20 @@

Vellum Relay

+
+ +
+ + +
+
+
diff --git a/clients/chrome-extension/popup/popup.ts b/clients/chrome-extension/popup/popup.ts index 01cb61f8b5f..6ae81046905 100644 --- a/clients/chrome-extension/popup/popup.ts +++ b/clients/chrome-extension/popup/popup.ts @@ -35,8 +35,12 @@ const cloudStatus = document.getElementById('cloud-status') as HTMLParagraphElem const btnPairLocal = document.getElementById('btn-pair-local') as HTMLButtonElement; const localStatus = document.getElementById('local-status') as HTMLParagraphElement; const cdpProxyToggle = document.getElementById('cdp-proxy-toggle') as HTMLInputElement; +const modeSelfHosted = document.getElementById('mode-self-hosted') as HTMLInputElement; +const modeCloud = document.getElementById('mode-cloud') as HTMLInputElement; const CDP_PROXY_ENABLED_KEY = 'vellum.cdpProxyEnabled'; +const RELAY_MODE_KEY = 'vellum.relayMode'; +type RelayModeKind = 'self-hosted' | 'cloud'; let manualMode = false; @@ -312,3 +316,43 @@ chrome.storage.local.get(CDP_PROXY_ENABLED_KEY).then((result) => { cdpProxyToggle.addEventListener('change', async () => { await chrome.storage.local.set({ [CDP_PROXY_ENABLED_KEY]: cdpProxyToggle.checked }); }); + +// ── Relay mode switcher (Phase 2 PR 14) ──────────────────────────── +// +// Flips `vellum.relayMode` in chrome.storage.local between "self-hosted" +// (default, back-compat) and "cloud". The service worker listens for +// storage changes via chrome.storage.onChanged and closes the current +// socket + reopens a new one against the selected transport. + +function isRelayModeKind(v: unknown): v is RelayModeKind { + return v === 'self-hosted' || v === 'cloud'; +} + +chrome.storage.local.get(RELAY_MODE_KEY).then((result) => { + const stored = result[RELAY_MODE_KEY]; + const mode: RelayModeKind = isRelayModeKind(stored) ? stored : 'self-hosted'; + if (mode === 'cloud') { + modeCloud.checked = true; + } else { + modeSelfHosted.checked = true; + } +}); + +async function handleModeChange(newMode: RelayModeKind): Promise { + await chrome.storage.local.set({ [RELAY_MODE_KEY]: newMode }); + // The service worker reacts to the storage change via + // chrome.storage.onChanged — we don't need to send an explicit + // disconnect/connect message here. +} + +modeSelfHosted.addEventListener('change', () => { + if (modeSelfHosted.checked) { + void handleModeChange('self-hosted'); + } +}); + +modeCloud.addEventListener('change', () => { + if (modeCloud.checked) { + void handleModeChange('cloud'); + } +}); diff --git a/clients/chrome-extension/tsconfig.json b/clients/chrome-extension/tsconfig.json index a575b64e6b3..7476c828e31 100644 --- a/clients/chrome-extension/tsconfig.json +++ b/clients/chrome-extension/tsconfig.json @@ -18,6 +18,7 @@ "background/cdp-proxy.ts", "background/cloud-auth.ts", "background/host-browser-dispatcher.ts", + "background/relay-connection.ts", "background/self-hosted-auth.ts", "background/__tests__/**/*.ts", "types/**/*.d.ts"