fix: improve event stream reconnect reliability

Recover stalled event streams without dropping the session Wait briefly for reconnection before showing connection lost errors Persist Electron server logs for easier disconnect debugging
2026-04-22 21:03:02 +03:00
parent 17dd526731
commit f24e6de21b
10 changed files with 493 additions and 139 deletions
@@ -666,6 +666,62 @@ describe('createEventPipeline', () => {
      },
    ]);
  });
+
+  it('marks the pipeline disconnected on heartbeat timeout and recovers on the next websocket connect', async () => {
+    installDomStubs();
+    globalThis.WebSocket = FakeWebSocket;
+
+    const disconnectReasons = [];
+    let reconnectCount = 0;
+
+    const sdk = {
+      global: {
+        event: async () => {
+          throw new Error('SSE should not be used in ws mode');
+        },
+      },
+    };
+
+    const recovered = new Promise((resolve) => {
+      const { cleanup } = createEventPipeline({
+        sdk,
+        transport: 'ws',
+        heartbeatTimeoutMs: 20,
+        reconnectDelayMs: 0,
+        wsReadyTimeoutMs: 20,
+        onEvent: () => {},
+        onDisconnect: (reason) => {
+          disconnectReasons.push(reason);
+        },
+        onReconnect: () => {
+          reconnectCount += 1;
+          if (reconnectCount === 2) {
+            cleanup();
+            resolve();
+          }
+        },
+      });
+    });
+
+    await Promise.resolve();
+
+    const firstSocket = FakeWebSocket.instances[0];
+    firstSocket.emitOpen();
+    firstSocket.emitMessage({ type: 'ready', scope: 'global' });
+
+    await new Promise((resolve) => setTimeout(resolve, 35));
+
+    const secondSocket = FakeWebSocket.instances[1];
+    expect(secondSocket).toBeDefined();
+
+    secondSocket.emitOpen();
+    secondSocket.emitMessage({ type: 'ready', scope: 'global' });
+
+    await recovered;
+
+    expect(disconnectReasons).toEqual(['ws_heartbeat_timeout']);
+    expect(reconnectCount).toBe(2);
+  });
 });

 // ---------------------------------------------------------------------------
@@ -21,10 +21,10 @@ export type FlushHandler = (events: QueuedEvent[]) => void

 const FLUSH_FRAME_MS = 33
 const STREAM_YIELD_MS = 8
-const RECONNECT_DELAY_MS = 250
-const HEARTBEAT_TIMEOUT_MS = 15_000
+const DEFAULT_RECONNECT_DELAY_MS = 250
+const DEFAULT_HEARTBEAT_TIMEOUT_MS = 30_000
 const WS_FALLBACK_WINDOW_MS = 60_000
