From 503c2c743d906a2436b11cfeb2f2b036abe96bd1 Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Sun, 19 Jul 2026 01:49:06 +0200 Subject: [PATCH] fix(sessions): harden notification inbox reconnects --- .../keep-session-notifications-visible.md | 5 + src/client/src/components/ChatView.ts | 4 +- .../sessionController.notifications.test.ts | 36 ++++++ .../sessionController.testSupport.ts | 15 ++- .../src/controllers/sessionController.ts | 11 +- .../sessionNotificationController.test.ts | 122 +++++++++++++++++- .../sessionNotificationController.ts | 90 +++++++++++-- src/client/src/sessionSocket.test.ts | 120 +++++++++++++++-- src/client/src/sessionSocket.ts | 42 ++++-- 9 files changed, 402 insertions(+), 43 deletions(-) create mode 100644 .changeset/keep-session-notifications-visible.md diff --git a/.changeset/keep-session-notifications-visible.md b/.changeset/keep-session-notifications-visible.md new file mode 100644 index 0000000..3080fe5 --- /dev/null +++ b/.changeset/keep-session-notifications-visible.md @@ -0,0 +1,5 @@ +--- +"@jmfederico/pi-web": patch +--- + +Keep extension notifications discoverable in a session-scoped inbox with background badges, reconnect recovery, and explicit dismissal. diff --git a/src/client/src/components/ChatView.ts b/src/client/src/components/ChatView.ts index 42d225b..3dc11a0 100644 --- a/src/client/src/components/ChatView.ts +++ b/src/client/src/components/ChatView.ts @@ -399,8 +399,8 @@ export class ChatView extends LitElement { const polite = announcements.filter((announcement) => announcement.severity !== "error"); const assertive = announcements.filter((announcement) => announcement.severity === "error"); return html` - - + + `; } diff --git a/src/client/src/controllers/sessionController.notifications.test.ts b/src/client/src/controllers/sessionController.notifications.test.ts index 4074507..9ec5c8e 100644 --- a/src/client/src/controllers/sessionController.notifications.test.ts +++ b/src/client/src/controllers/sessionController.notifications.test.ts @@ -33,6 +33,42 @@ function inboxEvent(): SessionNotificationInboxEvent { } describe("SessionController notification event boundary", () => { + it("refetches the bounded notification snapshot when the selected socket first opens", async () => { + const socket = new EmitSocket(); + let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, selectedSession: oldSession, sessions: [oldSession] }; + const refreshSelectedSession = vi.fn(() => Promise.resolve()); + const bridge: SessionNotificationSessionBridge = { + prepareSelectedSession: vi.fn(), + clearSelectedSession: vi.fn(), + refreshSelectedSession, + applyInboxEvent: vi.fn(), + shouldFilterLegacyNotification: vi.fn(() => true), + }; + const controller = new SessionController( + () => state, + (patch) => { state = { ...state, ...patch }; }, + () => undefined, + undefined, + { + socket, + notifications: bridge, + api: { + ...defaultApi, + messages: vi.fn(() => Promise.resolve(emptyPage)), + status: vi.fn(() => Promise.resolve(status(oldSession.id))), + streamSnapshot: vi.fn(() => Promise.resolve({ seq: 0, partial: null })), + }, + }, + ); + + await controller.selectSession(oldSession, { updateUrl: false }); + expect(refreshSelectedSession).toHaveBeenCalledOnce(); + + socket.open(); + expect(refreshSelectedSession).toHaveBeenCalledTimes(2); + expect(refreshSelectedSession).toHaveBeenLastCalledWith(oldSession, "local"); + }); + it("handles inbox events before transcript watermarking and filters only marked legacy output with support", async () => { const socket = new EmitSocket(); let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, selectedSession: oldSession, sessions: [oldSession] }; diff --git a/src/client/src/controllers/sessionController.testSupport.ts b/src/client/src/controllers/sessionController.testSupport.ts index 569eb08..dbcf583 100644 --- a/src/client/src/controllers/sessionController.testSupport.ts +++ b/src/client/src/controllers/sessionController.testSupport.ts @@ -55,10 +55,18 @@ export class FakeSocket implements SessionEventSocket { export class EmitSocket implements SessionEventSocket { readonly connectedSessionIds: string[] = []; private handler: ((event: SessionUiEvent) => void) | undefined; + private onInitialOpen: (() => void) | undefined; - connect(session: SessionRef, onEvent: (event: SessionUiEvent) => void): void { + connect( + session: SessionRef, + onEvent: (event: SessionUiEvent) => void, + _onReconnect?: () => void, + _machineId?: string, + onInitialOpen?: () => void, + ): void { this.connectedSessionIds.push(session.id); this.handler = onEvent; + this.onInitialOpen = onInitialOpen; } setHandler(onEvent: (event: SessionUiEvent) => void): void { @@ -69,8 +77,13 @@ export class EmitSocket implements SessionEventSocket { this.handler?.(event); } + open(): void { + this.onInitialOpen?.(); + } + close(): void { this.handler = undefined; + this.onInitialOpen = undefined; } } diff --git a/src/client/src/controllers/sessionController.ts b/src/client/src/controllers/sessionController.ts index b25f834..3063768 100644 --- a/src/client/src/controllers/sessionController.ts +++ b/src/client/src/controllers/sessionController.ts @@ -20,7 +20,13 @@ const MESSAGE_PAGE_SIZE = 100; const BULK_FALLBACK_CONCURRENCY = 4; export interface SessionEventSocket { - connect(session: SessionRef, onEvent: (event: SessionUiEvent) => void, onReconnect?: () => void, machineId?: string): void; + connect( + session: SessionRef, + onEvent: (event: SessionUiEvent) => void, + onReconnect?: () => void, + machineId?: string, + onInitialOpen?: () => void, + ): void; setHandler(onEvent: (event: SessionUiEvent) => void): void; close(): void; } @@ -206,7 +212,8 @@ export class SessionController { session, (event) => buffered.push(event), () => { void this.refreshSelectedSession(session.id); }, - selectedMachineId(this.getState()), + machineId, + () => { void this.notifications?.refreshSelectedSession(session, machineId); }, ); await this.requestSelectedSessionRefresh({ session, machineId, selectionSeq: seq }); if (!this.isCurrentRefreshTarget({ session, machineId, selectionSeq: seq })) return; diff --git a/src/client/src/controllers/sessionNotificationController.test.ts b/src/client/src/controllers/sessionNotificationController.test.ts index d8ad614..e6ffda0 100644 --- a/src/client/src/controllers/sessionNotificationController.test.ts +++ b/src/client/src/controllers/sessionNotificationController.test.ts @@ -123,7 +123,12 @@ function createHarness(initialState = capableState(), overrides: Partial { state = { ...state, ...patch }; }, { api, onBackgroundError: vi.fn() }, ); - return { controller, api, get state() { return state; } }; + return { + controller, + api, + get state() { return state; }, + replaceState(next: AppState) { state = next; }, + }; } describe("SessionNotificationController capability and joins", () => { @@ -184,6 +189,66 @@ describe("SessionNotificationController capability and joins", () => { expect(harness.controller.shouldFilterLegacyNotification("local", "notification-1")).toBe(true); }); + it("lets the selected live event announce before its matching global summary can trigger a snapshot", async () => { + const first = inboxSnapshot([entry(1)], { inboxRevision: 1, catalogRevision: 1 }); + const state = { + ...capableState(), + notificationCatalogsByMachine: { + local: { + machineId: "local", + status: "fresh" as const, + daemonInstanceId: first.daemonInstanceId, + catalogRevision: first.catalogRevision, + summariesBySessionId: { [session.id]: first.summary }, + }, + }, + }; + const notificationInbox = vi.fn(() => Promise.resolve(first)); + const harness = createHarness(state, { notificationInbox }); + harness.controller.prepareSelectedSession(session, "local"); + await harness.controller.refreshSelectedSession(session, "local"); + + const event = addedEvent(entry(2, "warning"), 2, 2); + harness.controller.applySummaryEvent("local", { + type: "notifications.summary", + daemonInstanceId: event.daemonInstanceId, + catalogRevision: event.catalogRevision, + summary: event.summary, + }); + await Promise.resolve(); + expect(notificationInbox).toHaveBeenCalledOnce(); + + harness.controller.applyInboxEvent("local", event); + expect(selectedNotificationView(harness.state.selectedNotificationInbox)?.announcements).toMatchObject([ + { severity: "warning", message: "notice 2" }, + ]); + expect(notificationInbox).toHaveBeenCalledOnce(); + }); + + it("refetches the selected inbox when a newly opened global socket finds a newer catalog revision", async () => { + const first = inboxSnapshot([entry(1)], { inboxRevision: 1, catalogRevision: 1 }); + const second = inboxSnapshot([entry(2, "error"), entry(1)], { inboxRevision: 2, catalogRevision: 2 }); + const notificationInbox = vi.fn() + .mockResolvedValueOnce(first) + .mockResolvedValueOnce(second); + const notificationCatalog = vi.fn(() => Promise.resolve({ + daemonInstanceId: second.daemonInstanceId, + catalogRevision: second.catalogRevision, + sessions: [second.summary], + })); + const harness = createHarness(capableState(), { notificationInbox, notificationCatalog }); + + harness.controller.prepareSelectedSession(session, "local"); + await harness.controller.refreshSelectedSession(session, "local"); + harness.controller.globalSocketOpened("local"); + + await vi.waitFor(() => { expect(notificationInbox).toHaveBeenCalledTimes(2); }); + expect(selectedNotificationView(harness.state.selectedNotificationInbox)?.notifications.map((notification) => notification.id)).toEqual([ + "daemon-a:2", + "daemon-a:1", + ]); + }); + it("hydrates missing selected-machine project workspaces without changing selection", async () => { const project = { id: "project-1", name: "Repo", path: "/repo", createdAt: "now" }; const next = { @@ -224,19 +289,70 @@ describe("SessionNotificationController capability and joins", () => { expect(harness.state.selectedNotificationInbox).toMatchObject({ sessionId: "session-2", status: "loading", notifications: [] }); }); + + it("hides a selected remote inbox when the machine becomes unreachable and ignores an in-flight snapshot", async () => { + const remoteMachine: Machine = { ...localMachine, id: "remote-a", name: "Remote", kind: "remote", baseUrl: "https://remote.example.test/" }; + const runtime = capableState().machineRuntimes["local"]; + if (runtime === undefined) throw new Error("expected capable runtime fixture"); + const initial = { + ...capableState(), + machines: [remoteMachine], + selectedMachine: remoteMachine, + machineRuntimes: { [remoteMachine.id]: { ...runtime, machineId: remoteMachine.id } }, + }; + const pendingInbox = deferred(); + const notificationInbox = vi.fn() + .mockResolvedValueOnce(inboxSnapshot()) + .mockImplementationOnce(() => pendingInbox.promise); + const harness = createHarness(initial, { notificationInbox }); + + harness.controller.prepareSelectedSession(session, remoteMachine.id); + await harness.controller.refreshSelectedSession(session, remoteMachine.id); + expect(selectedNotificationView(harness.state.selectedNotificationInbox)?.notifications).toHaveLength(1); + + const refresh = harness.controller.refreshSelectedSession(session, remoteMachine.id); + const previous = harness.state; + const offline: AppState = { + ...previous, + machineStatuses: { + [remoteMachine.id]: { + machineId: remoteMachine.id, + ok: false, + checkedAt: "2026-07-18T00:01:00.000Z", + status: "offline", + }, + }, + }; + harness.replaceState(offline); + harness.controller.syncEnvironment(previous, offline); + + expect(harness.state.selectedNotificationInbox?.status).toBe("stale"); + expect(selectedNotificationView(harness.state.selectedNotificationInbox)).toBeUndefined(); + + pendingInbox.resolve(inboxSnapshot([entry(2, "error")], { inboxRevision: 2, catalogRevision: 2 })); + await refresh; + expect(harness.state.selectedNotificationInbox?.status).toBe("stale"); + expect(selectedNotificationView(harness.state.selectedNotificationInbox)).toBeUndefined(); + }); }); describe("SessionNotificationController optimistic mutations", () => { it("optimistically dismisses one card, reconciles the response, and rolls back/refetches on failure", async () => { const dismiss = deferred(); const refreshAfterFailure = deferred(); + const initialInbox = inboxSnapshot([entry(2, "warning"), entry(1)]); const notificationInbox = vi.fn() - .mockResolvedValueOnce(inboxSnapshot([entry(2, "warning"), entry(1)])) + .mockResolvedValueOnce(initialInbox) .mockImplementationOnce(() => refreshAfterFailure.promise); const dismissNotification = vi.fn() .mockImplementationOnce(() => dismiss.promise) .mockRejectedValueOnce(new Error("offline")); - const harness = createHarness(capableState(), { notificationInbox, dismissNotification }); + const notificationCatalog = vi.fn(() => Promise.resolve({ + daemonInstanceId: initialInbox.daemonInstanceId, + catalogRevision: initialInbox.catalogRevision, + sessions: [initialInbox.summary], + })); + const harness = createHarness(capableState(), { notificationInbox, dismissNotification, notificationCatalog }); harness.controller.prepareSelectedSession(session, "local"); await harness.controller.refreshSelectedSession(session, "local"); diff --git a/src/client/src/controllers/sessionNotificationController.ts b/src/client/src/controllers/sessionNotificationController.ts index 785a718..80f633e 100644 --- a/src/client/src/controllers/sessionNotificationController.ts +++ b/src/client/src/controllers/sessionNotificationController.ts @@ -6,6 +6,7 @@ import { freshNotificationCatalog, installSelectedNotificationSnapshot, loadingSelectedNotificationInbox, + notificationSummaryIsEmpty, notificationTargetsEqual, selectedNotificationView, type SelectedSessionNotificationInbox, @@ -143,9 +144,11 @@ export class SessionNotificationController { applyInboxEvent(machineId: string, event: SessionNotificationInboxEvent): void { this.acceptedSupportByMachine.add(machineId); - this.applyCatalogSummary(machineId, inboxSummaryEvent(event)); const target = this.selectedTarget; - if (target?.machineId !== machineId || target.sessionId !== event.summary.sessionId || target.cwd !== event.summary.cwd) return; + if (target?.machineId !== machineId || target.sessionId !== event.summary.sessionId || target.cwd !== event.summary.cwd) { + this.applyCatalogSummary(machineId, inboxSummaryEvent(event)); + return; + } const join = this.selectedJoin; if (join?.generation === this.selectedGeneration) { join.events.push(event); @@ -153,6 +156,7 @@ export class SessionNotificationController { } const result = applySelectedNotificationEvent(this.getState().selectedNotificationInbox, target, event); if (result.changed) this.setState({ selectedNotificationInbox: result.value }); + this.applyCatalogSummary(machineId, inboxSummaryEvent(event)); if (result.needsRefresh) this.scheduleSelectedRefresh(target); } @@ -163,7 +167,10 @@ export class SessionNotificationController { join.events.push(event); return; } - this.applyCatalogSummary(machineId, event); + // The matching per-session event carries the notification text and must win + // the live-announcement race. Initial-open and reconnect snapshots still + // reconcile selected state through the catalog join. + this.applyCatalogSummary(machineId, event, false); this.ensureSelectedSupport(machineId); } @@ -193,10 +200,14 @@ export class SessionNotificationController { else if (!wasEligible || next.notificationCatalogsByMachine[machineId]?.status !== "fresh") void this.refreshCatalog(machineId); } const selected = this.selectedTarget; - if (selected !== undefined && this.machineSupportsNotifications(selected.machineId) && this.machineIsReachable(selected.machineId)) { - this.ensureSelectedProjection(selected); - if (!this.machineSupportsNotificationsInState(previous, selected.machineId) || next.selectedNotificationInbox?.status !== "fresh") { - void this.refreshSelectedSession({ id: selected.sessionId, cwd: selected.cwd }, selected.machineId); + if (selected !== undefined) { + if (!this.machineSupportsNotifications(selected.machineId) || !this.machineIsReachable(selected.machineId)) { + this.markSelectedStale(selected); + } else { + this.ensureSelectedProjection(selected); + if (!this.machineSupportsNotificationsInState(previous, selected.machineId) || next.selectedNotificationInbox?.status !== "fresh") { + void this.refreshSelectedSession({ id: selected.sessionId, cwd: selected.cwd }, selected.machineId); + } } } } @@ -285,16 +296,21 @@ export class SessionNotificationController { try { const snapshot = await this.api.notificationInbox({ id: target.sessionId, cwd: target.cwd }, target.machineId); if (!this.isCurrentTarget(target, operation.generation)) return; + if (!this.machineSupportsNotifications(target.machineId) || !this.machineIsReachable(target.machineId)) { + this.markSelectedStale(target); + return; + } this.acceptedSupportByMachine.add(target.machineId); let inbox = installSelectedNotificationSnapshot(this.getState().selectedNotificationInbox, target, snapshot); - this.applyCatalogSummary(target.machineId, snapshotSummaryEvent(snapshot)); + const catalogEvents = [snapshotSummaryEvent(snapshot)]; for (const event of [...join.events].sort((left, right) => left.summary.inboxRevision - right.summary.inboxRevision)) { const result = applySelectedNotificationEvent(inbox, target, event); inbox = result.value; - this.applyCatalogSummary(target.machineId, inboxSummaryEvent(event)); + catalogEvents.push(inboxSummaryEvent(event)); if (result.needsRefresh) operation.trailing = true; } this.setState({ selectedNotificationInbox: inbox }); + for (const event of catalogEvents) this.applyCatalogSummary(target.machineId, event); } catch (error) { if (this.isCurrentTarget(target, operation.generation)) { const current = this.getState().selectedNotificationInbox; @@ -350,7 +366,7 @@ export class SessionNotificationController { } while (operation.trailing && this.machineIsKnown(machineId) && this.machineIsReachable(machineId)); } - private applyCatalogSummary(machineId: string, event: SessionNotificationSummaryEvent): void { + private applyCatalogSummary(machineId: string, event: SessionNotificationSummaryEvent, reconcileSelected = true): void { const join = this.catalogJoins.get(machineId); if (join !== undefined) { join.events.push(event); @@ -358,7 +374,7 @@ export class SessionNotificationController { } const current = this.getState().notificationCatalogsByMachine[machineId]; const result = applyNotificationCatalogEvent(current, machineId, event); - if (result.changed) this.setCatalog(machineId, result.value); + if (result.changed) this.setCatalog(machineId, result.value, reconcileSelected); if (result.needsRefresh) this.scheduleCatalogRefresh(machineId); } @@ -369,6 +385,10 @@ export class SessionNotificationController { ): void { const current = this.getState().selectedNotificationInbox; if (current === undefined || !notificationTargetsEqual(current, target)) return; + if (!this.machineSupportsNotifications(target.machineId) || !this.machineIsReachable(target.machineId)) { + this.setState({ selectedNotificationInbox: { ...removeOverlay(current), status: "stale" } }); + return; + } const shouldInstall = current.daemonInstanceId !== snapshot.daemonInstanceId || current.summary === undefined || snapshot.summary.inboxRevision >= current.summary.inboxRevision; @@ -389,6 +409,8 @@ export class SessionNotificationController { private ensureSelectedSupport(machineId: string): void { const target = this.selectedTarget; if (target?.machineId !== machineId) return; + const current = this.getState().selectedNotificationInbox; + if (current !== undefined && notificationTargetsEqual(current, target) && current.status === "fresh") return; this.ensureSelectedProjection(target); void this.refreshSelectedSession({ id: target.sessionId, cwd: target.cwd }, machineId); } @@ -412,10 +434,36 @@ export class SessionNotificationController { }); } - private setCatalog(machineId: string, projection: SessionNotificationCatalogProjection): void { + private setCatalog(machineId: string, projection: SessionNotificationCatalogProjection, reconcileSelected = true): void { const current = this.getState().notificationCatalogsByMachine; if (current[machineId] === projection) return; this.setState({ notificationCatalogsByMachine: { ...current, [machineId]: projection } }); + if (reconcileSelected) this.reconcileSelectedWithCatalog(projection); + } + + private reconcileSelectedWithCatalog(catalog: SessionNotificationCatalogProjection): void { + if (catalog.status !== "fresh" || catalog.daemonInstanceId === undefined) return; + const target = this.selectedTarget; + if (target?.machineId !== catalog.machineId || !this.machineIsReachable(target.machineId)) return; + const inbox = this.getState().selectedNotificationInbox; + if (inbox === undefined || !notificationTargetsEqual(inbox, target) || inbox.status !== "fresh" || inbox.daemonInstanceId === undefined || inbox.summary === undefined) { + this.scheduleSelectedRefresh(target); + return; + } + if (inbox.daemonInstanceId !== catalog.daemonInstanceId) { + this.scheduleSelectedRefresh(target); + return; + } + const catalogSummary = catalog.summariesBySessionId[target.sessionId]; + if (catalogSummary === undefined) { + if (!notificationSummaryIsEmpty(inbox.summary)) this.scheduleSelectedRefresh(target); + return; + } + if (catalogSummary.cwd !== target.cwd + || catalogSummary.inboxRevision > inbox.summary.inboxRevision + || (catalogSummary.inboxRevision === inbox.summary.inboxRevision && !notificationSummariesEqual(catalogSummary, inbox.summary))) { + this.scheduleSelectedRefresh(target); + } } private markCatalogStale(machineId: string): void { @@ -424,6 +472,12 @@ export class SessionNotificationController { this.setCatalog(machineId, { ...current, status: "stale" }); } + private markSelectedStale(target: SessionNotificationTarget): void { + const current = this.getState().selectedNotificationInbox; + if (current === undefined || !notificationTargetsEqual(current, target) || current.status === "stale") return; + this.setState({ selectedNotificationInbox: { ...current, status: "stale" } }); + } + private pruneRemovedMachines(machines: readonly Machine[]): void { const machineIds = new Set(machines.map((machine) => machine.id)); if (machineIds.size === 0) machineIds.add("local"); @@ -540,6 +594,18 @@ function workspaceHydrationId(machineId: string, projectId: string): string { return JSON.stringify([machineId, projectId]); } +function notificationSummariesEqual( + left: SessionNotificationInboxSnapshot["summary"], + right: SessionNotificationInboxSnapshot["summary"], +): boolean { + return left.sessionId === right.sessionId + && left.cwd === right.cwd + && left.inboxRevision === right.inboxRevision + && left.retainedCount === right.retainedCount + && left.discardedCount === right.discardedCount + && left.highestSeverity === right.highestSeverity; +} + function errorMessage(error: unknown): string { return error instanceof Error ? error.message : String(error); } diff --git a/src/client/src/sessionSocket.test.ts b/src/client/src/sessionSocket.test.ts index 777d809..2488eba 100644 --- a/src/client/src/sessionSocket.test.ts +++ b/src/client/src/sessionSocket.test.ts @@ -1,5 +1,5 @@ -import { describe, expect, it } from "vitest"; -import { parseRealtimeSocketEvent, parseSessionSocketEvent } from "./sessionSocket"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { RealtimeSocket, SessionSocket, parseRealtimeSocketEvent, parseSessionSocketEvent } from "./sessionSocket"; function notification(order = 1) { return { @@ -23,16 +23,20 @@ function summary() { }; } +function inboxEvent() { + return { + type: "notifications.inbox", + daemonInstanceId: "daemon-a", + catalogRevision: 1, + summary: summary(), + dismissThrough: { order: 1, overflowWatermark: 0 }, + delta: { kind: "added", notification: notification() }, + }; +} + describe("notification socket guards", () => { it("accepts validated per-session and global notification events", () => { - expect(parseSessionSocketEvent({ - type: "notifications.inbox", - daemonInstanceId: "daemon-a", - catalogRevision: 1, - summary: summary(), - dismissThrough: { order: 1, overflowWatermark: 0 }, - delta: { kind: "added", notification: notification() }, - })).toMatchObject({ type: "notifications.inbox", delta: { kind: "added" } }); + expect(parseSessionSocketEvent(inboxEvent())).toMatchObject({ type: "notifications.inbox", delta: { kind: "added" } }); expect(parseRealtimeSocketEvent({ type: "notifications.summary", @@ -64,3 +68,99 @@ describe("notification socket guards", () => { expect(parseRealtimeSocketEvent({ type: "future.notification", payload: {} })).toBeUndefined(); }); }); + +class FakeWebSocket { + static readonly CONNECTING = 0; + static readonly instances: FakeWebSocket[] = []; + + readyState = 1; + onopen: (() => void) | null = null; + onmessage: ((event: { data: MessageEvent["data"] }) => void) | null = null; + onerror: (() => void) | null = null; + onclose: (() => void) | null = null; + + constructor(readonly url: string) { + FakeWebSocket.instances.push(this); + } + + close(): void { + this.readyState = 3; + } +} + +describe("socket instance isolation", () => { + const setTimeoutSpy = vi.fn(() => 1); + + beforeEach(() => { + FakeWebSocket.instances.length = 0; + setTimeoutSpy.mockClear(); + vi.stubGlobal("WebSocket", FakeWebSocket); + vi.stubGlobal("document", { baseURI: "https://pi.example.test/" }); + vi.stubGlobal("window", { clearTimeout: vi.fn(), setTimeout: setTimeoutSpy }); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("drops queued session frames and close callbacks from a replaced machine socket", async () => { + const socket = new SessionSocket(); + const oldHandler = vi.fn(); + const newHandler = vi.fn(); + const onInitialOpen = vi.fn(); + const target = { id: "session-1", cwd: "/repo" }; + socket.connect(target, oldHandler, undefined, "machine-a"); + const oldSocket = FakeWebSocket.instances[0]; + if (oldSocket === undefined) throw new Error("expected old session socket"); + const staleClose = oldSocket.onclose; + oldSocket.onmessage?.({ data: JSON.stringify(inboxEvent()) }); + + socket.connect(target, newHandler, undefined, "machine-b", onInitialOpen); + staleClose?.(); + await Promise.resolve(); + await Promise.resolve(); + + expect(oldHandler).not.toHaveBeenCalled(); + expect(newHandler).not.toHaveBeenCalled(); + expect(setTimeoutSpy).not.toHaveBeenCalled(); + + const newSocket = FakeWebSocket.instances[1]; + if (newSocket === undefined) throw new Error("expected replacement session socket"); + newSocket.onopen?.(); + expect(onInitialOpen).toHaveBeenCalledOnce(); + newSocket.onmessage?.({ data: JSON.stringify(inboxEvent()) }); + await Promise.resolve(); + await Promise.resolve(); + expect(newHandler).toHaveBeenCalledOnce(); + }); + + it("does not attribute a queued global frame to a replacement machine", async () => { + const socket = new RealtimeSocket(); + const oldHandler = vi.fn(); + const newHandler = vi.fn(); + const event = { + type: "notifications.summary", + daemonInstanceId: "daemon-a", + catalogRevision: 1, + summary: summary(), + }; + socket.connect(oldHandler, undefined, "machine-a"); + const oldSocket = FakeWebSocket.instances[0]; + if (oldSocket === undefined) throw new Error("expected old realtime socket"); + oldSocket.onmessage?.({ data: JSON.stringify(event) }); + + socket.connect(newHandler, undefined, "machine-b"); + await Promise.resolve(); + await Promise.resolve(); + + expect(oldHandler).not.toHaveBeenCalled(); + expect(newHandler).not.toHaveBeenCalled(); + + const newSocket = FakeWebSocket.instances[1]; + if (newSocket === undefined) throw new Error("expected replacement realtime socket"); + newSocket.onmessage?.({ data: JSON.stringify(event) }); + await Promise.resolve(); + await Promise.resolve(); + expect(newHandler).toHaveBeenCalledOnce(); + }); +}); diff --git a/src/client/src/sessionSocket.ts b/src/client/src/sessionSocket.ts index 95d9a1c..66b6015 100644 --- a/src/client/src/sessionSocket.ts +++ b/src/client/src/sessionSocket.ts @@ -13,14 +13,22 @@ export class SessionSocket { private shouldReconnect = false; private hasOpened = false; private onReconnect: (() => void) | undefined; + private onInitialOpen: (() => void) | undefined; private machineId = "local"; - connect(session: SessionRef, onEvent: (event: SessionUiEvent) => void, onReconnect?: () => void, machineId = "local"): void { + connect( + session: SessionRef, + onEvent: (event: SessionUiEvent) => void, + onReconnect?: () => void, + machineId = "local", + onInitialOpen?: () => void, + ): void { this.close(); this.machineId = machineId; this.session = session; this.onEvent = onEvent; this.onReconnect = onReconnect; + this.onInitialOpen = onInitialOpen; this.shouldReconnect = true; this.open(); } @@ -37,23 +45,29 @@ export class SessionSocket { this.session = undefined; this.onEvent = undefined; this.onReconnect = undefined; + this.onInitialOpen = undefined; this.hasOpened = false; this.machineId = "local"; } private open(): void { - if (this.session === undefined || this.session.id === "" || this.session.cwd === "" || !this.shouldReconnect) return; - const socket = sessionEvents(this.session, this.machineId); + const session = this.session; + if (session === undefined || session.id === "" || session.cwd === "" || !this.shouldReconnect) return; + const socket = sessionEvents(session, this.machineId); this.socket = socket; socket.onopen = () => { + if (this.socket !== socket) return; this.reconnectDelay = 500; - if (this.hasOpened) this.onReconnect?.(); + const isReconnect = this.hasOpened; this.hasOpened = true; + if (isReconnect) this.onReconnect?.(); + else this.onInitialOpen?.(); }; - socket.onmessage = (message) => void this.handleMessage(message.data, this.session); + socket.onmessage = (message) => void this.handleMessage(message.data, socket, session); socket.onerror = () => { socket.close(); }; socket.onclose = () => { - if (this.socket === socket) this.socket = undefined; + if (this.socket !== socket) return; + this.socket = undefined; this.scheduleReconnect(); }; } @@ -66,10 +80,10 @@ export class SessionSocket { this.reconnectTimer = window.setTimeout(() => { this.open(); }, delay); } - private async handleMessage(data: MessageEvent["data"], session: SessionRef | undefined): Promise { + private async handleMessage(data: MessageEvent["data"], socket: WebSocket, session: SessionRef): Promise { const event = parseSessionSocketEvent(await parseSocketEvent(data)); - if (event === undefined) return; - if (event.type === "notifications.inbox" && (session?.id !== event.summary.sessionId || session.cwd !== event.summary.cwd)) return; + if (this.socket !== socket || event === undefined) return; + if (event.type === "notifications.inbox" && (session.id !== event.summary.sessionId || session.cwd !== event.summary.cwd)) return; this.onEvent?.(event); } } @@ -107,13 +121,15 @@ export class RealtimeSocket { const socket = realtimeEvents(this.machineId); this.socket = socket; socket.onopen = () => { + if (this.socket !== socket) return; this.reconnectDelay = 500; this.onOpen?.(); }; - socket.onmessage = (message) => void this.handleMessage(message.data); + socket.onmessage = (message) => void this.handleMessage(message.data, socket); socket.onerror = () => { socket.close(); }; socket.onclose = () => { - if (this.socket === socket) this.socket = undefined; + if (this.socket !== socket) return; + this.socket = undefined; this.scheduleReconnect(); }; } @@ -126,9 +142,9 @@ export class RealtimeSocket { this.reconnectTimer = window.setTimeout(() => { this.open(); }, delay); } - private async handleMessage(data: MessageEvent["data"]): Promise { + private async handleMessage(data: MessageEvent["data"], socket: WebSocket): Promise { const event = parseRealtimeSocketEvent(await parseSocketEvent(data)); - if (event !== undefined) this.onEvent?.(event); + if (this.socket === socket && event !== undefined) this.onEvent?.(event); } }