fix(bb-realtime,bb-agent): refresh channel tokens on reconnect (survive 1h/2h TTLs)

.changeset/realtime-token-refresh.md

@@ -0,0 +1,39 @@
+---
+"@aws-blocks/bb-realtime": minor
+---
+
+feat(bb-realtime): refresh channel/connect tokens on reconnect so subscriptions outlive token TTLs
+
+Adds an optional `refresh` callback to `SubscribeOptions`:
+
+```ts
+refresh?: () => Promise<RealtimeChannelDescriptor>;
+```
+
+A reconnect opens a *new* WebSocket, which means API Gateway re-checks the
+connect token (carried in the socket URL, validated at `$connect`, ~2h TTL) and
+the server re-checks the channel token on resubscribe (~1h TTL, per `utils.ts`
+`mintChannelToken`'s 3600s default). Until now both the AWS and mock middlewares
+replayed the *original* stored `wsUrl` + channel token on every reconnect, so a
+reconnect more than ~1h after the descriptor was minted failed (the channel
+token had expired and the resubscribe was rejected), and more than ~2h after
+failed to even open the socket (`$connect` 403). Token minting is server-only
+(it needs the signing secret), so the client cannot re-sign locally — it must
+re-call the server method that produced the descriptor.
+
+When `refresh` is provided, both middlewares now call it **before** opening the
+reconnect socket (never on the initial subscribe), then open with the fresh
+connect token in the URL and resubscribe with the fresh channel token. The
+refresh-before-open ordering is required because the connect token lives in the
+socket URL and is validated at `$connect`, so it must be fresh at construction
+time. If `refresh` rejects, the middleware does not crash: it surfaces the
+failure via the existing `onDisconnect('error')` path and falls back to the
+normal exponential-backoff reconnect so a later attempt can retry.
+
+Fully backward compatible: with no `refresh` callback, a reconnect replays the
+stored `wsUrl` + token exactly as before, and the initial (non-reconnect) open
+stays synchronous and unchanged.
+
+This is a `minor` bump. `@aws-blocks/bb-realtime` is pre-1.0, where `minor` is
+this repo's signal for an API addition; the new option is optional and additive,
+and existing behavior is unchanged when it is omitted.

.changeset/usechat-token-refresh.md

@@ -0,0 +1,7 @@
+---
+"@aws-blocks/bb-agent": minor
+---
+
+useChat now forwards an optional consumer-supplied `refresh` callback to the Realtime subscription so reconnects mint fresh tokens and survive past the channel/connect token TTLs on long turns.
+
+useChat only holds the channelId plus the consumer's `subscribe` adapter; the channel descriptor is minted inside that adapter (via `api.agentGetChannel`), which useChat cannot reach — so it cannot self-mint. Instead, `UseChatOptions` accepts an optional `refresh?: () => Promise<ChatChannelDescriptor>` (typically `() => api.agentGetChannel(conversationId)`) that useChat forwards to the subscription (as `refresh` on the `ChatSubscribeOptions` object). The transport calls it before each reconnect (never on the initial subscribe) to obtain a freshly-minted connect + channel token, so a subscription can outlive the channel (~1h) and connect (~2h) token TTLs. Fully backward compatible: when omitted, a reconnect replays the original tokens exactly as before.

packages/bb-agent/src/index.hooks.ts

@@ -29,6 +29,20 @@
 /** Handler invoked for each streaming chunk delivered over the Realtime channel. */
 export type ChatChunkHandler = (chunk: AgentStreamChunk) => void;
 
+/**
+ * Minimal structural mirror of bb-realtime's `RealtimeChannelDescriptor` — the wire
+ * format an app's `subscribe` adapter hands to `channel.subscribe(...)` to hydrate a
+ * live channel. Mirrored here rather than imported, matching how {@link ChatSubscribeOptions}
+ * mirrors bb-realtime's `SubscribeOptions` structurally, so the client hooks take no hard
+ * type dependency on bb-realtime. Its fields match the descriptor exactly, so a `refresh`
+ * typed against it stays assignable to bb-realtime's `SubscribeOptions.refresh`.
+ */
+export interface ChatChannelDescriptor {
+	__blocks: 'realtime/channel';
+	channel: string;
+	[key: string]: unknown;
+}
+
 /**
  * Options form accepted by {@link UseChatOptions.subscribe}.
  *
@@ -52,6 +66,15 @@
 	 * resubscribed. useChat uses this to re-sync from the DB (see {@link UseChatOptions.subscribe}).
 	 */
 	onReconnect?: () => void;
+	/**
+	 * Called before each reconnect to obtain a freshly-minted channel descriptor (new
+	 * connect + channel token) so the subscription can outlive the token TTLs (channel
+	 * ~1h / connect ~2h). Mirrors bb-realtime's `SubscribeOptions.refresh`. useChat
+	 * forwards {@link UseChatOptions.refresh} here verbatim; the transport calls it on
+	 * reconnect only (never on the initial subscribe) and simply does not use it when
+	 * undefined.
+	 */
+	refresh?: () => Promise<ChatChannelDescriptor>;
 }
 
 /** Options for creating a chat instance. */
