Unified · Split
| @@ -18,6 +18,7 @@ import { · 17 unmodified lines | ||
| 18 | 18 | } from "./wsConnectionState"; |
| 19 | 19 | |
| 20 | 20 | export interface WsProtocolLifecycleHandlers { |
| 21 | + readonly isActive?: () => boolean; | |
| 21 | 22 | readonly onAttempt?: (socketUrl: string) => void; |
| 22 | 23 | readonly onOpen?: () => void; |
| 23 | 24 | readonly onError?: (message: string) => void; |
| @@ -49,6 +50,7 @@ function resolveWsRpcSocketUrl(rawUrl: string): string { · 25 unmodified lines | ||
| 49 | 50 | |
| 50 | 51 | function defaultLifecycleHandlers(): Required<WsProtocolLifecycleHandlers> { |
| 51 | 52 | return { |
| 53 | + isActive: () => true, | |
| 52 | 54 | onAttempt: recordWsConnectionAttempt, |
| 53 | 55 | onOpen: recordWsConnectionOpened, |
| 54 | 56 | onError: (message) => { |
| @@ -66,21 +68,35 @@ function composeLifecycleHandlers( · 11 unmodified lines | ||
| 66 | 68 | handlers?: WsProtocolLifecycleHandlers, |
| 67 | 69 | ): Required<WsProtocolLifecycleHandlers> { |
| 68 | 70 | const defaults = defaultLifecycleHandlers(); |
| 71 | + const isActive = handlers?.isActive ?? (() => true); | |
| 69 | 72 | |
| 70 | 73 | return { |
| 74 | + isActive, | |
| 71 | 75 | onAttempt: (socketUrl) => { |
| 76 | + if (!isActive()) { | |
| 77 | + return; | |
| 78 | + } | |
| 72 | 79 | defaults.onAttempt(socketUrl); |
| 73 | 80 | handlers?.onAttempt?.(socketUrl); |
| 74 | 81 | }, |
| 75 | 82 | onOpen: () => { |
| 83 | + if (!isActive()) { | |
| 84 | + return; | |
| 85 | + } | |
| 76 | 86 | defaults.onOpen(); |
| 77 | 87 | handlers?.onOpen?.(); |
| 78 | 88 | }, |
| 79 | 89 | onError: (message) => { |
| 90 | + if (!isActive()) { | |
| 91 | + return; | |
| 92 | + } | |
| 80 | 93 | defaults.onError(message); |
| 81 | 94 | handlers?.onError?.(message); |
| 82 | 95 | }, |
| 83 | 96 | onClose: (details) => { |
| 97 | + if (!isActive()) { | |
| 98 | + return; | |
| 99 | + } | |
| 84 | 100 | defaults.onClose(details); |
| 85 | 101 | handlers?.onClose?.(details); |
| 86 | 102 | }, |
| @@ -324,6 +324,11 @@ describe("WsTransport", () => { · 323 unmodified lines | ||
| 324 | 324 | const secondSocket = getSocket(); |
| 325 | 325 | expect(secondSocket).not.toBe(firstSocket); |
| 326 | 326 | expect(firstSocket.readyState).toBe(MockWebSocket.CLOSED); |
| 327 | + expect(getWsConnectionStatus()).toMatchObject({ | |
| 328 | + closeCode: null, | |
| 329 | + closeReason: null, | |
| 330 | + phase: "connecting", | |
| 331 | + }); | |
| 327 | 332 | |
| 328 | 333 | const requestPromise = transport.request((client) => |
| 329 | 334 | client[WS_METHODS.serverUpsertKeybinding]({ |
| @@ -361,6 +366,58 @@ describe("WsTransport", () => { · 31 unmodified lines | ||
| 361 | 366 | await transport.dispose(); |
| 362 | 367 | }); |
| 363 | 368 | |
| 369 | + it("ignores stale socket lifecycle events after a reconnect starts a new session", async () => { | |
| 370 | + const onClose = vi.fn(); | |
| 371 | + const transport = createTransport("ws://localhost:3020", { onClose }); | |
| 372 | + | |
| 373 | + await waitFor(() => { | |
| 374 | + expect(sockets).toHaveLength(1); | |
| 375 | + }); | |
| 376 | + | |
| 377 | + const firstSocket = getSocket(); | |
| 378 | + firstSocket.open(); | |
| 379 | + | |
| 380 | + await waitFor(() => { | |
| 381 | + expect(getWsConnectionStatus()).toMatchObject({ | |
| 382 | + hasConnected: true, | |
| 383 | + phase: "connected", | |
| 384 | + }); | |
| 385 | + }); | |
| 386 | + | |
| 387 | + await transport.reconnect(); | |
| 388 | + | |
| 389 | + await waitFor(() => { | |
| 390 | + expect(sockets).toHaveLength(2); | |
| 391 | + }); | |
| 392 | + | |
| 393 | + expect(onClose).not.toHaveBeenCalled(); | |
| 394 | + expect(getWsConnectionStatus()).toMatchObject({ | |
| 395 | + closeCode: null, | |
| 396 | + closeReason: null, | |
| 397 | + phase: "connecting", | |
| 398 | + }); | |
| 399 | + | |
| 400 | + const secondSocket = getSocket(); | |
| 401 | + secondSocket.open(); | |
| 402 | + | |
| 403 | + await waitFor(() => { | |
| 404 | + expect(getWsConnectionStatus()).toMatchObject({ | |
| 405 | + phase: "connected", | |
| 406 | + }); | |
| 407 | + }); | |
| 408 | + | |
| 409 | + firstSocket.close(1006, "stale close"); | |
| 410 | + | |
| 411 | + expect(onClose).not.toHaveBeenCalled(); | |
| 412 | + expect(getWsConnectionStatus()).toMatchObject({ | |
| 413 | + closeCode: null, | |
| 414 | + closeReason: null, | |
| 415 | + phase: "connected", | |
| 416 | + }); | |
| 417 | + | |
| 418 | + await transport.dispose(); | |
| 419 | + }); | |
| 420 | + | |
| 364 | 421 | it("marks unary requests as slow until the first server ack arrives", async () => { |
| 365 | 422 | const slowAckThresholdMs = 25; |
| 366 | 423 | setSlowRpcAckThresholdMsForTests(slowAckThresholdMs); |
| @@ -53,6 +53,8 @@ export class WsTransport { · 52 unmodified lines | ||
| 53 | 53 | private disposed = false; |
| 54 | 54 | private hasReportedTransportDisconnect = false; |
| 55 | 55 | private reconnectChain: Promise<void> = Promise.resolve(); |
| 56 | + private nextSessionId = 0; | |
| 57 | + private activeSessionId = 0; | |
| 56 | 58 | private session: TransportSession; |
| 57 | 59 | |
| 58 | 60 | constructor( |
| @@ -215,8 +217,17 @@ export class WsTransport { · 156 unmodified lines | ||
| 215 | 217 | } |
| 216 | 218 | |
| 217 | 219 | private createSession(): TransportSession { |
| 220 | + const sessionId = this.nextSessionId + 1; | |
| 221 | + this.nextSessionId = sessionId; | |
| 222 | + this.activeSessionId = sessionId; | |
| 218 | 223 | const runtime = ManagedRuntime.make( |
| 219 | − Layer.mergeAll(createWsRpcProtocolLayer(this.url, this.lifecycleHandlers), ClientTracingLive), | |
| 224 | + Layer.mergeAll( | |
| 225 | + createWsRpcProtocolLayer(this.url, { | |
| 226 | + ...this.lifecycleHandlers, | |
| 227 | + isActive: () => !this.disposed && this.activeSessionId === sessionId, | |
| 228 | + }), | |
| 229 | + ClientTracingLive, | |
| 230 | + ), | |
| 220 | 231 | ); |
| 221 | 232 | const clientScope = runtime.runSync(Scope.make()); |
| 222 | 233 | return { |