fix(bb-realtime,bb-agent): refresh channel tokens on reconnect (survive 1h/2h TTLs) - #503
Open
soberm wants to merge 7 commits into
Open
fix(bb-realtime,bb-agent): refresh channel tokens on reconnect (survive 1h/2h TTLs)#503soberm wants to merge 7 commits into
soberm wants to merge 7 commits into
Conversation
🦋 Changeset detectedLatest commit: ef4775e The changes in this PR will be included in the next version bump. This PR includes changesets to release 3 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
soberm
added a commit
that referenced
this pull request
Sep 7, 2026
…no zombie socket on reconnect) Address PR #503 review: the awaited refresh() continuation reopened a socket (and, on AWS, leaked a keep-alive interval) if unsubscribe/__resetConnectionsForTest landed during the await — reintroducing the hang PR1 fixed. - aws-middleware: add a tornDown flag to Connection, set it in every terminal teardown path (reset, last-channel unsubscribe, scheduleReconnect give-up + channel-less, terminal 1000/1005 onclose). The openSocket refresh .then now re-checks tornDown / pool-identity / subscriptions.size===0 before applyFreshDescriptor AND again before constructSocket. - mock-middleware: the doConnect refresh .then re-fetches the connection and bails on tornDown / missing / zero-subscription before applying the fresh token and opening — prevents getOrCreateConnection from resurrecting a torn connection (which re-armed timers and hung node --test). - applyFreshDescriptor: only re-key the pool when the target wsUrl is free or is this same connection — never clobber a colliding live entry. - types.ts JSDoc: note that a refresh failure surfaces an extra 'error' disconnect, and that refresh is connection-level (re-mints only its own channel token on a multiplexed connection). Tests: unsubscribe-during-pending-refresh opens no socket (aws + mock), persistently-rejecting refresh stops at MAX_RECONNECT. 96/96 pass, exit 0.
soberm
added a commit
that referenced
this pull request
Sep 7, 2026
… transport honors refresh/onReconnect Address PR #503 review (HIGH): useChat built its subscribe arg as a callable-with-props (Object.assign(fn, {onMessage,onReconnect,refresh})). Both realtime middlewares resolve subscribe with `typeof arg === 'function' ? handler : options` (function-FIRST), so the hybrid was treated as a bare handler and onMessage/onReconnect/onDisconnect/refresh were all silently dropped — making PR3's token refresh AND PR2's DB re-sync inert on the real transport. - useChat now passes a plain ChatSubscribeOptions object; the function-first middleware reads onReconnect/refresh correctly. subscribe param type still accepts a bare handler OR options (back-compat unchanged). - Test shim (index.test.ts) rewritten to mirror the real function-first precedence so a regression to a callable can no longer pass vacuously; added a test asserting useChat passes a NON-function options object with refresh forwarded. - aws-middleware: applyFreshDescriptor now returns a boolean; a refresh() result that fails the descriptor guard is treated like a refresh failure (onDisconnect('error') + backoff) instead of reopening with stale tokens. Added a malformed-descriptor test. - Fix the useChat JSDoc @example so conversationId is in scope. bb-agent 116/116, bb-realtime 97/97, biome 0 errors, cast-free. API.md unchanged (applyFreshDescriptor is private).
soberm
added a commit
that referenced
this pull request
Sep 7, 2026
…(PR #503 MEDIUM) The comprehensive app is the e2e customer-DX surface, so the token-refresh feature must be exercised there. Previously its useChat adapter forwarded the arg but wired no refresh, so the headline capability had no agent-path e2e. - aws-blocks/index.ts: add agentGetRawDescriptor(channelId), mirroring realtimeGetRawDescriptor but sourced from the agent's chunks channel (agent.getChannel -> realtime.getChannel('chunks', id)) so the fresh descriptor targets the same channel useChat subscribes to. - src/index.ts: wire refresh: () => agentGetRawDescriptor(conversationId) into the useChat({...}) call (re-adding __blocks + channel to reconstruct the descriptor, mirroring the realtime e2e). - test/agent.test.ts: add 'reconnect on the agent chunks channel invokes refresh and resubscribes' (force mid-stream close, assert onReconnect + refresh called) and 'agentGetRawDescriptor returns a fresh, well-formed chunks-channel descriptor' (fresh token each call). Both pass locally. Note: like the realtime refresh e2e, refresh returns the bare conversationId as channel (the full namespaced path is server-internal), so these prove the refresh path is invoked + reconnect works; literal TTL-expiry token application can't be waited out in a test.
…g turns survive TTLs A reconnect opens a NEW WebSocket, which re-checks the connect token (2h) at $connect and the channel token (1h) at subscribe. Both middlewares previously replayed the tokens captured at subscribe time, so any reconnect >1h after mint failed (channel token expired -> subscribe rejected) and >2h failed to open the socket ($connect 403). This broke long-running subscriptions -- exactly the 8h AgentCore turn case -- for any bb-realtime consumer. bb-realtime: - Add optional SubscribeOptions.refresh: () => Promise<RealtimeChannelDescriptor>. - On reconnect (before opening the new socket), if refresh is provided, await it for a freshly-minted descriptor and open with the fresh wsUrl (fresh connect token) + resubscribe with the fresh channel token. Restructured openSocket so a reconnect awaits refresh first; initial connect is unchanged. Falls back to replaying stored tokens when refresh is absent (backward compatible). - Mirror the wiring in the mock so local dev matches production. - Fix a stale comment that described channel tokens as ~2h (actual TTL ~1h). bb-agent: - useChat forwards an optional consumer-supplied refresh callback through to the Realtime subscription (typically () => api.agentGetChannel(conversationId)), so agent chats re-mint tokens on reconnect. useChat cannot mint itself -- the descriptor is produced by the consumer's subscribe adapter -- so it relays the callback (structural ChatChannelDescriptor type, no new cross-package dep). Tokens are minted server-side (HMAC secret), so the client re-calls the server method that established the channel; the middleware invokes the app-supplied callback rather than minting locally. Tests: bb-realtime 93/93 (fresh-descriptor reconnect, back-compat replay, refresh rejection fallback, mock refresh); bb-agent 115/115 (refresh forwarded to options). Changesets: @aws-blocks/bb-realtime minor, @aws-blocks/bb-agent minor.
…no zombie socket on reconnect) Address PR #503 review: the awaited refresh() continuation reopened a socket (and, on AWS, leaked a keep-alive interval) if unsubscribe/__resetConnectionsForTest landed during the await — reintroducing the hang PR1 fixed. - aws-middleware: add a tornDown flag to Connection, set it in every terminal teardown path (reset, last-channel unsubscribe, scheduleReconnect give-up + channel-less, terminal 1000/1005 onclose). The openSocket refresh .then now re-checks tornDown / pool-identity / subscriptions.size===0 before applyFreshDescriptor AND again before constructSocket. - mock-middleware: the doConnect refresh .then re-fetches the connection and bails on tornDown / missing / zero-subscription before applying the fresh token and opening — prevents getOrCreateConnection from resurrecting a torn connection (which re-armed timers and hung node --test). - applyFreshDescriptor: only re-key the pool when the target wsUrl is free or is this same connection — never clobber a colliding live entry. - types.ts JSDoc: note that a refresh failure surfaces an extra 'error' disconnect, and that refresh is connection-level (re-mints only its own channel token on a multiplexed connection). Tests: unsubscribe-during-pending-refresh opens no socket (aws + mock), persistently-rejecting refresh stops at MAX_RECONNECT. 96/96 pass, exit 0.
Adds an over-the-wire e2e that subscribes with a refresh callback, forces a mid-stream socket close, and asserts (a) refresh was invoked on the reconnect path to obtain fresh credentials and (b) a message published after the drop is delivered on the reconnected socket — exercising the PR3 token-refresh path end-to-end (literal TTL expiry is out of scope for a test). The refresh callback re-mints via api.realtimeGetRawDescriptor. Cleans up in a finally block. Passes locally against the mock middleware.
… transport honors refresh/onReconnect Address PR #503 review (HIGH): useChat built its subscribe arg as a callable-with-props (Object.assign(fn, {onMessage,onReconnect,refresh})). Both realtime middlewares resolve subscribe with `typeof arg === 'function' ? handler : options` (function-FIRST), so the hybrid was treated as a bare handler and onMessage/onReconnect/onDisconnect/refresh were all silently dropped — making PR3's token refresh AND PR2's DB re-sync inert on the real transport. - useChat now passes a plain ChatSubscribeOptions object; the function-first middleware reads onReconnect/refresh correctly. subscribe param type still accepts a bare handler OR options (back-compat unchanged). - Test shim (index.test.ts) rewritten to mirror the real function-first precedence so a regression to a callable can no longer pass vacuously; added a test asserting useChat passes a NON-function options object with refresh forwarded. - aws-middleware: applyFreshDescriptor now returns a boolean; a refresh() result that fails the descriptor guard is treated like a refresh failure (onDisconnect('error') + backoff) instead of reopening with stale tokens. Added a malformed-descriptor test. - Fix the useChat JSDoc @example so conversationId is in scope. bb-agent 116/116, bb-realtime 97/97, biome 0 errors, cast-free. API.md unchanged (applyFreshDescriptor is private).
…(PR #503 MEDIUM) The comprehensive app is the e2e customer-DX surface, so the token-refresh feature must be exercised there. Previously its useChat adapter forwarded the arg but wired no refresh, so the headline capability had no agent-path e2e. - aws-blocks/index.ts: add agentGetRawDescriptor(channelId), mirroring realtimeGetRawDescriptor but sourced from the agent's chunks channel (agent.getChannel -> realtime.getChannel('chunks', id)) so the fresh descriptor targets the same channel useChat subscribes to. - src/index.ts: wire refresh: () => agentGetRawDescriptor(conversationId) into the useChat({...}) call (re-adding __blocks + channel to reconstruct the descriptor, mirroring the realtime e2e). - test/agent.test.ts: add 'reconnect on the agent chunks channel invokes refresh and resubscribes' (force mid-stream close, assert onReconnect + refresh called) and 'agentGetRawDescriptor returns a fresh, well-formed chunks-channel descriptor' (fresh token each call). Both pass locally. Note: like the realtime refresh e2e, refresh returns the bare conversationId as channel (the full namespaced path is server-internal), so these prove the refresh path is invoked + reconnect works; literal TTL-expiry token application can't be waited out in a test.
…round-trip Same root cause as the realtime reconnect test: the mock fires onReconnect synchronously on frame-send, but the real AWS transport takes a full round-trip (backoff + $connect + Lambda cold start + Secrets Manager + token validation + DynamoDB + PostToConnection), which exceeds the mock-tuned 15s. Widen the refresh-test and agent-channel reconnect deadlines to 60s and post-reconnect delivery windows to 30s so E2E Sandbox/Production reflect a realistic AWS cold-start reconnect budget. Local (mock) stays fast.
The Realtime describe carried { timeout: 60_000 }, which budgets the WHOLE
suite; the slow reconnect tests consumed it and CANCELLED sibling tests
('test did not finish before its parent'). Move the long budget onto the
individual reconnect/refresh/agent reconnect tests (per-test { timeout: 120_000 })
and drop the describe-level cap so the rest of the Realtime suite runs normally.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Issue #, if available: internal — useChat WebSocket timeout handling for long agent turns
A reconnect opens a new WebSocket, which re-checks the connect token (2h) at
$connectand the channel token (1h) atsubscribe. Both middlewares previously replayed the tokens captured at subscribe time, so any reconnect >1h after mint failed (channel token expired → subscribe rejected) and >2h failed to even open the socket ($connect403). This breaks any long-running bb-realtime subscription — exactly the 8h AgentCore turn case — the moment PR1's auto-reconnect fires past the token TTLs.PR 3 of 3 — stacked on #501. #497 made production reconnect; #501 made useChat re-sync content on reconnect; this PR makes the reconnect's credentials fresh so it actually authorizes past the 1h/2h boundaries.
Changes
bb-realtime
SubscribeOptions.refresh: () => Promise<RealtimeChannelDescriptor>.refreshis provided the middleware awaits a freshly-minted descriptor and opens with the freshwsUrl(fresh connect token) then resubscribes with the fresh channel token.openSocketrestructured so a reconnect awaitsrefreshfirst; initial connect unchanged. Falls back to replaying stored tokens whenrefreshis absent (backward compatible).bb-agent
useChatforwards an optional consumer-suppliedrefreshcallback to the Realtime subscription (typically() => api.agentGetChannel(conversationId)). useChat can't mint itself — the descriptor is produced by the consumer'ssubscribeadapter — so it relays the callback (structuralChatChannelDescriptortype, no new cross-package dep).Tokens are minted server-side (HMAC secret), so the client re-calls the server method that established the channel; the middleware invokes the app-supplied callback rather than minting locally.
Validation
refresh;refreshrejection falls back to backoff/onDisconnect; mock reconnect callsrefresh). Suite exits cleanly (no hang).useChatforwards a consumer-suppliedrefreshto the subscription options).tsc --buildclean ·biome lint0 errors · cast-free · bb-realtime API.md regenerated (+ refresh?).@aws-blocks/bb-realtimeminor,@aws-blocks/bb-agentminor.Maintainer notes
minor.refreshis opt-in; existing consumers are unaffected (they keep the current replay behavior, which only works within the first hour).Checklist
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.