-const WS_READY_TIMEOUT_MS = 2_000
+const DEFAULT_WS_READY_TIMEOUT_MS = 2_000
 const ABSOLUTE_URL_PATTERN = /^[a-zA-Z][a-zA-Z\d+\-.]*:\/\//

 export type EventPipelineInput = {
@@ -38,6 +38,9 @@ export type EventPipelineInput = {
  /** Called when transport switches (e.g. WS timeout → SSE fallback) without actual disconnection. */
  onTransportSwitch?: () => void
  transport?: "auto" | "ws" | "sse"
+  heartbeatTimeoutMs?: number
+  reconnectDelayMs?: number
+  wsReadyTimeoutMs?: number
 }

 type MessageStreamWsFrame = {
@@ -145,8 +148,25 @@ type DirectoryQueue = {
  last: number
 }

+type AttemptAbortReason =
+  | "pipeline_stopped"
+  | "ws_heartbeat_timeout"
+  | "sse_heartbeat_timeout"
+  | null
+
 export function createEventPipeline(input: EventPipelineInput) {
-  const { sdk, onEvent, onReconnect, onDisconnect, onTransportSwitch, routeDirectory, transport = "auto" } = input
+  const {
+    sdk,
+    onEvent,
+    onReconnect,
+    onDisconnect,
+    onTransportSwitch,
+    routeDirectory,
+    transport = "auto",
+    heartbeatTimeoutMs = DEFAULT_HEARTBEAT_TIMEOUT_MS,
+    reconnectDelayMs = DEFAULT_RECONNECT_DELAY_MS,
+    wsReadyTimeoutMs = DEFAULT_WS_READY_TIMEOUT_MS,
+  } = input
  const abort = new AbortController()
  let disconnected = false
  let lastEventId: string | undefined
@@ -244,6 +264,16 @@ export function createEventPipeline(input: EventPipelineInput) {
  let attempt: AbortController | undefined
  let lastEventAt = Date.now()
  let heartbeat: ReturnType<typeof setTimeout> | undefined
+  let activeTransport: "ws" | "sse" = transport === "ws" ? "ws" : "sse"
+  let attemptAbortReason: AttemptAbortReason = null
+
+  const notifyDisconnected = (reason: string) => {
+    if (disconnected) {
+      return
+    }
+    disconnected = true
+    onDisconnect?.(reason)
+  }

  const markConnected = () => {
    disconnected = false
@@ -295,8 +325,9 @@ export function createEventPipeline(input: EventPipelineInput) {
    lastEventAt = Date.now()
    if (heartbeat) clearTimeout(heartbeat)
    heartbeat = setTimeout(() => {
+      attemptAbortReason = `${activeTransport}_heartbeat_timeout`
      attempt?.abort()
-    }, HEARTBEAT_TIMEOUT_MS)
+    }, heartbeatTimeoutMs)
  }

  const clearHeartbeat = () => {
@@ -359,7 +390,7 @@ export function createEventPipeline(input: EventPipelineInput) {
        } catch {
          // ignore
        }
-      }, WS_READY_TIMEOUT_MS)
+      }, wsReadyTimeoutMs)

      const cleanup = () => {
        if (readyTimer) {
@@ -499,9 +530,12 @@ export function createEventPipeline(input: EventPipelineInput) {
    while (!abort.signal.aborted) {
      attempt = new AbortController()
      lastEventAt = Date.now()
-      let retryDelayMs = RECONNECT_DELAY_MS
+      attemptAbortReason = null
+      let retryDelayMs = reconnectDelayMs
      const currentTransport = resolveTransport()
+      activeTransport = currentTransport
      const onAbort = () => {
+        attemptAbortReason = "pipeline_stopped"
        attempt?.abort()
      }
      abort.signal.addEventListener("abort", onAbort)
@@ -531,21 +565,18 @@ export function createEventPipeline(input: EventPipelineInput) {
          // update connection state (e.g. set isConnected = false).
          // Guard: only fire once per disconnection cycle to avoid repeated
          // setState calls on every failed retry attempt.
-          if (!disconnected) {
-            disconnected = true
-            const taggedReason = typeof error === "object" && error !== null
-              ? (error as { reason?: unknown }).reason
-              : undefined
-            const message = typeof error === "object" && error !== null
-              ? (error as { message?: unknown }).message
-              : undefined
-            const reason = typeof taggedReason === "string" && taggedReason.length > 0
-              ? taggedReason
-              : typeof message === "string" && message.length > 0
-                ? `${currentTransport}_error:${message.slice(0, 80)}`
-                : `${currentTransport}_error:unknown`
-            onDisconnect?.(reason)
-          }
+          const taggedReason = typeof error === "object" && error !== null
+            ? (error as { reason?: unknown }).reason
+            : undefined
+          const message = typeof error === "object" && error !== null
+            ? (error as { message?: unknown }).message
+            : undefined
+          const reason = typeof taggedReason === "string" && taggedReason.length > 0
+            ? taggedReason
+            : typeof message === "string" && message.length > 0
+              ? `${currentTransport}_error:${message.slice(0, 80)}`
+              : `${currentTransport}_error:unknown`
+          notifyDisconnected(reason)
        }
      } finally {
        abort.signal.removeEventListener("abort", onAbort)
@@ -554,6 +585,11 @@ export function createEventPipeline(input: EventPipelineInput) {
      }

      if (abort.signal.aborted) return
+      if (attemptAbortReason && attemptAbortReason !== "pipeline_stopped") {
+        notifyDisconnected(attemptAbortReason)
+        retryDelayMs = 0
+        attemptAbortReason = null
+      }
      if (retryDelayMs > 0) {
        await wait(retryDelayMs)
      }
@@ -563,7 +599,7 @@ export function createEventPipeline(input: EventPipelineInput) {
  const onVisibility = () => {
    if (typeof document === "undefined") return
    if (document.visibilityState !== "visible") return
-    if (Date.now() - lastEventAt < HEARTBEAT_TIMEOUT_MS) return
+    if (Date.now() - lastEventAt < heartbeatTimeoutMs) return
    attempt?.abort()
  }

@@ -55,11 +55,30 @@ function dir() {
 }

 function connectionLostError(): Error {
-  const reason = useConfigStore.getState().lastDisconnectReason
-  const suffix = reason ? ` (${reason})` : " (never connected)"
+  const { hasEverConnected, lastDisconnectReason } = useConfigStore.getState()
+  const suffix = lastDisconnectReason
+    ? ` (${lastDisconnectReason})`
+    : hasEverConnected
+      ? ""
+      : " (never connected)"
  return new Error(`Connection lost${suffix}. Please wait for reconnection.`)
 }

+// Wait briefly for the pipeline to re-establish connection before failing a
+// send. Transient reconnects (heartbeat race, WS→SSE fallback, brief network
+// blip) otherwise surface as a hard "Connection lost" toast even though the
+// pipeline recovers within a second. Poll isConnected at 100ms intervals.
+const CONNECTION_GRACE_MS = 2000
+export async function waitForConnectionOrThrow(): Promise<void> {
+  if (useConfigStore.getState().isConnected) return
+  const deadline = Date.now() + CONNECTION_GRACE_MS
+  while (Date.now() < deadline) {
+    await new Promise((resolve) => setTimeout(resolve, 100))
+    if (useConfigStore.getState().isConnected) return
+  }
+  throw connectionLostError()
+}
+
 function getSessionDirectory(sessionId: string): string | undefined {
  return useSessionUIStore.getState().getDirectoryForSession(sessionId) || dir()
 }
@@ -332,9 +351,7 @@ export async function optimisticSend(input: {
    throw new Error("Optimistic refs not set — is useSync() mounted?")
  }

-  if (!useConfigStore.getState().isConnected) {
-    throw connectionLostError()
-  }
+  await waitForConnectionOrThrow()

  const store = dirStore()
  const messageID = ascendingId("msg")
@@ -419,9 +436,7 @@ export async function respondToPermission(
  requestId: string,
  response: "once" | "always" | "reject",
 ): Promise<void> {
-  if (!useConfigStore.getState().isConnected) {
-    throw connectionLostError()
-  }
+  await waitForConnectionOrThrow()
  const result = await getRequestReplyClient("permission", sessionId, requestId).permission.reply({
    requestID: requestId,
    reply: response,
@@ -435,9 +450,7 @@ export async function dismissPermission(
  sessionId: string,
  requestId: string,
 ): Promise<void> {
-  if (!useConfigStore.getState().isConnected) {
-    throw connectionLostError()
-  }
+  await waitForConnectionOrThrow()
  const result = await getRequestReplyClient("permission", sessionId, requestId).permission.reply({
    requestID: requestId,
    reply: "reject",
@@ -456,9 +469,7 @@ export async function respondToQuestion(
  requestId: string,
  answers: string[] | string[][],
 ): Promise<void> {
-  if (!useConfigStore.getState().isConnected) {
-    throw connectionLostError()
-  }
+  await waitForConnectionOrThrow()
  const result = await getRequestReplyClient("question", sessionId, requestId).question.reply({
    requestID: requestId,
    answers: answers as Array<Array<string>>,
@@ -472,9 +483,7 @@ export async function rejectQuestion(
  sessionId: string,
  requestId: string,
 ): Promise<void> {
-  if (!useConfigStore.getState().isConnected) {
-    throw connectionLostError()
-  }
+  await waitForConnectionOrThrow()
  const result = await getRequestReplyClient("question", sessionId, requestId).question.reject({
    requestID: requestId,
  })
@@ -1376,7 +1376,11 @@ export function SyncProvider(props: {
        handleEvent(directory, payload, childStores, routingIndex)
      },
      onReconnect: () => {
-        useConfigStore.setState({ isConnected: true, lastDisconnectReason: null })
+        useConfigStore.setState({
+          isConnected: true,
+          hasEverConnected: true,
+          connectionPhase: "connected",
+        })
        for (const [dir, store] of childStores.children) {
          if (reconnectResyncing.has(dir)) continue
          if (getReconnectCandidateSessionIds(store.getState()).length === 0) continue
@@ -1392,13 +1396,22 @@ export function SyncProvider(props: {
        }
      },
      onDisconnect: (reason) => {
-        useConfigStore.setState({ isConnected: false, lastDisconnectReason: reason })
+        const { hasEverConnected } = useConfigStore.getState()
+        useConfigStore.setState({
+          isConnected: false,
+          connectionPhase: hasEverConnected ? "reconnecting" : "connecting",
+          lastDisconnectReason: reason,
+        })
      },
      onTransportSwitch: () => {
        // Transport switched (e.g. WS timeout → SSE fallback) without
        // actual disconnection. No events lost — just update connection
        // state without triggering a full directory resync.
-        useConfigStore.setState({ isConnected: true })
+        useConfigStore.setState({
+          isConnected: true,
+          hasEverConnected: true,
+          connectionPhase: "connected",
+        })
      },
    })
    return cleanup