Fix chat replay recovery during WebSocket reconnect (resubscribe)

apps/web/src/orchestrationRecovery.ts

@@ -1,4 +1,8 @@
-export type OrchestrationRecoveryReason = "bootstrap" | "sequence-gap" | "replay-failed";
+export type OrchestrationRecoveryReason =
+  | "bootstrap"
+  | "sequence-gap"
+  | "resubscribe"
+  | "replay-failed";
 
 export interface OrchestrationRecoveryPhase {
   kind: "snapshot" | "replay";

apps/web/src/routes/__root.tsx

@@ -428,8 +428,8 @@ function EventRouter() {
       queueMicrotask(flushPendingDomainEvents);
     };
 
-    const recoverFromSequenceGap = async (): Promise<void> => {
-      if (!recovery.beginReplayRecovery("sequence-gap")) {
+    const runReplayRecovery = async (reason: "sequence-gap" | "resubscribe"): Promise<void> => {
+      if (!recovery.beginReplayRecovery(reason)) {
         return;
       }
 
@@ -466,7 +466,7 @@ function EventRouter() {
               return;
             }
           }
-          void recoverFromSequenceGap();
+          void runReplayRecovery(reason);
         } else if (replayCompletion.shouldReplay && import.meta.env.MODE !== "test") {
           console.warn(
             "[orchestration-recovery]",
@@ -505,7 +505,7 @@ function EventRouter() {
           syncServerReadModel(snapshot);
           reconcileSnapshotDerivedState();
           if (recovery.completeSnapshotRecovery(snapshot.snapshotSequence)) {
-            void recoverFromSequenceGap();
+            void runReplayRecovery("sequence-gap");
           }
         }
       } catch {
@@ -522,18 +522,29 @@ function EventRouter() {
     const fallbackToSnapshotRecovery = async (): Promise<void> => {
       await runSnapshotRecovery("replay-failed");
     };
-    const unsubDomainEvent = api.orchestration.onDomainEvent((event) => {
-      const action = recovery.classifyDomainEvent(event.sequence);
-      if (action === "apply") {
-        pendingDomainEvents.push(event);
-        schedulePendingDomainEventFlush();
-        return;
-      }
-      if (action === "recover") {
-        flushPendingDomainEvents();
-        void recoverFromSequenceGap();
-      }
-    });
+    const unsubDomainEvent = api.orchestration.onDomainEvent(
+      (event) => {
+        const action = recovery.classifyDomainEvent(event.sequence);
+        if (action === "apply") {
+          pendingDomainEvents.push(event);
+          schedulePendingDomainEventFlush();
+          return;
+        }
+        if (action === "recover") {
+          flushPendingDomainEvents();
+          void runReplayRecovery("sequence-gap");
+        }
+      },
+      {
+        onResubscribe: () => {
+          if (disposed) {
+            return;
+          }
+          flushPendingDomainEvents();
+          void runReplayRecovery("resubscribe");
+        },
+      },
+    );
     const unsubTerminalEvent = api.terminal.onEvent((event) => {
       const thread = useStore.getState().threads.find((entry) => entry.id === event.threadId);
       if (thread && thread.archivedAt !== null) {

apps/web/src/wsNativeApi.test.ts

@@ -243,6 +243,20 @@ describe("wsNativeApi", () => {
     expect(onDomainEvent).toHaveBeenCalledWith(orchestrationEvent);
   });
 
+  it("forwards orchestration stream subscription options to the RPC client", async () => {
+    const { createWsNativeApi } = await import("./wsNativeApi");
+
+    const api = createWsNativeApi();
+    const onDomainEvent = vi.fn();
+    const onResubscribe = vi.fn();
+
+    api.orchestration.onDomainEvent(onDomainEvent, { onResubscribe });
+
+    expect(rpcClientMock.orchestration.onDomainEvent).toHaveBeenCalledWith(onDomainEvent, {
+      onResubscribe,
+    });
+  });
+
   it("sends orchestration dispatch commands as the direct RPC payload", async () => {
     rpcClientMock.orchestration.dispatchCommand.mockResolvedValue({ sequence: 1 });
     const { createWsNativeApi } = await import("./wsNativeApi");

apps/web/src/wsNativeApi.ts

@@ -98,7 +98,8 @@ export function createWsNativeApi(): NativeApi {
         rpcClient.orchestration
           .replayEvents({ fromSequenceExclusive })
           .then((events) => [...events]),
-      onDomainEvent: (callback) => rpcClient.orchestration.onDomainEvent(callback),
+      onDomainEvent: (callback, options) =>
+        rpcClient.orchestration.onDomainEvent(callback, options),
     },
   };
 

apps/web/src/wsRpcClient.ts

@@ -16,6 +16,10 @@ type RpcTag = keyof WsRpcProtocolClient & string;
 type RpcMethod<TTag extends RpcTag> = WsRpcProtocolClient[TTag];
 type RpcInput<TTag extends RpcTag> = Parameters<RpcMethod<TTag>>[0];
 
+interface StreamSubscriptionOptions {
+  readonly onResubscribe?: () => void;
+}
+
 type RpcUnaryMethod<TTag extends RpcTag> =
   RpcMethod<TTag> extends (input: any, options?: any) => Effect.Effect<infer TSuccess, any, any>
     ? (input: RpcInput<TTag>) => Promise<TSuccess>
@@ -28,7 +32,7 @@ type RpcUnaryNoArgMethod<TTag extends RpcTag> =
 
 type RpcStreamMethod<TTag extends RpcTag> =
   RpcMethod<TTag> extends (input: any, options?: any) => Stream.Stream<infer TEvent, any, any>
-    ? (listener: (event: TEvent) => void) => () => void
+    ? (listener: (event: TEvent) => void, options?: StreamSubscriptionOptions) => () => void
     : never;
 
 interface GitRunStackedActionOptions {
@@ -120,8 +124,12 @@ export function createWsRpcClient(transport = new WsTransport()): WsRpcClient {
       clear: (input) => transport.request((client) => client[WS_METHODS.terminalClear](input)),
       restart: (input) => transport.request((client) => client[WS_METHODS.terminalRestart](input)),
       close: (input) => transport.request((client) => client[WS_METHODS.terminalClose](input)),
-      onEvent: (listener) =>
-        transport.subscribe((client) => client[WS_METHODS.subscribeTerminalEvents]({}), listener),
+      onEvent: (listener, options) =>
+        transport.subscribe(
+          (client) => client[WS_METHODS.subscribeTerminalEvents]({}),
+          listener,
+          options,
+        ),
     },
     projects: {
       searchEntries: (input) =>
@@ -179,10 +187,18 @@ export function createWsRpcClient(transport = new WsTransport()): WsRpcClient {
       getSettings: () => transport.request((client) => client[WS_METHODS.serverGetSettings]({})),
       updateSettings: (patch) =>
         transport.request((client) => client[WS_METHODS.serverUpdateSettings]({ patch })),
-      subscribeConfig: (listener) =>
-        transport.subscribe((client) => client[WS_METHODS.subscribeServerConfig]({}), listener),
-      subscribeLifecycle: (listener) =>
-        transport.subscribe((client) => client[WS_METHODS.subscribeServerLifecycle]({}), listener),
+      subscribeConfig: (listener, options) =>
+        transport.subscribe(
+          (client) => client[WS_METHODS.subscribeServerConfig]({}),
+          listener,
+          options,
+        ),
+      subscribeLifecycle: (listener, options) =>
+        transport.subscribe(
+          (client) => client[WS_METHODS.subscribeServerLifecycle]({}),
+          listener,
+          options,
+        ),
     },
     orchestration: {
       getSnapshot: () =>
@@ -197,10 +213,11 @@ export function createWsRpcClient(transport = new WsTransport()): WsRpcClient {
         transport
           .request((client) => client[ORCHESTRATION_WS_METHODS.replayEvents](input))
           .then((events) => [...events]),
-      onDomainEvent: (listener) =>
+      onDomainEvent: (listener, options) =>
         transport.subscribe(
           (client) => client[WS_METHODS.subscribeOrchestrationDomainEvents]({}),
           listener,
+          options,
         ),
     },
   };

apps/web/src/wsTransport.test.ts

@@ -250,10 +250,12 @@ describe("WsTransport", () => {
   it("re-subscribes stream listeners after the stream exits", async () => {
     const transport = new WsTransport("ws://localhost:3020");
     const listener = vi.fn();
+    const onResubscribe = vi.fn();
 
     const unsubscribe = transport.subscribe(
       (client) => client[WS_METHODS.subscribeServerLifecycle]({}),
       listener,
+      { onResubscribe },
     );
     await waitFor(() => {
       expect(sockets).toHaveLength(1);
@@ -301,6 +303,7 @@ describe("WsTransport", () => {
         .find((message) => message._tag === "Request" && message.id !== firstRequest.id);
       expect(nextRequest).toBeDefined();
     });
+    expect(onResubscribe).toHaveBeenCalledOnce();
 
     const secondRequest = socket.sent
       .map((message) => JSON.parse(message) as { _tag?: string; id?: string; tag?: string })
@@ -339,6 +342,52 @@ describe("WsTransport", () => {
     await transport.dispose();
   });
 
+  it("does not fire onResubscribe when the first stream attempt exits before any value", async () => {
+    const transport = new WsTransport("ws://localhost:3020");
+    const listener = vi.fn();
+    const onResubscribe = vi.fn();
+
+    const unsubscribe = transport.subscribe(
+      (client) => client[WS_METHODS.subscribeServerLifecycle]({}),
+      listener,
+      { onResubscribe },
+    );
+    await waitFor(() => {
+      expect(sockets).toHaveLength(1);
+    });
+
+    const socket = getSocket();
+    socket.open();
+
+    await waitFor(() => {
+      expect(socket.sent).toHaveLength(1);
+    });
+
+    const firstRequest = JSON.parse(socket.sent[0] ?? "{}") as { id: string };
+    socket.serverMessage(
+      JSON.stringify({
+        _tag: "Exit",
+        requestId: firstRequest.id,
+        exit: {
+          _tag: "Success",
+          value: null,
+        },
+      }),
+    );
+
+    await waitFor(() => {
+      const nextRequest = socket.sent
+        .map((message) => JSON.parse(message) as { _tag?: string; id?: string })
+        .find((message) => message._tag === "Request" && message.id !== firstRequest.id);
+      expect(nextRequest).toBeDefined();
+    });
+    expect(onResubscribe).not.toHaveBeenCalled();
+    expect(listener).not.toHaveBeenCalled();
+
+    unsubscribe();
+    await transport.dispose();
+  });
+
   it("streams finite request events without re-subscribing", async () => {
     const transport = new WsTransport("ws://localhost:3020");
     const listener = vi.fn();

apps/web/src/wsTransport.ts

@@ -9,6 +9,7 @@ import { RpcClient } from "effect/unstable/rpc";
 
 interface SubscribeOptions {
   readonly retryDelay?: Duration.Input;
+  readonly onResubscribe?: () => void;
 }
 
 interface RequestOptions {
@@ -82,15 +83,27 @@ export class WsTransport {
     }
 
     let active = true;
+    let hasReceivedValue = false;
     const retryDelayMs = options?.retryDelay ?? DEFAULT_SUBSCRIPTION_RETRY_DELAY_MS;
     const cancel = this.runtime.runCallback(
-      Effect.promise(() => this.clientPromise).pipe(
+      Effect.sync(() => {
+        if (!hasReceivedValue) {
+          return;
+        }
+        try {
+          options?.onResubscribe?.();
+        } catch {
+          // Swallow reconnect hook errors so the stream can recover.
+        }
+      }).pipe(
+        Effect.andThen(Effect.promise(() => this.clientPromise)),
         Effect.flatMap((client) =>
           Stream.runForEach(connect(client), (value) =>
             Effect.sync(() => {
               if (!active) {
                 return;
               }
+              hasReceivedValue = true;
               try {
                 listener(value);
               } catch {

packages/contracts/src/ipc.ts

@@ -180,6 +180,11 @@ export interface NativeApi {
       input: OrchestrationGetFullThreadDiffInput,
     ) => Promise<OrchestrationGetFullThreadDiffResult>;
     replayEvents: (fromSequenceExclusive: number) => Promise<OrchestrationEvent[]>;
-    onDomainEvent: (callback: (event: OrchestrationEvent) => void) => () => void;
+    onDomainEvent: (
+      callback: (event: OrchestrationEvent) => void,
+      options?: {
+        onResubscribe?: () => void;
+      },
+    ) => () => void;
   };
 }