Skip to content

Commit 4f5fcd9

Browse files
GreyCGreyC
authored
fix(extension): keep active daemon websocket
Keep stale Browser Bridge WebSocket events from clobbering the active daemon connection.\n\nCo-authored-by: Jeff Chen <jeff@adtiming.com>
1 parent 40b2f75 commit 4f5fcd9

3 files changed

Lines changed: 131 additions & 24 deletions

File tree

‎extension/dist/background.js‎

Lines changed: 25 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -671,13 +671,21 @@ const _origLog = console.log.bind(console);
671671
const _origWarn = console.warn.bind(console);
672672
const _origError = console.error.bind(console);
673673
function forwardLog(level, args) {
674-
if (!ws || ws.readyState !== WebSocket.OPEN) return;
675674
try {
676675
const msg = args.map((a) => typeof a === "string" ? a : JSON.stringify(a)).join(" ");
677-
ws.send(JSON.stringify({ type: "log", level, msg, ts: Date.now() }));
676+
safeSend(ws, { type: "log", level, msg, ts: Date.now() });
678677
} catch {
679678
}
680679
}
680+
function safeSend(socket, payload) {
681+
if (!socket || socket.readyState !== WebSocket.OPEN) return false;
682+
try {
683+
socket.send(JSON.stringify(payload));
684+
return true;
685+
} catch {
686+
return false;
687+
}
688+
}
681689
console.log = (...args) => {
682690
_origLog(...args);
683691
forwardLog("info", args);
@@ -698,44 +706,50 @@ async function connect() {
698706
} catch {
699707
return;
700708
}
709+
let thisWs;
701710
try {
702711
const contextId = await getCurrentContextId();
703-
ws = new WebSocket(DAEMON_WS_URL);
712+
thisWs = new WebSocket(DAEMON_WS_URL);
713+
ws = thisWs;
704714
currentContextId = contextId;
705715
} catch {
706716
scheduleReconnect();
707717
return;
708718
}
709-
ws.onopen = () => {
719+
thisWs.onopen = () => {
720+
if (ws !== thisWs) return;
710721
console.log("[opencli] Connected to daemon");
711722
reconnectAttempts = 0;
712723
if (reconnectTimer) {
713724
clearTimeout(reconnectTimer);
714725
reconnectTimer = null;
715726
}
716-
ws?.send(JSON.stringify({
727+
safeSend(thisWs, {
717728
type: "hello",
718729
contextId: currentContextId,
719730
version: chrome.runtime.getManifest().version,
720731
compatRange: ">=1.7.0"
721-
}));
732+
});
722733
};
723-
ws.onmessage = async (event) => {
734+
thisWs.onmessage = async (event) => {
735+
if (ws !== thisWs) return;
724736
try {
725737
const command = JSON.parse(event.data);
726738
const result = await handleCommand(command);
727-
ws?.send(JSON.stringify(result));
739+
if (ws !== thisWs) return;
740+
safeSend(thisWs, result);
728741
} catch (err) {
729742
console.error("[opencli] Message handling error:", err);
730743
}
731744
};
732-
ws.onclose = () => {
745+
thisWs.onclose = () => {
746+
if (ws !== thisWs) return;
733747
console.log("[opencli] Disconnected from daemon");
734748
ws = null;
735749
scheduleReconnect();
736750
};
737-
ws.onerror = () => {
738-
ws?.close();
751+
thisWs.onerror = () => {
752+
thisWs.close();
739753
};
740754
}
741755
const MAX_EAGER_ATTEMPTS = 6;

‎extension/src/background.test.ts‎

Lines changed: 80 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,15 +31,23 @@ const adapterKey = (session: string): string => leaseKey('adapter', session);
3131
class MockWebSocket {
3232
static OPEN = 1;
3333
static CONNECTING = 0;
34+
static CLOSED = 3;
35+
static instances: MockWebSocket[] = [];
3436
readyState = MockWebSocket.CONNECTING;
37+
sent: string[] = [];
3538
onopen: (() => void) | null = null;
3639
onmessage: ((event: { data: string }) => void) | null = null;
3740
onclose: (() => void) | null = null;
3841
onerror: (() => void) | null = null;
3942

40-
constructor(_url: string) {}
41-
send(_data: string): void {}
43+
constructor(_url: string) {
44+
MockWebSocket.instances.push(this);
45+
}
46+
send(data: string): void {
47+
this.sent.push(data);
48+
}
4249
close(): void {
50+
this.readyState = MockWebSocket.CLOSED;
4351
this.onclose?.();
4452
}
4553
}
@@ -194,6 +202,7 @@ describe('background tab isolation', () => {
194202
beforeEach(() => {
195203
vi.resetModules();
196204
vi.useRealTimers();
205+
MockWebSocket.instances = [];
197206
vi.stubGlobal('WebSocket', MockWebSocket);
198207
});
199208

@@ -649,6 +658,75 @@ describe('background tab isolation', () => {
649658
});
650659
});
651660

661+
it('keeps the active daemon connection when a superseded WebSocket closes later', async () => {
662+
const { chrome } = createChromeMock();
663+
vi.stubGlobal('chrome', chrome);
664+
vi.stubGlobal('fetch', vi.fn(async () => ({ ok: true })));
665+
666+
await import('./background');
667+
await vi.waitFor(() => {
668+
expect(MockWebSocket.instances).toHaveLength(1);
669+
});
670+
const firstWs = MockWebSocket.instances[0];
671+
firstWs.readyState = 3;
672+
673+
const onAlarmListener = chrome.alarms.onAlarm.addListener.mock.calls[0][0];
674+
await onAlarmListener({ name: 'keepalive' });
675+
await vi.waitFor(() => {
676+
expect(MockWebSocket.instances).toHaveLength(2);
677+
});
678+
const secondWs = MockWebSocket.instances[1];
679+
secondWs.readyState = MockWebSocket.OPEN;
680+
681+
firstWs.onclose?.();
682+
secondWs.onmessage?.({
683+
data: JSON.stringify({
684+
id: 'sessions-after-stale-close',
685+
action: 'tabs',
686+
op: 'list',
687+
session: 'work',
688+
surface: 'browser',
689+
}),
690+
});
691+
692+
await vi.waitFor(() => {
693+
expect(secondWs.sent.some((entry) => entry.includes('sessions-after-stale-close'))).toBe(true);
694+
});
695+
});
696+
697+
it('ignores daemon commands delivered to a superseded WebSocket', async () => {
698+
const { chrome } = createChromeMock();
699+
vi.stubGlobal('chrome', chrome);
700+
vi.stubGlobal('fetch', vi.fn(async () => ({ ok: true })));
701+
702+
await import('./background');
703+
await vi.waitFor(() => {
704+
expect(MockWebSocket.instances).toHaveLength(1);
705+
});
706+
const firstWs = MockWebSocket.instances[0];
707+
firstWs.readyState = MockWebSocket.OPEN;
708+
709+
const onAlarmListener = chrome.alarms.onAlarm.addListener.mock.calls[0][0];
710+
firstWs.readyState = MockWebSocket.CLOSED;
711+
await onAlarmListener({ name: 'keepalive' });
712+
await vi.waitFor(() => {
713+
expect(MockWebSocket.instances).toHaveLength(2);
714+
});
715+
firstWs.readyState = MockWebSocket.OPEN;
716+
717+
await firstWs.onmessage?.({
718+
data: JSON.stringify({
719+
id: 'stale-command',
720+
action: 'tabs',
721+
op: 'list',
722+
session: 'work',
723+
surface: 'browser',
724+
}),
725+
});
726+
727+
expect(firstWs.sent.some((entry) => entry.includes('stale-command'))).toBe(false);
728+
});
729+
652730
it('can execute concurrently on two pages in the same session', async () => {
653731
const { chrome, tabs } = createChromeMock();
654732
tabs.push({

‎extension/src/background.ts‎

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -70,13 +70,22 @@ const _origWarn = console.warn.bind(console);
7070
const _origError = console.error.bind(console);
7171

7272
function forwardLog(level: 'info' | 'warn' | 'error', args: unknown[]): void {
73-
if (!ws || ws.readyState !== WebSocket.OPEN) return;
7473
try {
7574
const msg = args.map(a => typeof a === 'string' ? a : JSON.stringify(a)).join(' ');
76-
ws.send(JSON.stringify({ type: 'log', level, msg, ts: Date.now() }));
75+
safeSend(ws, { type: 'log', level, msg, ts: Date.now() });
7776
} catch { /* don't recurse */ }
7877
}
7978

79+
function safeSend(socket: WebSocket | null | undefined, payload: unknown): boolean {
80+
if (!socket || socket.readyState !== WebSocket.OPEN) return false;
81+
try {
82+
socket.send(JSON.stringify(payload));
83+
return true;
84+
} catch {
85+
return false;
86+
}
87+
}
88+
8089
console.log = (...args: unknown[]) => { _origLog(...args); forwardLog('info', args); };
8190
console.warn = (...args: unknown[]) => { _origWarn(...args); forwardLog('warn', args); };
8291
console.error = (...args: unknown[]) => { _origError(...args); forwardLog('error', args); };
@@ -100,49 +109,55 @@ async function connect(): Promise<void> {
100109
return; // daemon not running — skip WebSocket to avoid console noise
101110
}
102111

112+
let thisWs: WebSocket;
103113
try {
104114
const contextId = await getCurrentContextId();
105-
ws = new WebSocket(DAEMON_WS_URL);
115+
thisWs = new WebSocket(DAEMON_WS_URL);
116+
ws = thisWs;
106117
currentContextId = contextId;
107118
} catch {
108119
scheduleReconnect();
109120
return;
110121
}
111122

112-
ws.onopen = () => {
123+
thisWs.onopen = () => {
124+
if (ws !== thisWs) return;
113125
console.log('[opencli] Connected to daemon');
114126
reconnectAttempts = 0; // Reset on successful connection
115127
if (reconnectTimer) {
116128
clearTimeout(reconnectTimer);
117129
reconnectTimer = null;
118130
}
119131
// Send version + compatibility range so the daemon can report mismatches to the CLI
120-
ws?.send(JSON.stringify({
132+
safeSend(thisWs, {
121133
type: 'hello',
122134
contextId: currentContextId,
123135
version: chrome.runtime.getManifest().version,
124136
compatRange: __OPENCLI_COMPAT_RANGE__,
125-
}));
137+
});
126138
};
127139

128-
ws.onmessage = async (event) => {
140+
thisWs.onmessage = async (event) => {
141+
if (ws !== thisWs) return;
129142
try {
130143
const command = JSON.parse(event.data as string) as Command;
131144
const result = await handleCommand(command);
132-
ws?.send(JSON.stringify(result));
145+
if (ws !== thisWs) return;
146+
safeSend(thisWs, result);
133147
} catch (err) {
134148
console.error('[opencli] Message handling error:', err);
135149
}
136150
};
137151

138-
ws.onclose = () => {
152+
thisWs.onclose = () => {
153+
if (ws !== thisWs) return;
139154
console.log('[opencli] Disconnected from daemon');
140155
ws = null;
141156
scheduleReconnect();
142157
};
143158

144-
ws.onerror = () => {
145-
ws?.close();
159+
thisWs.onerror = () => {
160+
thisWs.close();
146161
};
147162
}
148163

0 commit comments

Comments
 (0)