From b3569ce667b953381216cc7d1316e39adc815c2a Mon Sep 17 00:00:00 2001 From: Michael Sober Date: Mon, 7 Sep 2026 12:26:46 +0000 Subject: [PATCH 1/3] fix(bb-agent): useChat survives mid-turn WebSocket reconnect and send-path failures Long agent turns (up to 8h on AgentCore) can outlive API Gateway WebSocket limits (2h connection, 10-min idle). useChat subscribed once and assumed the socket stayed healthy, so a reconnect gap could swallow the done chunk (loading stuck true forever) and a rejected/504 send left a hung spinner + orphaned empty assistant bubble. - Widen the subscribe option to accept { onMessage, onDisconnect?, onReconnect? } (backward compatible with the bare-handler form). - On reconnect, re-sync authoritative state from the DB: getConversation to recover the final assistant text if the turn completed during the gap, and getPendingInterrupts to recover a missed interrupt. If the turn is still running, keep loading and resume streaming on the resubscribed channel. - Wrap api.sendMessage / api.resume in try/catch: reset loading, drop the empty placeholder, and surface onError on rejection. - Bounded post-reconnect failsafe clears loading if no terminal chunk arrives within a window, so the spinner can never hang indefinitely. Adds unit tests for send-path reset, reconnect re-sync (done missed / still running), pending-interrupt re-check, and the bounded failsafe. --- .changeset/usechat-ws-reconnect-resync.md | 12 ++ packages/bb-agent/src/index.hooks.ts | 184 +++++++++++++++++- packages/bb-agent/src/index.test.ts | 217 +++++++++++++++++++++- 3 files changed, 395 insertions(+), 18 deletions(-) create mode 100644 .changeset/usechat-ws-reconnect-resync.md diff --git a/.changeset/usechat-ws-reconnect-resync.md b/.changeset/usechat-ws-reconnect-resync.md new file mode 100644 index 000000000..6970dfb9c --- /dev/null +++ b/.changeset/usechat-ws-reconnect-resync.md @@ -0,0 +1,12 @@ +--- +"@aws-blocks/bb-agent": minor +--- + +useChat now survives a mid-turn Realtime WebSocket disconnect/reconnect and send-path failures. + +Long-running agent turns (up to 8h on AgentCore) can outlive API Gateway's WebSocket limits (2h max connection, 10-min idle). Previously useChat subscribed once and assumed the socket stayed healthy for the whole turn, so a reconnect gap could swallow the `done` chunk and leave `loading` stuck true, and a rejected/timed-out send (e.g. a 504 cold dispatch) left the spinner hanging with an orphaned empty assistant bubble. + +- The `subscribe` option now accepts an options object (`{ onMessage, onDisconnect?, onReconnect? }`) in addition to a bare handler — backward compatible. +- On reconnect, useChat re-syncs authoritative state from the database (`getConversation` to recover the final assistant text if the turn completed during the gap; `getPendingInterrupts` to recover a missed interrupt). If the turn is still running, loading is preserved and streaming resumes on the resubscribed channel. +- `sendMessage` and `respondToInterrupt` now reset `loading` and surface `onError` when the underlying RPC rejects, dropping any empty placeholder. +- A bounded failsafe clears `loading` if no terminal chunk arrives within a window after a reconnect, so the spinner can never hang indefinitely. diff --git a/packages/bb-agent/src/index.hooks.ts b/packages/bb-agent/src/index.hooks.ts index 22dd8168c..fae3d5c34 100644 --- a/packages/bb-agent/src/index.hooks.ts +++ b/packages/bb-agent/src/index.hooks.ts @@ -26,6 +26,34 @@ export interface ChatMessage { metadata?: Record; } +/** Handler invoked for each streaming chunk delivered over the Realtime channel. */ +export type ChatChunkHandler = (chunk: AgentStreamChunk) => void; + +/** + * Options form accepted by {@link UseChatOptions.subscribe}. + * + * Mirrors bb-realtime's `SubscribeOptions` shape so an app's `subscribe` adapter can + * forward this object straight to `channel.subscribe(...)`. useChat passes this object + * (rather than a bare handler) so it can react to a mid-turn transport disconnect and, + * on reconnect, re-sync authoritative state from the DB — the correctness backstop for + * chunks lost while the socket was down. + */ +export interface ChatSubscribeOptions { + /** Called for each incoming chunk (the streaming handler). */ + onMessage: ChatChunkHandler; + /** + * Called when the connection is lost for any reason (including a client-initiated + * `unsubscribe()`). Optional — useChat does not require it, but forwards it so an + * adapter can surface drops. + */ + onDisconnect?: (reason: string) => void; + /** + * Called after the transport transparently reconnects and this channel has been + * resubscribed. useChat uses this to re-sync from the DB (see {@link UseChatOptions.subscribe}). + */ + onReconnect?: () => void; +} + /** Options for creating a chat instance. */ export interface UseChatOptions { api: { @@ -36,11 +64,17 @@ export interface UseChatOptions { getPendingInterrupts?(conversationId: string): Promise<{ interrupts: Array<{ id: string; name: string; reason?: any }> }>; }; /** - * Subscribe to a Realtime channel. Must return an object with: + * Subscribe to a Realtime channel. Called with the channel id and either a bare + * chunk handler or a {@link ChatSubscribeOptions} object — useChat always passes the + * options object so it can react to disconnect/reconnect, but the bare-handler form + * is still accepted for backward compatibility. Adapters typically forward the second + * argument straight to `channel.subscribe(...)`, which accepts both shapes. + * + * Must return an object with: * - unsubscribe(): stop receiving messages * - established: Promise that resolves when the WS subscription is confirmed */ - subscribe: (channelId: string, handler: (chunk: AgentStreamChunk) => void) => Promise<{ unsubscribe(): void; established: Promise }>; + subscribe: (channelId: string, handlerOrOptions: ChatChunkHandler | ChatSubscribeOptions) => Promise<{ unsubscribe(): void; established: Promise }>; /** Called whenever the message list changes. */ onMessagesChange?: (messages: ChatMessage[]) => void; /** Called whenever loading state changes. */ @@ -76,6 +110,15 @@ function nextId(): string { return `msg-${++messageCounter}-${Date.now()}`; } +/** + * Last-resort window (ms) after a reconnect. If neither the terminal chunk on the + * resubscribed channel nor the getConversation re-sync resolves the turn within this + * bound, useChat stops the spinner and surfaces an error so the UI can never hang + * forever. This is a backstop, NOT the primary recovery (which is the done chunk / + * DB re-sync). Kept generous so it only fires when both of those genuinely fail. + */ +const RECONNECT_FAILSAFE_MS = 30_000; + /** * Create a chat instance for managing agent conversations. * @@ -87,9 +130,11 @@ function nextId(): string { * createConversation: () => api.agentCreateConversationId(), * getConversation: (id) => api.agentGetConversation(id), * }, - * subscribe: async (channelId, handler) => { + * subscribe: async (channelId, sub) => { * const result = await api.agentGetChannel(channelId); - * return result.channel.subscribe(handler); + * // `sub` is a ChatSubscribeOptions object (onMessage/onReconnect/onDisconnect); + * // channel.subscribe accepts it directly and wires reconnect handling for us. + * return result.channel.subscribe(sub); * }, * onMessagesChange: (msgs) => renderMessages(msgs), * onLoadingChange: (loading) => updateSpinner(loading), @@ -106,6 +151,106 @@ export function useChat(options: UseChatOptions): ChatInstance { let activeSub: { unsubscribe(): void } | null = null; let assistantId: string | null = null; let assistantText = ''; + /** Timer id for the post-reconnect failsafe (see RECONNECT_FAILSAFE_MS). null when disarmed. */ + let failsafeTimer: ReturnType | null = null; + + /** Cancel the post-reconnect failsafe timer if one is armed. */ + function clearFailsafe() { + if (failsafeTimer !== null) { + clearTimeout(failsafeTimer); + failsafeTimer = null; + } + } + + /** + * Arm the bounded post-reconnect failsafe. If the turn is still running after a + * reconnect and the terminal chunk never arrives (lost a second time), this stops + * the spinner and surfaces an error rather than hanging loading=true forever. + */ + function armFailsafe() { + clearFailsafe(); + failsafeTimer = setTimeout(() => { + failsafeTimer = null; + if (loading) { + loading = false; + options.onLoadingChange?.(loading); + options.onError?.('Timed out waiting for the agent to respond after reconnect.'); + } + }, RECONNECT_FAILSAFE_MS); + } + + /** + * Drop the optimistic empty assistant placeholder (if present and still empty). + * Used on a send-path failure so no orphaned empty bubble is left in the UI. + */ + function removeEmptyAssistantPlaceholder() { + if (!assistantId) return; + const placeholder = messages.find(m => m.id === assistantId); + if (placeholder && !placeholder.content) { + messages = messages.filter(m => m.id !== assistantId); + options.onMessagesChange?.(messages); + } + assistantId = null; + } + + /** + * Shared handler for a rejected send path (api.sendMessage / api.resume): reset + * loading, drop the dangling empty placeholder, and surface the error via onError + * (swallowed, matching how the `error` chunk is handled — no re-throw). + */ + function handleSendFailure(err: unknown) { + loading = false; + options.onLoadingChange?.(loading); + removeEmptyAssistantPlaceholder(); + options.onError?.(err instanceof Error ? err.message : String(err)); + } + + /** + * Re-sync authoritative state from the DB after the transport transparently + * reconnects. Any chunks published while the socket was down were missed, so the + * persisted conversation is the source of truth: + * - If the turn completed server-side (history ends with a non-empty assistant + * message), adopt that final text into the in-flight bubble and clear loading — + * the `done` chunk was lost in the gap. + * - If the turn is still running (no final assistant message yet), keep loading + * true and wait for the terminal chunk on the resubscribed channel, arming a + * bounded failsafe so the spinner can't hang if that chunk is also lost. + * Also re-checks pending interrupts, which may have been raised during the gap. + */ + async function handleReconnect() { + if (!conversationId) return; + try { + const { messages: history } = await options.api.getConversation(conversationId); + const last = history[history.length - 1]; + const turnComplete = !!last && last.role === 'assistant' && !!last.content; + + if (turnComplete && assistantId) { + // Turn finished while we were disconnected; the terminal `done` chunk was lost. + // Replace the in-flight assistant bubble with the persisted final text. + assistantText = last.content; + messages = messages.map(m => (m.id === assistantId ? { ...m, content: last.content } : m)); + options.onMessagesChange?.(messages); + assistantId = null; + clearFailsafe(); + loading = false; + options.onLoadingChange?.(loading); + } else if (loading) { + // Turn still running server-side — do NOT clear loading. Wait for the terminal + // chunk on the resubscribed channel, guarded by the bounded failsafe. + armFailsafe(); + } + + // A pending interrupt may have been raised while the socket was down. + if (options.api.getPendingInterrupts) { + const { interrupts } = await options.api.getPendingInterrupts(conversationId); + if (interrupts.length) options.onInterrupt?.(interrupts); + } + } catch (err) { + // Re-sync itself failed. Surface it, and keep the spinner honest via the failsafe. + if (loading) armFailsafe(); + options.onError?.(err instanceof Error ? err.message : String(err)); + } + } /** Handle a chunk from the Realtime subscription. */ function handleChunk(chunk: AgentStreamChunk) { @@ -122,11 +267,13 @@ export function useChat(options: UseChatOptions): ChatInstance { messages = messages.map(m => m.id === assistantId ? { ...m, content: chunk.text! } : m); options.onMessagesChange?.(messages); } + clearFailsafe(); loading = false; options.onLoadingChange?.(loading); } if (chunk.type === 'error') { + clearFailsafe(); loading = false; options.onLoadingChange?.(loading); options.onError?.(chunk.error ?? 'Unknown error'); @@ -142,6 +289,7 @@ export function useChat(options: UseChatOptions): ChatInstance { } } assistantId = null; + clearFailsafe(); loading = false; options.onLoadingChange?.(loading); options.onInterrupt?.(chunk.interrupts); @@ -152,13 +300,20 @@ export function useChat(options: UseChatOptions): ChatInstance { async function ensureSubscribed(channelId: string) { if (activeSub) { activeSub.unsubscribe(); activeSub = null; } - const sub = await options.subscribe(channelId, handleChunk); + // Pass an options object (not a bare handler) so the transport can notify us on + // reconnect — we re-sync authoritative state from the DB in handleReconnect(). + const subscribeOptions: ChatSubscribeOptions = { + onMessage: handleChunk, + onReconnect: () => { void handleReconnect(); }, + }; + + const sub = await options.subscribe(channelId, subscribeOptions); try { await sub.established; } catch (err) { console.warn('Subscription failed, retrying with fresh token:', err); sub.unsubscribe(); - const retrySub = await options.subscribe(channelId, handleChunk); + const retrySub = await options.subscribe(channelId, subscribeOptions); await retrySub.established; activeSub = retrySub; return; @@ -192,7 +347,13 @@ export function useChat(options: UseChatOptions): ChatInstance { options.onLoadingChange?.(loading); // Submit — chunks arrive via the already-open subscription - await options.api.sendMessage(conversationId, text, conversationId); + try { + await options.api.sendMessage(conversationId, text, conversationId); + } catch (err) { + // Send failed (e.g. 504) — the turn never started server-side. Reset loading, + // drop the empty assistant placeholder, and surface the error via onError. + handleSendFailure(err); + } }, async respondToInterrupt(responses: Array<{ interruptId: string; approved: boolean; trust?: boolean; toolName?: string; input?: any }>) { @@ -216,7 +377,13 @@ export function useChat(options: UseChatOptions): ChatInstance { loading = true; options.onLoadingChange?.(loading); if (!options.api.resume) throw new Error('respondToInterrupt requires api.resume to be configured'); - await options.api.resume(conversationId, responses, conversationId); + try { + await options.api.resume(conversationId, responses, conversationId); + } catch (err) { + // Resume failed — reset loading, drop the empty assistant placeholder, and + // surface the error via onError (consistent with the sendMessage failsafe). + handleSendFailure(err); + } }, getMessages() { return messages; }, @@ -250,6 +417,7 @@ export function useChat(options: UseChatOptions): ChatInstance { }, destroy() { + clearFailsafe(); if (activeSub) { activeSub.unsubscribe(); activeSub = null; } }, }; diff --git a/packages/bb-agent/src/index.test.ts b/packages/bb-agent/src/index.test.ts index b2bfa36bf..dc73896d6 100644 --- a/packages/bb-agent/src/index.test.ts +++ b/packages/bb-agent/src/index.test.ts @@ -1185,6 +1185,36 @@ describe('model-factory', () => { // ── useChat ────────────────────────────────────────────────────────────────── import { useChat } from './index.hooks.js'; +import type { AgentStreamChunk, ChatMessage, UseChatOptions } from './index.hooks.js'; + +/** Flush pending microtasks so an async onReconnect handler settles before assertions. */ +function flush(): Promise { + return new Promise((resolve) => setImmediate(resolve)); +} + +/** + * useChat calls `subscribe` with a {@link ChatSubscribeOptions} object. Capture its + * callbacks (onMessage / onReconnect / onDisconnect) cast-free so tests can drive chunks + * and simulate a transport reconnect. Also tolerates the bare-handler form for safety. + */ +function subscribeCapture() { + const cap: { + handler?: (chunk: AgentStreamChunk) => void; + reconnect?: () => void; + disconnect?: (reason: string) => void; + } = {}; + const subscribe: UseChatOptions['subscribe'] = async (_channelId, handlerOrOptions) => { + if (typeof handlerOrOptions === 'function') { + cap.handler = handlerOrOptions; + } else { + cap.handler = handlerOrOptions.onMessage; + cap.reconnect = handlerOrOptions.onReconnect; + cap.disconnect = handlerOrOptions.onDisconnect; + } + return { unsubscribe() {}, established: Promise.resolve() }; + }; + return { cap, subscribe }; +} describe('useChat', () => { test('onError is called when error chunk arrives', async () => { @@ -1198,8 +1228,8 @@ describe('useChat', () => { createConversation: async () => ({ conversationId: 'conv-1' }), getConversation: async () => ({ messages: [] }), }, - subscribe: async (_channelId, handler) => { - chunkHandler = handler; + subscribe: async (_channelId, handlerOrOptions) => { + chunkHandler = typeof handlerOrOptions === 'function' ? handlerOrOptions : handlerOrOptions.onMessage; return { unsubscribe() {}, established: Promise.resolve() }; }, onLoadingChange: (l) => { loadingStates.push(l); }, @@ -1225,8 +1255,8 @@ describe('useChat', () => { createConversation: async () => ({ conversationId: 'conv-1' }), getConversation: async () => ({ messages: [] }), }, - subscribe: async (_channelId, handler) => { - chunkHandler = handler; + subscribe: async (_channelId, handlerOrOptions) => { + chunkHandler = typeof handlerOrOptions === 'function' ? handlerOrOptions : handlerOrOptions.onMessage; return { unsubscribe() {}, established: Promise.resolve() }; }, onLoadingChange: (l) => { loadingStates.push(l); }, @@ -1254,8 +1284,8 @@ describe('useChat', () => { getConversation: async () => ({ messages: [] }), resume: async (channelId, responses, convId) => { resumeCalled = true; resumeArgs = { channelId, responses, convId }; }, }, - subscribe: async (_channelId, handler) => { - chunkHandler = handler; + subscribe: async (_channelId, handlerOrOptions) => { + chunkHandler = typeof handlerOrOptions === 'function' ? handlerOrOptions : handlerOrOptions.onMessage; return { unsubscribe() {}, established: Promise.resolve() }; }, }); @@ -1280,8 +1310,8 @@ describe('useChat', () => { createConversation: async () => ({ conversationId: 'conv-1' }), getConversation: async () => ({ messages: [] }), }, - subscribe: async (_channelId, handler) => { - chunkHandler = handler; + subscribe: async (_channelId, handlerOrOptions) => { + chunkHandler = typeof handlerOrOptions === 'function' ? handlerOrOptions : handlerOrOptions.onMessage; return { unsubscribe() {}, established: Promise.resolve() }; }, }); @@ -1301,8 +1331,8 @@ describe('useChat', () => { createConversation: async () => ({ conversationId: 'conv-1' }), getConversation: async () => ({ messages: [] }), }, - subscribe: async (_channelId, handler) => { - chunkHandler = handler; + subscribe: async (_channelId, handlerOrOptions) => { + chunkHandler = typeof handlerOrOptions === 'function' ? handlerOrOptions : handlerOrOptions.onMessage; return { unsubscribe() {}, established: Promise.resolve() }; }, onMessagesChange: (msgs) => { lastMessages = msgs; }, @@ -1315,6 +1345,173 @@ describe('useChat', () => { chunkHandler!({ type: 'interrupt', interrupts: [{ id: 'int-1', name: 'approve:delete' }] }); assert.ok(!lastMessages.some(m => m.role === 'assistant' && m.content === ''), 'empty placeholder should be removed'); }); + + // ── Send-path failsafes + mid-turn reconnect re-sync (PR2 / Option A) ──────── + + test('sendMessage rejection (504) resets loading and calls onError', async () => { + let errorReceived: string | undefined; + const loadingStates: boolean[] = []; + let lastMessages: ChatMessage[] = []; + const { subscribe } = subscribeCapture(); + + const chat = useChat({ + api: { + sendMessage: async () => { throw new Error('504 Gateway Timeout'); }, + createConversation: async () => ({ conversationId: 'conv-1' }), + getConversation: async () => ({ messages: [] }), + }, + subscribe, + onLoadingChange: (l) => { loadingStates.push(l); }, + onError: (e) => { errorReceived = e; }, + onMessagesChange: (m) => { lastMessages = m; }, + }); + + await chat.sendMessage('hello'); + + assert.strictEqual(chat.isLoading(), false, 'loading should be reset after send failure'); + assert.strictEqual(loadingStates.at(-1), false); + assert.ok(errorReceived, 'onError should be called'); + assert.match(errorReceived!, /504/); + assert.ok(!lastMessages.some(m => m.role === 'assistant' && m.content === ''), 'no orphaned empty assistant bubble'); + }); + + test('respondToInterrupt rejection resets loading and calls onError', async () => { + let errorReceived: string | undefined; + let lastMessages: ChatMessage[] = []; + const { cap, subscribe } = subscribeCapture(); + + const chat = useChat({ + api: { + sendMessage: async () => {}, + createConversation: async () => ({ conversationId: 'conv-1' }), + getConversation: async () => ({ messages: [] }), + resume: async () => { throw new Error('resume failed'); }, + }, + subscribe, + onError: (e) => { errorReceived = e; }, + onMessagesChange: (m) => { lastMessages = m; }, + }); + + await chat.sendMessage('hello'); + cap.handler!({ type: 'interrupt', interrupts: [{ id: 'int-1', name: 'approve:delete' }] }); + await chat.respondToInterrupt([{ interruptId: 'int-1', approved: true }]); + + assert.strictEqual(chat.isLoading(), false, 'loading should be reset after resume failure'); + assert.ok(errorReceived, 'onError should be called'); + assert.match(errorReceived!, /resume failed/); + assert.ok(!lastMessages.some(m => m.role === 'assistant' && m.content === ''), 'no orphaned empty assistant bubble'); + }); + + test('reconnect re-syncs final assistant text from getConversation when the done chunk was missed', async () => { + const { cap, subscribe } = subscribeCapture(); + + const chat = useChat({ + api: { + sendMessage: async () => {}, + createConversation: async () => ({ conversationId: 'conv-1' }), + // Turn completed server-side: history ends with a non-empty assistant message. + getConversation: async () => ({ messages: [ + { role: 'user', content: 'hello' }, + { role: 'assistant', content: 'Hello! The final persisted answer.' }, + ] }), + }, + subscribe, + }); + + await chat.sendMessage('hello'); + // Two deltas arrive, then the socket drops before `done`. + cap.handler!({ type: 'text-delta', text: 'Hel' }); + cap.handler!({ type: 'text-delta', text: 'lo' }); + // Transport reconnects — useChat re-syncs from the DB. + cap.reconnect!(); + await flush(); + + assert.strictEqual(chat.isLoading(), false, 'loading cleared after re-sync of completed turn'); + const assistant = chat.getMessages().find(m => m.role === 'assistant'); + assert.ok(assistant, 'assistant message should exist'); + assert.strictEqual(assistant!.content, 'Hello! The final persisted answer.', 'in-flight bubble replaced with persisted final text'); + chat.destroy(); + }); + + test('reconnect while turn still running keeps loading true (no premature clear)', async () => { + const { cap, subscribe } = subscribeCapture(); + + const chat = useChat({ + api: { + sendMessage: async () => {}, + createConversation: async () => ({ conversationId: 'conv-1' }), + // Turn still running: history ends with the user message (no final assistant yet). + getConversation: async () => ({ messages: [{ role: 'user', content: 'hello' }] }), + }, + subscribe, + }); + + await chat.sendMessage('hello'); + cap.reconnect!(); + await flush(); + + assert.strictEqual(chat.isLoading(), true, 'loading stays true while the turn is still running'); + + // The terminal chunk finally arrives on the resubscribed channel. + cap.handler!({ type: 'done', text: 'done at last' }); + assert.strictEqual(chat.isLoading(), false, 'a later done chunk clears loading'); + chat.destroy(); + }); + + test('reconnect re-checks pending interrupts', async () => { + let interruptsReceived: Array<{ id: string; name: string; reason?: unknown }> | undefined; + const { cap, subscribe } = subscribeCapture(); + + const chat = useChat({ + api: { + sendMessage: async () => {}, + createConversation: async () => ({ conversationId: 'conv-1' }), + getConversation: async () => ({ messages: [{ role: 'user', content: 'hello' }] }), + // An interrupt was raised server-side while the socket was down. + getPendingInterrupts: async () => ({ interrupts: [{ id: 'int-9', name: 'approve:refund' }] }), + }, + subscribe, + onInterrupt: (ints) => { interruptsReceived = ints; }, + }); + + await chat.sendMessage('hello'); + cap.reconnect!(); + await flush(); + + assert.ok(interruptsReceived, 'pending interrupt should surface after reconnect'); + assert.strictEqual(interruptsReceived!.length, 1); + assert.strictEqual(interruptsReceived![0].name, 'approve:refund'); + chat.destroy(); + }); + + test('bounded failsafe clears loading if no terminal chunk arrives after reconnect', async (t) => { + t.mock.timers.enable({ apis: ['setTimeout'] }); + let errorReceived: string | undefined; + const { cap, subscribe } = subscribeCapture(); + + const chat = useChat({ + api: { + sendMessage: async () => {}, + createConversation: async () => ({ conversationId: 'conv-1' }), + // Turn still running on reconnect — arms the failsafe. + getConversation: async () => ({ messages: [{ role: 'user', content: 'hello' }] }), + }, + subscribe, + onError: (e) => { errorReceived = e; }, + }); + + await chat.sendMessage('hello'); + cap.reconnect!(); + await flush(); + assert.strictEqual(chat.isLoading(), true, 'still loading right after reconnect (turn running)'); + + // No terminal chunk ever arrives — advance past the failsafe window. + t.mock.timers.tick(60_000); + + assert.strictEqual(chat.isLoading(), false, 'failsafe should clear loading'); + assert.ok(errorReceived, 'failsafe should surface an error'); + chat.destroy(); + }); }); describe('checkModelHealth', () => { From d902daf1c9bb47e1be911c07d8cd4b2d4c22578c Mon Sep 17 00:00:00 2001 From: Michael Sober Date: Mon, 7 Sep 2026 13:07:05 +0000 Subject: [PATCH 2/3] =?UTF-8?q?fix(bb-agent):=20address=20useChat=20reconn?= =?UTF-8?q?ect=20review=20=E2=80=94=20failsafe=20liveness,=20turn-identity?= =?UTF-8?q?=20guard,=20single=20onError?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolve code-review findings on the useChat reconnect/re-sync path: - Failsafe no longer misfires on legitimately long/silent turns: any received chunk re-arms the post-reconnect failsafe, so it only fires after a window of COMPLETE silence, and the window is raised to 11 min (RECONNECT_FAILSAFE_MS = 660_000) to sit above the Realtime 10-min idle-timeout + reconnect budget. - Reconnect re-sync can no longer clobber correct text or adopt a stale/previous turn: assistantId is nulled on done/error (reliable in-flight signal), and handleReconnect captures turnAtStart before the getConversation await and only adopts the persisted text when it is still the same in-flight turn AND the persisted content extends what was already streamed. - A send-path rejection that may have started the turn server-side no longer double-reports: onError is guarded to fire at most once per turn (reportError), matching the error-chunk path. - The widened subscribe arg is a callable that also carries onMessage/onReconnect/onDisconnect, so existing adapters that invoke arg2 keep working while new adapters can read the reconnect hooks. - Teardown guard: destroy() sets `destroyed`, and handleReconnect/armFailsafe early-return so no callback fires after unmount. Tests updated/added for failsafe liveness, live-done-vs-late-getConversation, stale-snapshot rejection, and single-onError. --- packages/bb-agent/src/index.hooks.ts | 161 +++++++++++++++++++++------ packages/bb-agent/src/index.test.ts | 134 ++++++++++++++++++++-- 2 files changed, 251 insertions(+), 44 deletions(-) diff --git a/packages/bb-agent/src/index.hooks.ts b/packages/bb-agent/src/index.hooks.ts index fae3d5c34..9069e33e1 100644 --- a/packages/bb-agent/src/index.hooks.ts +++ b/packages/bb-agent/src/index.hooks.ts @@ -111,13 +111,19 @@ function nextId(): string { } /** - * Last-resort window (ms) after a reconnect. If neither the terminal chunk on the - * resubscribed channel nor the getConversation re-sync resolves the turn within this - * bound, useChat stops the spinner and surfaces an error so the UI can never hang - * forever. This is a backstop, NOT the primary recovery (which is the done chunk / - * DB re-sync). Kept generous so it only fires when both of those genuinely fail. + * Last-resort window (ms) of COMPLETE SILENCE after a reconnect. The failsafe is a + * backstop, NOT the primary recovery (which is the done chunk / DB re-sync). It is + * re-armed by every received chunk (see handleChunk), so it only fires after a window + * in which NO chunk at all arrived — not during a long tool-call/thinking gap or a slow + * post-reconnect stream. On the multi-hour turns this feature targets, such gaps are + * normal. + * + * This MUST comfortably exceed the Realtime transport's idle-timeout + reconnect budget + * (API Gateway WebSocket: 10-min idle timeout, 2h max connection duration), otherwise a + * perfectly normal idle → disconnect → reconnect cycle would trip a spurious 'Timed out'. + * 11 minutes sits just above the 10-min idle timeout with margin for the reconnect. */ -const RECONNECT_FAILSAFE_MS = 30_000; +const RECONNECT_FAILSAFE_MS = 660_000; /** * Create a chat instance for managing agent conversations. @@ -153,6 +159,18 @@ export function useChat(options: UseChatOptions): ChatInstance { let assistantText = ''; /** Timer id for the post-reconnect failsafe (see RECONNECT_FAILSAFE_MS). null when disarmed. */ let failsafeTimer: ReturnType | null = null; + /** + * True once destroy() has run. Guards async callbacks (the reconnect re-sync and the + * failsafe timer) so no onLoadingChange/onError/onMessagesChange fires after teardown. + */ + let destroyed = false; + /** + * Whether onError has already been reported for the CURRENT turn. A send-path + * rejection (e.g. a 504) may mean the turn actually STARTED server-side, so a later + * `error` chunk can arrive for the same turn; this guard reports onError at most once + * per turn. Reset when a new turn begins (sendMessage / respondToInterrupt). + */ + let errorReported = false; /** Cancel the post-reconnect failsafe timer if one is armed. */ function clearFailsafe() { @@ -171,6 +189,8 @@ export function useChat(options: UseChatOptions): ChatInstance { clearFailsafe(); failsafeTimer = setTimeout(() => { failsafeTimer = null; + // Teardown guard: destroy() may have run while the timer was pending. + if (destroyed) return; if (loading) { loading = false; options.onLoadingChange?.(loading); @@ -193,16 +213,32 @@ export function useChat(options: UseChatOptions): ChatInstance { assistantId = null; } + /** + * Surface an error to the consumer at most ONCE per turn (see errorReported). A 504 + * send-rejection and a later `error` chunk can both describe the same failed turn; we + * must not fire onError twice for it. + */ + function reportError(message: string) { + if (errorReported) return; + errorReported = true; + options.onError?.(message); + } + /** * Shared handler for a rejected send path (api.sendMessage / api.resume): reset * loading, drop the dangling empty placeholder, and surface the error via onError * (swallowed, matching how the `error` chunk is handled — no re-throw). + * + * NOTE: a 504 (or similar) rejection often means the turn DID start server-side, so a + * later done/delta/error chunk may still arrive for it. Recovery of such a started + * turn happens via the normal reconnect → getConversation re-sync path; here we only + * guard against a DOUBLE onError (reportError) if that later chunk is an `error`. */ function handleSendFailure(err: unknown) { loading = false; options.onLoadingChange?.(loading); removeEmptyAssistantPlaceholder(); - options.onError?.(err instanceof Error ? err.message : String(err)); + reportError(err instanceof Error ? err.message : String(err)); } /** @@ -218,42 +254,84 @@ export function useChat(options: UseChatOptions): ChatInstance { * Also re-checks pending interrupts, which may have been raised during the gap. */ async function handleReconnect() { + if (destroyed) return; if (!conversationId) return; + // Capture the in-flight turn identity BEFORE the await. getConversation reads + // DynamoDB, which is eventually consistent, so this read can (a) resolve LATE — + // after a terminal chunk already resolved the turn on the resubscribed channel — + // and (b) reflect a STALE view whose last row is a previous turn's assistant + // message. Both are guarded against once the read resolves, using this snapshot. + const turnAtStart = assistantId; try { const { messages: history } = await options.api.getConversation(conversationId); + if (destroyed) return; const last = history[history.length - 1]; const turnComplete = !!last && last.role === 'assistant' && !!last.content; - if (turnComplete && assistantId) { - // Turn finished while we were disconnected; the terminal `done` chunk was lost. - // Replace the in-flight assistant bubble with the persisted final text. - assistantText = last.content; - messages = messages.map(m => (m.id === assistantId ? { ...m, content: last.content } : m)); - options.onMessagesChange?.(messages); - assistantId = null; - clearFailsafe(); - loading = false; - options.onLoadingChange?.(loading); - } else if (loading) { - // Turn still running server-side — do NOT clear loading. Wait for the terminal - // chunk on the resubscribed channel, guarded by the bounded failsafe. + // Turn-identity + liveness guard: only act on the persisted read if the turn we + // captured is STILL the in-flight one and we're still loading. If a terminal + // chunk (done/error) resolved the turn while getConversation was in flight, it + // already nulled assistantId (and set the authoritative final text), so we must + // ignore this — possibly stale — DB result rather than clobber a live outcome. + const stillSameInFlightTurn = assistantId !== null && assistantId === turnAtStart && loading; + + if (turnComplete && stillSameInFlightTurn) { + // Adopt the persisted final text ONLY if it extends what we've already + // streamed this turn (nothing streamed yet, or the stream is a prefix of the + // persisted text). If it does NOT extend our stream, `last` is a stale / + // previous-turn assistant row that does not continue the current bubble; + // adopting it would overwrite live text with prior-turn content and halt + // streaming, so we treat the turn as still running and wait for the terminal + // chunk instead. + const extendsStream = assistantText === '' || last.content.startsWith(assistantText); + if (extendsStream) { + // Turn finished while we were disconnected; the terminal `done` chunk was lost. + // Replace the in-flight assistant bubble with the persisted final text. + assistantText = last.content; + messages = messages.map(m => (m.id === assistantId ? { ...m, content: last.content } : m)); + options.onMessagesChange?.(messages); + assistantId = null; + clearFailsafe(); + loading = false; + options.onLoadingChange?.(loading); + } else { + // Snapshot doesn't extend our stream — keep waiting for the terminal chunk, + // guarded by the bounded failsafe. + armFailsafe(); + } + } else if (stillSameInFlightTurn) { + // Turn still running server-side (no final assistant message yet) — do NOT clear + // loading. Wait for the terminal chunk on the resubscribed channel, guarded by + // the bounded failsafe. armFailsafe(); } + // else: the turn was already resolved by a terminal chunk during the await + // (assistantId nulled / loading cleared) — nothing to adopt. // A pending interrupt may have been raised while the socket was down. if (options.api.getPendingInterrupts) { const { interrupts } = await options.api.getPendingInterrupts(conversationId); + if (destroyed) return; if (interrupts.length) options.onInterrupt?.(interrupts); } } catch (err) { - // Re-sync itself failed. Surface it, and keep the spinner honest via the failsafe. - if (loading) armFailsafe(); + if (destroyed) return; + // Re-sync itself failed. Surface it via onError only. We deliberately do NOT also + // arm the failsafe here (NIT): the channel is resubscribed, so a terminal chunk can + // still resolve the turn; arming would fire a second, misleading 'Timed out' error + // ~11min later on top of the error we just surfaced. options.onError?.(err instanceof Error ? err.message : String(err)); } } /** Handle a chunk from the Realtime subscription. */ function handleChunk(chunk: AgentStreamChunk) { + // Liveness: ANY received chunk proves the stream is alive. If the post-reconnect + // failsafe is armed, re-arm it (reset its countdown) so it only ever fires after a + // window of COMPLETE silence — not during a long tool-call/thinking gap or a slow + // post-reconnect stream. Terminal chunks below still clearFailsafe outright. + if (failsafeTimer !== null) armFailsafe(); + options.onChunk?.(chunk); if (chunk.type === 'text-delta' && chunk.text && assistantId) { @@ -267,16 +345,25 @@ export function useChat(options: UseChatOptions): ChatInstance { messages = messages.map(m => m.id === assistantId ? { ...m, content: chunk.text! } : m); options.onMessagesChange?.(messages); } + // Null assistantId so it is a reliable in-flight signal: a later-resolving + // reconnect re-sync must not adopt a (possibly stale) DB snapshot over this + // already-resolved turn (see handleReconnect's turn-identity guard). + assistantId = null; clearFailsafe(); loading = false; options.onLoadingChange?.(loading); } if (chunk.type === 'error') { + // Null assistantId (mirroring done/interrupt) so a later reconnect re-sync treats + // the turn as resolved and won't clobber the bubble with a stale getConversation read. + assistantId = null; clearFailsafe(); loading = false; options.onLoadingChange?.(loading); - options.onError?.(chunk.error ?? 'Unknown error'); + // Report at most once per turn: a prior send-rejection may have already surfaced + // onError for this same (server-started) turn. + reportError(chunk.error ?? 'Unknown error'); } if (chunk.type === 'interrupt' && chunk.interrupts) { @@ -300,20 +387,27 @@ export function useChat(options: UseChatOptions): ChatInstance { async function ensureSubscribed(channelId: string) { if (activeSub) { activeSub.unsubscribe(); activeSub = null; } - // Pass an options object (not a bare handler) so the transport can notify us on - // reconnect — we re-sync authoritative state from the DB in handleReconnect(). - const subscribeOptions: ChatSubscribeOptions = { - onMessage: handleChunk, - onReconnect: () => { void handleReconnect(); }, - }; - - const sub = await options.subscribe(channelId, subscribeOptions); + // Pass a CALLABLE that ALSO carries the ChatSubscribeOptions properties. This keeps + // both shapes of the `subscribe` union valid at runtime: an old adapter that invokes + // arg2 as a bare chunk handler still works (it's a function that dispatches to + // handleChunk), and a new adapter that reads arg2.onMessage/onReconnect/onDisconnect + // still works (those are attached as properties). Cast-free: the intersection type + // is assignable to either member of the ChatChunkHandler | ChatSubscribeOptions union. + const subscribeArg: ChatChunkHandler & ChatSubscribeOptions = Object.assign( + (chunk: AgentStreamChunk) => { handleChunk(chunk); }, + { + onMessage: handleChunk, + onReconnect: () => { void handleReconnect(); }, + }, + ); + + const sub = await options.subscribe(channelId, subscribeArg); try { await sub.established; } catch (err) { console.warn('Subscription failed, retrying with fresh token:', err); sub.unsubscribe(); - const retrySub = await options.subscribe(channelId, subscribeOptions); + const retrySub = await options.subscribe(channelId, subscribeArg); await retrySub.established; activeSub = retrySub; return; @@ -343,6 +437,7 @@ export function useChat(options: UseChatOptions): ChatInstance { assistantText = ''; messages = [...messages, userMsg, aMsg]; options.onMessagesChange?.(messages); + errorReported = false; // fresh turn — allow one onError report loading = true; options.onLoadingChange?.(loading); @@ -374,6 +469,7 @@ export function useChat(options: UseChatOptions): ChatInstance { } assistantText = ''; options.onMessagesChange?.(messages); + errorReported = false; // fresh turn — allow one onError report loading = true; options.onLoadingChange?.(loading); if (!options.api.resume) throw new Error('respondToInterrupt requires api.resume to be configured'); @@ -417,6 +513,7 @@ export function useChat(options: UseChatOptions): ChatInstance { }, destroy() { + destroyed = true; clearFailsafe(); if (activeSub) { activeSub.unsubscribe(); activeSub = null; } }, diff --git a/packages/bb-agent/src/index.test.ts b/packages/bb-agent/src/index.test.ts index dc73896d6..0133ca3a3 100644 --- a/packages/bb-agent/src/index.test.ts +++ b/packages/bb-agent/src/index.test.ts @@ -1185,7 +1185,7 @@ describe('model-factory', () => { // ── useChat ────────────────────────────────────────────────────────────────── import { useChat } from './index.hooks.js'; -import type { AgentStreamChunk, ChatMessage, UseChatOptions } from './index.hooks.js'; +import type { AgentStreamChunk, ChatMessage, UseChatOptions, ChatChunkHandler, ChatSubscribeOptions } from './index.hooks.js'; /** Flush pending microtasks so an async onReconnect handler settles before assertions. */ function flush(): Promise { @@ -1193,10 +1193,17 @@ function flush(): Promise { } /** - * useChat calls `subscribe` with a {@link ChatSubscribeOptions} object. Capture its - * callbacks (onMessage / onReconnect / onDisconnect) cast-free so tests can drive chunks - * and simulate a transport reconnect. Also tolerates the bare-handler form for safety. + * useChat calls `subscribe` with a callable that ALSO carries the {@link ChatSubscribeOptions} + * properties (onMessage / onReconnect / onDisconnect) — a backward-compatible hybrid. Capture + * those callbacks cast-free (via a type guard, not a cast) so tests can drive chunks and + * simulate a transport reconnect. Also tolerates a plain options object and a truly-bare handler. */ +function hasSubscribeOptions(arg: ChatChunkHandler | ChatSubscribeOptions): arg is ChatSubscribeOptions { + // A bare handler is a function with no onMessage property; the options object and the + // callable-with-props hybrid both carry onMessage. + return typeof arg !== 'function' || 'onMessage' in arg; +} + function subscribeCapture() { const cap: { handler?: (chunk: AgentStreamChunk) => void; @@ -1204,12 +1211,12 @@ function subscribeCapture() { disconnect?: (reason: string) => void; } = {}; const subscribe: UseChatOptions['subscribe'] = async (_channelId, handlerOrOptions) => { - if (typeof handlerOrOptions === 'function') { - cap.handler = handlerOrOptions; - } else { + if (hasSubscribeOptions(handlerOrOptions)) { cap.handler = handlerOrOptions.onMessage; cap.reconnect = handlerOrOptions.onReconnect; cap.disconnect = handlerOrOptions.onDisconnect; + } else { + cap.handler = handlerOrOptions; } return { unsubscribe() {}, established: Promise.resolve() }; }; @@ -1484,7 +1491,7 @@ describe('useChat', () => { chat.destroy(); }); - test('bounded failsafe clears loading if no terminal chunk arrives after reconnect', async (t) => { + test('bounded failsafe: chunks re-arm it; it fires only after a fully silent window', async (t) => { t.mock.timers.enable({ apis: ['setTimeout'] }); let errorReceived: string | undefined; const { cap, subscribe } = subscribeCapture(); @@ -1505,11 +1512,114 @@ describe('useChat', () => { await flush(); assert.strictEqual(chat.isLoading(), true, 'still loading right after reconnect (turn running)'); - // No terminal chunk ever arrives — advance past the failsafe window. - t.mock.timers.tick(60_000); + // A >30s tool-call/thinking gap (past the OLD 30s window) must NOT trip the failsafe. + t.mock.timers.tick(120_000); + assert.strictEqual(chat.isLoading(), true, 'a long silent gap under the new window must not fire'); + assert.strictEqual(errorReceived, undefined, 'no premature timeout during a normal long gap'); + + // A text-delta proves the stream is alive and RE-ARMS the window from now. + cap.handler!({ type: 'text-delta', text: 'still working…' }); + // Advance almost a full window since that delta — still alive, still no fire. + t.mock.timers.tick(600_000); + assert.strictEqual(chat.isLoading(), true, 'the delta re-armed the window, so it has not elapsed'); + assert.strictEqual(errorReceived, undefined); + + // Now go fully silent past the whole window — the failsafe finally fires. + t.mock.timers.tick(660_001); + assert.strictEqual(chat.isLoading(), false, 'failsafe clears loading after a fully silent window'); + assert.ok(errorReceived, 'failsafe surfaces an error'); + chat.destroy(); + }); + + test('reconnect does not clobber final text delivered by a live done that arrived before getConversation resolved', async () => { + const { cap, subscribe } = subscribeCapture(); + // Gate getConversation so the reconnect re-sync resolves AFTER a live `done` chunk. + let resolveGet!: (v: { messages: { role: string; content: string }[] }) => void; + const getGate = new Promise<{ messages: { role: string; content: string }[] }>((r) => { resolveGet = r; }); + + const chat = useChat({ + api: { + sendMessage: async () => {}, + createConversation: async () => ({ conversationId: 'conv-1' }), + getConversation: () => getGate, + }, + subscribe, + }); + + await chat.sendMessage('hello'); + // Reconnect kicks off getConversation, which stays pending on the gate. + cap.reconnect!(); + await flush(); + + // The live terminal `done` arrives on the resubscribed channel BEFORE the DB read resolves. + cap.handler!({ type: 'done', text: 'LIVE final answer' }); + assert.strictEqual(chat.isLoading(), false, 'the live done cleared loading'); + + // The late getConversation now resolves with a DIFFERENT (stale/eventually-consistent) view. + resolveGet({ messages: [ + { role: 'user', content: 'hello' }, + { role: 'assistant', content: 'STALE db answer' }, + ] }); + await flush(); + + const assistant = chat.getMessages().find(m => m.role === 'assistant'); + assert.strictEqual(assistant!.content, 'LIVE final answer', 'the live done text is preserved; the late DB read is ignored'); + assert.strictEqual(chat.isLoading(), false, 'loading stays cleared'); + chat.destroy(); + }); + + test('reconnect ignores a stale/previous-turn getConversation snapshot', async () => { + const { cap, subscribe } = subscribeCapture(); + + const chat = useChat({ + api: { + sendMessage: async () => {}, + createConversation: async () => ({ conversationId: 'conv-1' }), + // The turn is still running, but the DB read's last row is a PRIOR turn's + // assistant message that does NOT extend what we've streamed this turn. + getConversation: async () => ({ messages: [ + { role: 'user', content: 'previous question' }, + { role: 'assistant', content: 'answer to a PRIOR turn' }, + { role: 'user', content: 'hello' }, + { role: 'assistant', content: 'answer to a PRIOR turn' }, + ] }), + }, + subscribe, + }); + + await chat.sendMessage('hello'); + // Live text for the CURRENT turn. + cap.handler!({ type: 'text-delta', text: 'live streaming answer' }); + cap.reconnect!(); + await flush(); + + assert.strictEqual(chat.isLoading(), true, 'a stale snapshot must not resolve the still-running turn'); + const assistant = chat.getMessages().find(m => m.role === 'assistant' && m.content !== ''); + assert.strictEqual(assistant!.content, 'live streaming answer', 'in-flight bubble is not overwritten by prior-turn text'); + chat.destroy(); + }); + + test('a send rejection then a later error chunk reports onError exactly once', async () => { + const errors: string[] = []; + const { cap, subscribe } = subscribeCapture(); + + const chat = useChat({ + api: { + sendMessage: async () => { throw new Error('504 Gateway Timeout'); }, + createConversation: async () => ({ conversationId: 'conv-1' }), + getConversation: async () => ({ messages: [] }), + }, + subscribe, + onError: (e) => { errors.push(e); }, + }); + + await chat.sendMessage('hello'); + // The 504 rejection already surfaced onError. The turn DID start server-side, so a + // later `error` chunk arrives for the same turn — it must NOT double-report. + cap.handler!({ type: 'error', error: 'server-side failure' }); - assert.strictEqual(chat.isLoading(), false, 'failsafe should clear loading'); - assert.ok(errorReceived, 'failsafe should surface an error'); + assert.strictEqual(errors.length, 1, 'onError fires exactly once for the same failed turn'); + assert.match(errors[0], /504/, 'the first (send-rejection) error is the one surfaced'); chat.destroy(); }); }); From 8dd65352f1281f44658dffef12b3cf60a31b19ad Mon Sep 17 00:00:00 2001 From: Michael Sober Date: Mon, 7 Sep 2026 16:10:37 +0000 Subject: [PATCH 3/3] fix(bb-agent): useChat must pass an options object so reconnect re-sync reaches the transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address PR #501 review (BLOCKER + HIGH + MEDIUM): useChat built its subscribe arg as a callable-with-props (Object.assign(fn, {onMessage,onReconnect})). Both bb-realtime middlewares resolve subscribe with `typeof arg === 'function'` FIRST and treat any function as a bare handler, never reading its properties — so onReconnect/onDisconnect were silently dropped and the DB re-sync + bounded failsafe (this PR's headline stuck-spinner fix) never fired on the real transport for the documented adapter usage. - index.hooks.ts: pass a plain ChatSubscribeOptions object; the function-first middleware now hits its object branch and reads onReconnect. subscribe param type still accepts a bare handler OR options (back-compat). - index.test.ts: the subscribeCapture shim was options-first (read props off a function), the reverse of the real middleware, so every reconnect test passed vacuously. Rewrote it to mirror the real function-first precedence — a function is always a bare handler — so a regression to a callable now fails the tests. - README.md: both useChat adapter snippets now forward the subscribe argument (`sub`) verbatim to channel.subscribe, with a note that a bare handler drops the reconnect callbacks. bb-agent 114/114, build clean, biome 0 errors, cast-free. --- packages/bb-agent/README.md | 14 ++++++++++---- packages/bb-agent/src/index.hooks.ts | 23 ++++++++++------------- packages/bb-agent/src/index.test.ts | 9 ++++++--- 3 files changed, 26 insertions(+), 20 deletions(-) diff --git a/packages/bb-agent/README.md b/packages/bb-agent/README.md index 2ed229025..f5ea7caef 100644 --- a/packages/bb-agent/README.md +++ b/packages/bb-agent/README.md @@ -724,9 +724,12 @@ const chat = useChat({ getConversation: (id) => api.getConversation(id), resume: (chId, responses, convId) => api.resume(chId, responses, convId), }, - subscribe: async (channelId, handler) => { + // Forward the subscribe argument (`sub`) VERBATIM to the channel — it is an + // options object carrying onMessage/onReconnect/onDisconnect. Passing only a + // bare handler would drop the reconnect callbacks the transport needs. + subscribe: async (channelId, sub) => { const channel = await api.getChannel(channelId); - return channel.subscribe(handler); + return channel.subscribe(sub); }, onMessagesChange: (msgs) => renderMessages(msgs), onLoadingChange: (loading) => updateSpinner(loading), @@ -788,9 +791,12 @@ const chat = useChat({ createConversation: () => api.createConversation(userId), getConversation: (id) => api.getConversation(id), }, - subscribe: async (channelId, handler) => { + // Forward the subscribe argument (`sub`) VERBATIM to the channel — it is an + // options object carrying onMessage/onReconnect/onDisconnect. Passing only a + // bare handler would drop the reconnect callbacks the transport needs. + subscribe: async (channelId, sub) => { const channel = await api.getChannel(channelId); - return channel.subscribe(handler); + return channel.subscribe(sub); }, onMessagesChange: (msgs) => renderMessages(msgs), onLoadingChange: (loading) => updateSpinner(loading), diff --git a/packages/bb-agent/src/index.hooks.ts b/packages/bb-agent/src/index.hooks.ts index 9069e33e1..957f49b81 100644 --- a/packages/bb-agent/src/index.hooks.ts +++ b/packages/bb-agent/src/index.hooks.ts @@ -387,19 +387,16 @@ export function useChat(options: UseChatOptions): ChatInstance { async function ensureSubscribed(channelId: string) { if (activeSub) { activeSub.unsubscribe(); activeSub = null; } - // Pass a CALLABLE that ALSO carries the ChatSubscribeOptions properties. This keeps - // both shapes of the `subscribe` union valid at runtime: an old adapter that invokes - // arg2 as a bare chunk handler still works (it's a function that dispatches to - // handleChunk), and a new adapter that reads arg2.onMessage/onReconnect/onDisconnect - // still works (those are attached as properties). Cast-free: the intersection type - // is assignable to either member of the ChatChunkHandler | ChatSubscribeOptions union. - const subscribeArg: ChatChunkHandler & ChatSubscribeOptions = Object.assign( - (chunk: AgentStreamChunk) => { handleChunk(chunk); }, - { - onMessage: handleChunk, - onReconnect: () => { void handleReconnect(); }, - }, - ); + // Pass a plain options object (NOT a callable-with-props). Both bb-realtime + // middlewares resolve subscribe with `typeof handlerOrOptions === 'function'` + // FIRST — a function is treated as a bare handler and its onMessage/onReconnect/ + // onDisconnect properties are never read. A hybrid callable would therefore + // silently drop onReconnect, making the reconnect re-sync + failsafe dead on the + // real transport. The options object hits the transport's object branch. + const subscribeArg: ChatSubscribeOptions = { + onMessage: handleChunk, + onReconnect: () => { void handleReconnect(); }, + }; const sub = await options.subscribe(channelId, subscribeArg); try { diff --git a/packages/bb-agent/src/index.test.ts b/packages/bb-agent/src/index.test.ts index 0133ca3a3..254177bf1 100644 --- a/packages/bb-agent/src/index.test.ts +++ b/packages/bb-agent/src/index.test.ts @@ -1199,9 +1199,12 @@ function flush(): Promise { * simulate a transport reconnect. Also tolerates a plain options object and a truly-bare handler. */ function hasSubscribeOptions(arg: ChatChunkHandler | ChatSubscribeOptions): arg is ChatSubscribeOptions { - // A bare handler is a function with no onMessage property; the options object and the - // callable-with-props hybrid both carry onMessage. - return typeof arg !== 'function' || 'onMessage' in arg; + // Mirror the REAL bb-realtime middleware precedence: it branches on + // `typeof arg === 'function'` FIRST and treats any function as a bare handler, + // never reading properties off it. So onReconnect/onDisconnect are only honored + // when arg is a non-function options object. Testing it any other way would let a + // regression to a callable-with-props pass vacuously. + return typeof arg !== 'function'; } function subscribeCapture() {