@@ -75,6 +98,17 @@
 	 * - established: Promise that resolves when the WS subscription is confirmed
 	 */
 	subscribe: (channelId: string, handlerOrOptions: ChatChunkHandler | ChatSubscribeOptions) => Promise<{ unsubscribe(): void; established: Promise<void> }>;
+	/**
+	 * Optional consumer-supplied callback to re-mint a fresh channel descriptor when the
+	 * Realtime transport reconnects. useChat only holds the channelId (== conversationId)
+	 * plus your `subscribe` adapter; the channel descriptor is minted INSIDE that adapter
+	 * (via `api.agentGetChannel`), which useChat cannot reach — so it cannot self-mint.
+	 * Provide this and useChat forwards it to the subscription (as `refresh`) so long turns
+	 * survive the channel (~1h) / connect (~2h) token TTLs: a reconnect mints fresh tokens
+	 * instead of replaying expired ones. Typically `() => api.agentGetChannel(conversationId)`.
+	 * When omitted, a reconnect replays the original tokens (fine for short turns).
+	 */
+	refresh?: () => Promise<ChatChannelDescriptor>;
 	/** Called whenever the message list changes. */
 	onMessagesChange?: (messages: ChatMessage[]) => void;
 	/** Called whenever loading state changes. */
@@ -138,10 +172,14 @@
  *   },
  *   subscribe: async (channelId, sub) => {
  *     const result = await api.agentGetChannel(channelId);
- *     // `sub` is a ChatSubscribeOptions object (onMessage/onReconnect/onDisconnect);
- *     // channel.subscribe accepts it directly and wires reconnect handling for us.
+ *     // `sub` is a ChatSubscribeOptions object (onMessage/onReconnect/onDisconnect/refresh);
+ *     // channel.subscribe accepts it directly and wires reconnect handling — including
+ *     // calling sub.refresh to re-mint fresh tokens before each reconnect — for us.
  *     return result.channel.subscribe(sub);
  *   },
+ *   // Re-mint a fresh channel descriptor on reconnect so long turns outlive the channel
+ *   // (~1h) / connect (~2h) token TTLs. useChat forwards this to the subscription as sub.refresh.
+ *   refresh: () => api.agentGetChannel(conversationId),
  *   onMessagesChange: (msgs) => renderMessages(msgs),
  *   onLoadingChange: (loading) => updateSpinner(loading),
  * });
@@ -342,7 +380,7 @@
 
 		if (chunk.type === 'done') {
 			if (chunk.text && assistantId) {
 				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
@@ -390,12 +428,15 @@
 		// 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.
+		// onDisconnect/refresh properties are never read. A hybrid callable would therefore
+		// silently drop them, making the reconnect re-sync + failsafe + token refresh dead
+		// on the real transport. The options object hits the transport's object branch.
 		const subscribeArg: ChatSubscribeOptions = {
 			onMessage: handleChunk,
 			onReconnect: () => { void handleReconnect(); },
+			// Forward the consumer-supplied re-mint callback (if any). When undefined the
+			// transport simply replays the original tokens on reconnect (back-compat).
+			refresh: options.refresh,
 		};
 
 		const sub = await options.subscribe(channelId, subscribeArg);

packages/bb-agent/src/index.test.ts

@@ -1185,7 +1185,7 @@
 // ── useChat ──────────────────────────────────────────────────────────────────
 
 import { useChat } from './index.hooks.js';
-import type { AgentStreamChunk, ChatMessage, UseChatOptions, ChatChunkHandler, ChatSubscribeOptions } from './index.hooks.js';
+import type { AgentStreamChunk, ChatMessage, UseChatOptions, ChatChunkHandler, ChatSubscribeOptions, ChatChannelDescriptor } from './index.hooks.js';
 
 /** Flush pending microtasks so an async onReconnect handler settles before assertions. */
 function flush(): Promise<void> {
@@ -1212,12 +1212,14 @@
 		handler?: (chunk: AgentStreamChunk) => void;
 		reconnect?: () => void;
 		disconnect?: (reason: string) => void;
+		refresh?: () => Promise<ChatChannelDescriptor>;
 	} = {};
 	const subscribe: UseChatOptions['subscribe'] = async (_channelId, handlerOrOptions) => {
 		if (hasSubscribeOptions(handlerOrOptions)) {
 			cap.handler = handlerOrOptions.onMessage;
 			cap.reconnect = handlerOrOptions.onReconnect;
 			cap.disconnect = handlerOrOptions.onDisconnect;
+			cap.refresh = handlerOrOptions.refresh;
 		} else {
 			cap.handler = handlerOrOptions;
 		}
@@ -1248,7 +1250,7 @@
 
 		await chat.sendMessage('hello');
 		// Simulate error chunk from server
 		chunkHandler!({ type: 'error', error: 'model throttled' });
 
 		assert.strictEqual(errorReceived, 'model throttled');
 		assert.strictEqual(loadingStates.at(-1), false, 'loading should be false after error');
@@ -1274,11 +1276,11 @@
 		});
 
 		await chat.sendMessage('hello');
 		chunkHandler!({ type: 'interrupt', interrupts: [{ id: 'int-1', name: 'approve:deleteRecords', reason: { tool: 'deleteRecords' } }] });
 
 		assert.ok(interruptsReceived, 'onInterrupt should be called');
 		assert.strictEqual(interruptsReceived!.length, 1);
 		assert.strictEqual(interruptsReceived![0].name, 'approve:deleteRecords');
 		assert.strictEqual(loadingStates.at(-1), false, 'loading should be false after interrupt');
 	});
 
