Handle intentional WS closes and heartbeat freshness

- Ignore intentional close events during transport disposal
- Track heartbeat freshness from pong traffic and reset on reconnect
- Propagate request lifecycle hooks through the RPC client

Author: juliusmarminge

apps/web/src/rpc/wsTransport.test.ts

@@ -242,10 +242,15 @@ describe("WsTransport (web instrumentation)", () => {
     socket.close(1012, "service restart");
 
     await waitFor(() => {
-      expect(onClose).toHaveBeenCalledWith({
-        code: 1012,
-        reason: "service restart",
-      });
+      expect(onClose).toHaveBeenCalledWith(
+        {
+          code: 1012,
+          reason: "service restart",
+        },
+        {
+          intentional: false,
+        },
+      );
       expect(getWsConnectionStatus()).toMatchObject({
         attemptCount: 2,
         closeReason: "service restart",

apps/web/src/rpc/wsTransport.ts

@@ -31,8 +31,11 @@ function createWsRpcProtocolLayer(
         clearAllTrackedRpcRequests();
         recordWsConnectionErrored(message);
       },
-      onClose: (details) => {
+      onClose: (details, context) => {
         clearAllTrackedRpcRequests();
+        if (context.intentional) {
+          return;
+        }
         recordWsConnectionClosed(details);
       },
     },

packages/client-runtime/src/wsRpcClient.test.ts

@@ -43,12 +43,13 @@ describe("createWsRpcClient", () => {
       reconnect: vi.fn(async () => {
         order.push("reconnect");
       }),
+      isHeartbeatFresh: vi.fn(() => true),
       request: vi.fn(),
       requestStream: vi.fn(),
       subscribe: vi.fn(() => () => undefined),
     } satisfies Pick<
       WsTransport,
-      "dispose" | "reconnect" | "request" | "requestStream" | "subscribe"
+      "dispose" | "isHeartbeatFresh" | "reconnect" | "request" | "requestStream" | "subscribe"
     >;
 
     const client = createWsRpcClient(transport as unknown as WsTransport, {
@@ -61,6 +62,26 @@ describe("createWsRpcClient", () => {
     expect(order).toEqual(["beforeReconnect", "reconnect"]);
   });
 
+  it("delegates heartbeat freshness to the transport", () => {
+    const isHeartbeatFresh = vi.fn(() => true);
+    const transport = {
+      dispose: vi.fn(async () => undefined),
+      reconnect: vi.fn(async () => undefined),
+      isHeartbeatFresh,
+      request: vi.fn(),
+      requestStream: vi.fn(),
+      subscribe: vi.fn(() => () => undefined),
+    } satisfies Pick<
+      WsTransport,
+      "dispose" | "isHeartbeatFresh" | "reconnect" | "request" | "requestStream" | "subscribe"
+    >;
+
+    const client = createWsRpcClient(transport as unknown as WsTransport);
+
+    expect(client.isHeartbeatFresh()).toBe(true);
+    expect(isHeartbeatFresh).toHaveBeenCalledOnce();
+  });
+
   it("reduces vcs status stream events into flat status snapshots", () => {
     const subscribe = vi.fn(<TValue>(_connect: unknown, listener: (value: TValue) => void) => {
       for (const event of [
@@ -89,12 +110,13 @@ describe("createWsRpcClient", () => {
     const transport = {
       dispose: vi.fn(async () => undefined),
       reconnect: vi.fn(async () => undefined),
+      isHeartbeatFresh: vi.fn(() => true),
       request: vi.fn(),
       requestStream: vi.fn(),
       subscribe,
     } satisfies Pick<
       WsTransport,
-      "dispose" | "reconnect" | "request" | "requestStream" | "subscribe"
+      "dispose" | "isHeartbeatFresh" | "reconnect" | "request" | "requestStream" | "subscribe"
     >;
 
     const client = createWsRpcClient(transport as unknown as WsTransport);
@@ -134,12 +156,13 @@ describe("createWsRpcClient", () => {
     const transport = {
       dispose: vi.fn(async () => undefined),
       reconnect: vi.fn(async () => undefined),
+      isHeartbeatFresh: vi.fn(() => true),
       request: vi.fn(),
       requestStream: vi.fn(),
       subscribe,
     } satisfies Pick<
       WsTransport,
-      "dispose" | "reconnect" | "request" | "requestStream" | "subscribe"
+      "dispose" | "isHeartbeatFresh" | "reconnect" | "request" | "requestStream" | "subscribe"
     >;
 
     const client = createWsRpcClient(transport as unknown as WsTransport);

packages/client-runtime/src/wsRpcClient.ts

@@ -125,7 +125,9 @@ export interface WsRpcClient {
   };
   readonly server: {
     readonly getConfig: RpcUnaryNoArgMethod<typeof WS_METHODS.serverGetConfig>;
-    readonly refreshProviders: RpcUnaryNoArgMethod<typeof WS_METHODS.serverRefreshProviders>;
+    readonly refreshProviders: (
+      input?: RpcInput<typeof WS_METHODS.serverRefreshProviders>,
+    ) => ReturnType<RpcUnaryMethod<typeof WS_METHODS.serverRefreshProviders>>;
     readonly discoverSourceControl: RpcUnaryNoArgMethod<
       typeof WS_METHODS.serverDiscoverSourceControl
     >;
@@ -171,7 +173,7 @@ export function createWsRpcClient(
 ): WsRpcClient {
   return {
     dispose: () => transport.dispose(),
-    isHeartbeatFresh: () => false,
+    isHeartbeatFresh: () => transport.isHeartbeatFresh(),
     reconnect: async () => {
       options?.beforeReconnect?.();
       await transport.reconnect();
@@ -278,8 +280,8 @@ export function createWsRpcClient(
     },
     server: {
       getConfig: () => transport.request((client) => client[WS_METHODS.serverGetConfig]({})),
-      refreshProviders: () =>
-        transport.request((client) => client[WS_METHODS.serverRefreshProviders]({})),
+      refreshProviders: (input) =>
+        transport.request((client) => client[WS_METHODS.serverRefreshProviders](input ?? {})),
       discoverSourceControl: () =>
         transport.request((client) => client[WS_METHODS.serverDiscoverSourceControl]({})),
       updateProvider: (input) =>

packages/client-runtime/src/wsRpcProtocol.ts

@@ -15,10 +15,34 @@ import {
 export interface WsProtocolLifecycleHandlers {
   readonly getConnectionLabel?: () => string | null;
   readonly getVersionMismatchHint?: () => string | null;
+  readonly isCloseIntentional?: () => boolean;
+  readonly isActive?: () => boolean;
   readonly onAttempt?: (socketUrl: string) => void;
   readonly onOpen?: () => void;
+  readonly onHeartbeatPing?: () => void;
+  readonly onHeartbeatPong?: () => void;
+  readonly onHeartbeatTimeout?: () => void;
+  readonly onRequestStart?: (info: {
+    readonly id: string;
+    readonly tag: string;
+    readonly stream: boolean;
+  }) => void;
+  readonly onRequestChunk?: (info: {
+    readonly id: string;
+    readonly tag: string;
+    readonly chunkCount: number;
+  }) => void;
+  readonly onRequestExit?: (info: {
+    readonly id: string;
+    readonly tag: string;
+    readonly stream: boolean;
+  }) => void;
+  readonly onRequestInterrupt?: (info: { readonly id: string; readonly tag?: string }) => void;
   readonly onError?: (message: string) => void;
-  readonly onClose?: (details: { readonly code: number; readonly reason: string }) => void;
+  readonly onClose?: (
+    details: { readonly code: number; readonly reason: string },
+    context: { readonly intentional: boolean },
+  ) => void;
 }
 
 export interface WsRpcProtocolRequestTelemetry {
@@ -63,48 +87,107 @@ function resolveWsRpcSocketUrl(rawUrl: string): string {
   return resolved.toString();
 }
 
-function defaultLifecycleHandlers(): Required<WsProtocolLifecycleHandlers> {
+type ResolvedLifecycleHandlers = Required<
+  Pick<
+    WsProtocolLifecycleHandlers,
+    | "getConnectionLabel"
+    | "getVersionMismatchHint"
+    | "isCloseIntentional"
+    | "isActive"
+    | "onAttempt"
+    | "onOpen"
+    | "onHeartbeatPing"
+    | "onHeartbeatPong"
+    | "onHeartbeatTimeout"
+    | "onError"
+    | "onClose"
+  >
+>;
+
+function defaultLifecycleHandlers(): ResolvedLifecycleHandlers {
   return {
     onAttempt: () => undefined,
     onOpen: () => undefined,
+    onHeartbeatPing: () => undefined,
+    onHeartbeatPong: () => undefined,
+    onHeartbeatTimeout: () => undefined,
     onError: () => undefined,
     onClose: () => undefined,
     getConnectionLabel: () => null,
     getVersionMismatchHint: () => null,
+    isCloseIntentional: () => false,
+    isActive: () => true,
   };
 }
 
 function resolveLifecycleHandlers(
   handlers: WsProtocolLifecycleHandlers | undefined,
   telemetryLifecycle: WsProtocolLifecycleHandlers | undefined,
-): Required<WsProtocolLifecycleHandlers> {
-  if (telemetryLifecycle === undefined) {
-    return {
-      ...defaultLifecycleHandlers(),
-      ...handlers,
-    };
-  }
+): ResolvedLifecycleHandlers {
+  const defaults = defaultLifecycleHandlers();
+  const isActive = handlers?.isActive ?? telemetryLifecycle?.isActive ?? defaults.isActive;
+  const isCloseIntentional =
+    handlers?.isCloseIntentional ??
+    telemetryLifecycle?.isCloseIntentional ??
+    defaults.isCloseIntentional;
 
   return {
     getConnectionLabel: () =>
-      handlers?.getConnectionLabel?.() ?? telemetryLifecycle.getConnectionLabel?.() ?? null,
+      handlers?.getConnectionLabel?.() ?? telemetryLifecycle?.getConnectionLabel?.() ?? null,
     getVersionMismatchHint: () =>
-      handlers?.getVersionMismatchHint?.() ?? telemetryLifecycle.getVersionMismatchHint?.() ?? null,
+      handlers?.getVersionMismatchHint?.() ??
+      telemetryLifecycle?.getVersionMismatchHint?.() ??
+      null,
+    isActive,
+    isCloseIntentional,
     onAttempt: (socketUrl) => {
-      telemetryLifecycle.onAttempt?.(socketUrl);
+      if (!isActive()) {
+        return;
+      }
+      telemetryLifecycle?.onAttempt?.(socketUrl);
       handlers?.onAttempt?.(socketUrl);
     },
     onOpen: () => {
-      telemetryLifecycle.onOpen?.();
+      if (!isActive()) {
+        return;
+      }
+      telemetryLifecycle?.onOpen?.();
       handlers?.onOpen?.();
     },
+    onHeartbeatPing: () => {
+      if (!isActive()) {
+        return;
+      }
+      telemetryLifecycle?.onHeartbeatPing?.();
+      handlers?.onHeartbeatPing?.();
+    },
+    onHeartbeatPong: () => {
+      if (!isActive()) {
+        return;
+      }
+      telemetryLifecycle?.onHeartbeatPong?.();
+      handlers?.onHeartbeatPong?.();
+    },
+    onHeartbeatTimeout: () => {
+      if (!isActive()) {
+        return;
+      }
+      telemetryLifecycle?.onHeartbeatTimeout?.();
+      handlers?.onHeartbeatTimeout?.();
+    },
     onError: (message) => {
-      telemetryLifecycle.onError?.(message);
+      if (!isActive()) {
+        return;
+      }
+      telemetryLifecycle?.onError?.(message);
       handlers?.onError?.(message);
     },
-    onClose: (details) => {
-      telemetryLifecycle.onClose?.(details);
-      handlers?.onClose?.(details);
+    onClose: (details, context) => {
+      if (!isActive()) {
+        return;
+      }
+      telemetryLifecycle?.onClose?.(details, context);
+      handlers?.onClose?.(details, context);
     },
   };
 }
@@ -158,10 +241,15 @@ export function createWsRpcProtocolLayer(
       socket.addEventListener(
         "close",
         (event) => {
-          lifecycle.onClose({
-            code: event.code,
-            reason: event.reason,
-          });
+          lifecycle.onClose(
+            {
+              code: event.code,
+              reason: event.reason,
+            },
+            {
+              intentional: lifecycle.isCloseIntentional(),
+            },
+          );
         },
         { once: true },
       );
@@ -210,6 +298,77 @@ export function createWsRpcProtocolLayer(
           retryTransientErrors: true,
         }),
   );
+  const requestHooksLayer = Layer.succeed(
+    RpcClient.RequestHooks,
+    RpcClient.RequestHooks.of({
+      onRequestStart: (info) =>
+        Effect.sync(() => {
+          if (!lifecycle.isActive()) {
+            return;
+          }
+          handlers?.onRequestStart?.({
+            id: String(info.id),
+            tag: info.tag,
+            stream: info.stream,
+          });
+        }),
+      onRequestChunk: (info) =>
+        Effect.sync(() => {
+          if (!lifecycle.isActive()) {
+            return;
+          }
+          handlers?.onRequestChunk?.({
+            id: String(info.id),
+            tag: info.tag,
+            chunkCount: info.chunkCount,
+          });
+        }),
+      onRequestExit: (info) =>
+        Effect.sync(() => {
+          if (!lifecycle.isActive()) {
+            return;
+          }
+          handlers?.onRequestExit?.({
+            id: String(info.id),
+            tag: info.tag,
+            stream: info.stream,
+          });
+        }),
+      onRequestInterrupt: (info) =>
+        Effect.sync(() => {
+          if (!lifecycle.isActive()) {
+            return;
+          }
+          handlers?.onRequestInterrupt?.({
+            id: String(info.id),
+            ...(info.tag === undefined ? {} : { tag: info.tag }),
+          });
+        }),
+    }),
+  );
+  const connectionHooksLayer = Layer.succeed(
+    RpcClient.ConnectionHooks,
+    RpcClient.ConnectionHooks.of({
+      onConnect: Effect.void,
+      onDisconnect: Effect.void,
+      onPing: Effect.sync(() => {
+        lifecycle.onHeartbeatPing();
+      }),
+      onPong: Effect.sync(() => {
+        lifecycle.onHeartbeatPong();
+      }),
+      onPingTimeout: Effect.sync(() => {
+        requestTelemetry?.onClearTrackedRequests?.();
+        lifecycle.onHeartbeatTimeout();
+      }),
+    }),
+  );
 
-  return protocolLayer.pipe(Layer.provide(Layer.mergeAll(socketLayer, RpcSerialization.layerJson)));
+  return Layer.mergeAll(
+    protocolLayer.pipe(
+      Layer.provide(Layer.mergeAll(socketLayer, RpcSerialization.layerJson, connectionHooksLayer)),
+    ),
+    requestHooksLayer,
+    connectionHooksLayer,
+  );
 }