diff --git a/packages/code-link-cli/src/helpers/connection.ts b/packages/code-link-cli/src/helpers/connection.ts index fe63730d9..907f484fb 100644 --- a/packages/code-link-cli/src/helpers/connection.ts +++ b/packages/code-link-cli/src/helpers/connection.ts @@ -4,7 +4,7 @@ * Wrapper around ws.Server that normalizes handshake and surfaces callbacks. */ -import type { CliToPluginMessage, PluginToCliMessage } from "@code-link/shared" +import { CLOSE_CODE_REPLACED, type CliToPluginMessage, type PluginToCliMessage } from "@code-link/shared" import https from "node:https" import { WebSocket, WebSocketServer } from "ws" import type { CertBundle } from "./certs.ts" @@ -89,7 +89,7 @@ export function initConnection(port: number, certs: CertBundle): Promise { dispatch({ type: "set-mode", mode: "syncing" }) } + const handleReplaced = () => { + dispatch({ type: "set-mode", mode: "replaced" }) + } const handleMessage = createMessageHandler({ dispatch, api, syncTracker }) const controller = createSocketConnectionController({ @@ -175,6 +178,7 @@ export function App() { onMessage: handleMessage, onConnected: handleConnected, onDisconnected: handleDisconnected, + onReplaced: handleReplaced, }) controller.start() @@ -262,6 +266,10 @@ export function App() { }) return + case "replaced": + return framer.closePlugin("Replaced by another Plugin connection", { + variant: "info", + }) default: void framer.setBackgroundMessage(backgroundStatusFromMode(state.mode)) void framer.hideUI() diff --git a/plugins/code-link/src/utils/sockets.test.ts b/plugins/code-link/src/utils/sockets.test.ts new file mode 100644 index 000000000..0cad1a919 --- /dev/null +++ b/plugins/code-link/src/utils/sockets.test.ts @@ -0,0 +1,134 @@ +import { CLOSE_CODE_REPLACED, type ProjectInfo } from "@code-link/shared" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" +import { createSocketConnectionController } from "./sockets.ts" + +vi.mock("framer-plugin", () => ({ + framer: { + getProjectInfo: vi.fn(() => Promise.resolve({ id: "project-id", name: "Project Name" })), + }, +})) + +class MockEventTarget { + private listeners = new Map>() + + addEventListener = vi.fn((type: string, listener: EventListener) => { + const listenersForType = this.listeners.get(type) ?? new Set() + listenersForType.add(listener) + this.listeners.set(type, listenersForType) + }) + + removeEventListener = vi.fn((type: string, listener: EventListener) => { + this.listeners.get(type)?.delete(listener) + }) + + dispatch(type: string) { + for (const listener of this.listeners.get(type) ?? []) { + listener(new Event(type)) + } + } +} + +class MockWebSocket { + static readonly CONNECTING = 0 + static readonly OPEN = 1 + static readonly CLOSING = 2 + static readonly CLOSED = 3 + static instances: MockWebSocket[] = [] + + readonly url: string + readyState = MockWebSocket.CONNECTING + onopen: ((event: Event) => void) | null = null + onclose: ((event: CloseEventLike) => void) | null = null + onerror: ((event: Event) => void) | null = null + onmessage: ((event: MessageEvent) => void) | null = null + send = vi.fn() + + constructor(url: string) { + this.url = url + MockWebSocket.instances.push(this) + } + + close(code = 1000, reason = "") { + this.readyState = MockWebSocket.CLOSED + this.onclose?.({ code, reason, wasClean: true }) + } + + emitClose(code = 1000, reason = "") { + this.readyState = MockWebSocket.CLOSED + this.onclose?.({ code, reason, wasClean: true }) + } +} + +interface CloseEventLike { + code: number + reason: string + wasClean: boolean +} + +describe("createSocketConnectionController", () => { + const originalDocument = globalThis.document + const originalWindow = globalThis.window + const originalWebSocket = globalThis.WebSocket + let mockDocument: MockEventTarget & { visibilityState: DocumentVisibilityState } + let mockWindow: MockEventTarget + + beforeEach(() => { + vi.useFakeTimers() + MockWebSocket.instances = [] + + mockDocument = Object.assign(new MockEventTarget(), { + visibilityState: "visible" as DocumentVisibilityState, + }) + Object.defineProperty(mockDocument, "visibilityState", { + value: "visible", + configurable: true, + writable: true, + }) + + mockWindow = new MockEventTarget() + + globalThis.document = mockDocument as unknown as Document + globalThis.window = mockWindow as unknown as Window & typeof globalThis + globalThis.WebSocket = MockWebSocket as unknown as typeof WebSocket + }) + + afterEach(() => { + vi.useRealTimers() + vi.restoreAllMocks() + globalThis.document = originalDocument + globalThis.window = originalWindow + globalThis.WebSocket = originalWebSocket + }) + + it("cleans up listeners and timers after a replaced close followed by stop", () => { + const setSocket = vi.fn() + const onReplaced = vi.fn() + + const controller = createSocketConnectionController({ + project: { id: "project-id", name: "Project Name" } satisfies ProjectInfo, + setSocket, + onMessage: vi.fn(() => Promise.resolve()), + onConnected: vi.fn(), + onDisconnected: vi.fn(), + onReplaced, + }) + + controller.start() + + expect(mockDocument.addEventListener).toHaveBeenCalledWith("visibilitychange", expect.any(Function)) + expect(mockWindow.addEventListener).toHaveBeenCalledWith("focus", expect.any(Function)) + expect(MockWebSocket.instances).toHaveLength(1) + + mockWindow.dispatch("focus") + expect(vi.getTimerCount()).toBe(2) + + MockWebSocket.instances[0]?.emitClose(CLOSE_CODE_REPLACED, "replaced") + controller.stop() + + expect(onReplaced).toHaveBeenCalledOnce() + expect(mockDocument.removeEventListener).toHaveBeenCalledWith("visibilitychange", expect.any(Function)) + expect(mockWindow.removeEventListener).toHaveBeenCalledWith("focus", expect.any(Function)) + expect(vi.getTimerCount()).toBe(0) + expect(setSocket).toHaveBeenLastCalledWith(null) + }) +}) diff --git a/plugins/code-link/src/utils/sockets.ts b/plugins/code-link/src/utils/sockets.ts index ad4f7d062..a939421f6 100644 --- a/plugins/code-link/src/utils/sockets.ts +++ b/plugins/code-link/src/utils/sockets.ts @@ -1,4 +1,5 @@ import { + CLOSE_CODE_REPLACED, type CliToPluginMessage, getPortFromHash, isCliToPluginMessage, @@ -30,12 +31,14 @@ export function createSocketConnectionController({ onMessage, onConnected, onDisconnected, + onReplaced, }: { project: ProjectInfo setSocket: (socket: WebSocket | null) => void onMessage: (message: CliToPluginMessage, socket: WebSocket) => Promise onConnected: () => void onDisconnected: (message: string) => void + onReplaced: () => void }): SocketConnectionController { const RECONNECT_BASE_MS = 500 const RECONNECT_MAX_MS = 5000 @@ -52,6 +55,7 @@ export function createSocketConnectionController({ let hasNotifiedDisconnected = false let activeSocket: WebSocket | null = null let messageQueue: Promise = Promise.resolve() + let hasCleanedUp = false const protocol = "wss" const timers: Record | null> = { connectTrigger: null, @@ -117,6 +121,25 @@ export function createSocketConnectionController({ setSocket(socket) } + const cleanupResources = () => { + if (hasCleanedUp) return + hasCleanedUp = true + + document.removeEventListener("visibilitychange", onVisibilityChange) + window.removeEventListener("focus", onFocus) + + clearAllTimers() + + if (activeSocket) { + clearSocket(activeSocket) + } + } + + const dispose = () => { + setLifecycle("disposed") + cleanupResources() + } + const clearSocket = (socket: WebSocket) => { clearTimer("connectTimeout") detachSocketHandlers(socket) @@ -306,7 +329,6 @@ export function createSocketConnectionController({ if (isStale()) return setActiveSocket(null) - failureCount += 1 log.debug("WebSocket closed", { code: event.code, @@ -319,6 +341,16 @@ export function createSocketConnectionController({ failureCount, }) + // Another plugin tab took over this connection — stop reconnecting. + if (event.code === CLOSE_CODE_REPLACED) { + log.debug("Connection replaced by another plugin tab, disposing") + onReplaced() + dispose() + return + } + + failureCount += 1 + if ( !hasNotifiedDisconnected && failureCount >= DISCONNECTED_NOTICE_FAILURE_THRESHOLD && @@ -391,20 +423,7 @@ export function createSocketConnectionController({ } }, stop: () => { - if (isDisposed()) return - setLifecycle("disposed") - - document.removeEventListener("visibilitychange", onVisibilityChange) - window.removeEventListener("focus", onFocus) - - clearAllTimers() - const socket = activeSocket - - if (socket) { - clearSocket(socket) - } else { - setActiveSocket(null) - } + dispose() }, } }