@@ -1301,7 +1303,7 @@
 		});
 
 		await chat.sendMessage('hello');
 		chunkHandler!({ type: 'interrupt', interrupts: [{ id: 'int-1', name: 'approve:delete' }] });
 		await chat.respondToInterrupt([{ interruptId: 'int-1', approved: true }]);
 
 		assert.ok(resumeCalled, 'api.resume should be called');
@@ -1327,7 +1329,7 @@
 		});
 
 		await chat.sendMessage('hello');
 		chunkHandler!({ type: 'interrupt', interrupts: [{ id: 'int-1', name: 'approve:delete' }] });
 		await assert.rejects(() => chat.respondToInterrupt([{ interruptId: 'int-1', approved: true }]), /api.resume/);
 	});
 
@@ -1352,7 +1354,7 @@
 		// At this point there's a user message + empty assistant placeholder
 		assert.ok(lastMessages.some(m => m.role === 'assistant' && m.content === ''), 'should have empty placeholder');
 		// Interrupt arrives — placeholder should be removed
 		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');
 	});
 
@@ -1555,7 +1557,7 @@
 		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.
@@ -1625,6 +1627,36 @@
 		assert.match(errors[0], /504/, 'the first (send-rejection) error is the one surfaced');
 		chat.destroy();
 	});
+
+	test('useChat forwards a consumer-supplied refresh to the subscription options', async () => {
+		// The consumer owns api.agentGetChannel, so it supplies the actual mint; useChat only
+		// forwards it. A fresh descriptor is what the transport uses to re-mint tokens on reconnect.
+		const descriptor: ChatChannelDescriptor = { __blocks: 'realtime/channel', channel: 'conv-1' };
+		let refreshCalls = 0;
+		const mockRefresh = async (): Promise<ChatChannelDescriptor> => { refreshCalls++; return descriptor; };
+		const { cap, subscribe } = subscribeCapture();
+
+		const chat = useChat({
+			api: {
+				sendMessage: async () => {},
+				createConversation: async () => ({ conversationId: 'conv-1' }),
+				getConversation: async () => ({ messages: [] }),
+			},
+			subscribe,
+			refresh: mockRefresh,
+		});
+
+		// Sending triggers ensureSubscribed, which builds subscribeArg with refresh attached.
+		await chat.sendMessage('hello');
+
+		assert.strictEqual(typeof cap.refresh, 'function', 'refresh should be forwarded to the subscription options');
+		assert.strictEqual(cap.refresh, mockRefresh, 'the exact consumer-supplied fn is forwarded verbatim');
+		// Invoking the forwarded fn calls the consumer mint and yields its descriptor.
+		const result = await cap.refresh!();
+		assert.strictEqual(refreshCalls, 1, 'the forwarded refresh invokes the consumer mint');
+		assert.strictEqual(result, descriptor, 'and returns the freshly-minted descriptor');
+		chat.destroy();
+	});
 });
 
 describe('checkModelHealth', () => {
@@ -1697,7 +1729,7 @@
 			assert.strictEqual(healthy, false, 'non-JSON 200 response should be treated as unhealthy, not throw');
 			const warned = warnings.find(w => w.meta && 'bodySnippet' in w.meta);
 			assert.ok(warned, 'should warn about the non-JSON body');
 			assert.match(warned!.meta.bodySnippet, /<html>/, 'warning should include a snippet of the offending body');
 		} finally {
 			server.close();
 		}

packages/bb-realtime/API.md

@@ -72,6 +72,7 @@ export interface SubscribeOptions<T = unknown> {
     onDisconnect?: (reason: DisconnectReason) => void;
     onMessage: (message: T) => void;
     onReconnect?: () => void;
+    refresh?: () => Promise<RealtimeChannelDescriptor>;
 }
 
 // (No @packageDocumentation comment for this package)