Skip to content

feat: WebSocket transport β€” full-duplex resumable streams#969

Open
AlemTuzlak wants to merge 24 commits into
mainfrom
feat/websocket-transport
Open

feat: WebSocket transport β€” full-duplex resumable streams#969
AlemTuzlak wants to merge 24 commits into
mainfrom
feat/websocket-transport

Conversation

@AlemTuzlak

@AlemTuzlak AlemTuzlak commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

WebSocket transport

Adds WebSockets as a third streaming transport alongside SSE and NDJSON, reusing the delivery-durability seam that landed in #955. Full-duplex, conversation-scoped, resumable.

Server (@tanstack/ai)

  • toWebSocketStream(socket, request, { onRun, durability?, batch?, heartbeatMs?, idleTimeoutMs?, debug? }) β€” portable core that pumps a conversation over an already-accepted WHATWG WebSocketLike server socket (Node via ws, Bun, etc.).
  • toWebSocketResponse(request, init) β€” thin wrapper that upgrades via WebSocketPair and returns a 101 Response on Cloudflare Workers/Durable Objects; throws elsewhere, pointing you to toWebSocketStream.
  • resumeWebSocketStream(socket, { adapter }) / resumeWebSocketResponse({ adapter }) β€” read-only replay of a run from the durability log (no model call).
  • Because one socket outlives many chat() turns (client-tool resubmits, follow-up messages), you pass an onRun(ctx) => AsyncIterable<StreamChunk> factory instead of a prebuilt stream β€” it's called per inbound RunAgentInput frame. Durability is keyed per turn and reuses the existing durableStreamSource (now exported), so serverβ†’client frames carry the same { id, chunk } envelope as NDJSON.
  • Conversation-scoped lifecycle: stays open across turns; closes on client close, an { type: 'abort', runId } frame (aborts only that turn), or idle timeout; periodic { type: 'ping' } heartbeat.

Client (@tanstack/ai-client, re-exported from ai-react/-solid/-vue/-svelte/-angular)

  • webSocket(url, options) β€” full-duplex subscribe + send connection adapter for useChat. send() writes a RunAgentInput frame; subscribe() yields inbound chunks, ignores heartbeats, unwraps durable envelopes, and auto-reconnects a dropped durable run by reopening with ?runId=&offset= (browsers can't set a Last-Event-ID handshake header, so the offset rides in the URL).
  • Reconnect bookkeeping (offset de-dupe, no-progress ceiling β†’ StreamReconnectLimitError) is shared with the HTTP adapters via the new createReconnectTracker. A fatal drop surfaces to the consumer instead of hanging.

Testing

  • Unit: @tanstack/ai stream-to-websocket 18/18; @tanstack/ai-client connection-adapter suites 63/63 (incl. reconnect, fatal-drop-surfacing, and open-promise-race regressions).
  • E2E: a provider-free WebSocket arm on the delivery-durability harness (Vite upgrade hook + ws) β€” ordered stream + reconnect-resume, 2/2; SSE/NDJSON regression 6/6.
  • Example: a /websocket-chat route in examples/ts-react-chat (type-checks + builds).
  • Docs: docs/resumable-streams/ overview + advanced + new WebSockets page; kiira-typechecked snippets, both server and client halves.

Notes / follow-ups (v2, documented)

  • Client stop() does not emit an { type: 'abort' } frame today β€” per-turn abort is via socket close (aborts all turns). The protocol primitive exists server-side.
  • One shared webSocket() connection assumes a single in-flight run at a time.

πŸ€– Generated with Claude Code

Summary by CodeRabbit

  • New Features

    • Added full-duplex WebSocket transport for chat and resumable streams.
    • Added automatic reconnection, replay, heartbeat, idle timeout, and per-turn abort support.
    • Added WebSocket adapters across supported client frameworks.
    • Added a WebSocket chat example with durable delivery and resume support.
  • Bug Fixes

    • Improved reconnect handling to prevent duplicate messages and surface stalled or failed connections.
  • Documentation

    • Added comprehensive WebSocket setup, protocol, hosting, and resumable-stream guidance.
  • Tests

    • Added unit, integration, and end-to-end coverage for WebSocket streaming and reconnection.

Add a full-duplex WebSocket chat demo to ts-react-chat: a Vite dev-server
plugin (websocket-chat-plugin.ts) hooks the Node http server's `upgrade`
event and wires toWebSocketStream/resumeWebSocketStream around chat() with
gpt-5.5, mirroring the pattern in testing/e2e's durable-delivery-ws-plugin
(no WebSocketPair on Node/Nitro). The /websocket-chat route uses useChat with
the webSocket() connection adapter, linked from the nav.

@tanstack/ai-react didn't re-export webSocket/WebSocketConnectionOptions from
@tanstack/ai-client yet (only the other connection adapters were), so this
adds that re-export as it's required for the example's `import { useChat,
webSocket } from '@tanstack/ai-react'`.
Document the new WebSocket transport (toWebSocketStream/toWebSocketResponse,
resumeWebSocketStream/resumeWebSocketResponse, client webSocket()): a new
standalone WebSockets page, a short intro + snippet in the Overview, a
reconnect/lifecycle summary in Advanced, and a pointer from Connection
Adapters to the built-in adapter instead of the old hand-rolled example.
@github-actions

Copy link
Copy Markdown
Contributor

πŸš€ Changeset Version Preview

8 package(s) bumped directly, 39 bumped as dependents.

πŸŸ₯ Major bumps

Package Version Reason
@tanstack/ai-angular 0.3.1 β†’ 1.0.0 Changeset
@tanstack/ai-durable-stream 0.0.0 β†’ 1.0.0 Changeset
@tanstack/ai-react 0.18.1 β†’ 1.0.0 Changeset
@tanstack/ai-solid 0.15.1 β†’ 1.0.0 Changeset
@tanstack/ai-svelte 0.15.1 β†’ 1.0.0 Changeset
@tanstack/ai-vue 0.15.1 β†’ 1.0.0 Changeset
@tanstack/ai-acp 0.2.3 β†’ 1.0.0 Dependent
@tanstack/ai-anthropic 0.16.3 β†’ 1.0.0 Dependent
@tanstack/ai-bedrock 0.1.4 β†’ 1.0.0 Dependent
@tanstack/ai-claude-code 0.2.3 β†’ 1.0.0 Dependent
@tanstack/ai-code-mode 0.3.8 β†’ 1.0.0 Dependent
@tanstack/ai-code-mode-skills 0.3.11 β†’ 1.0.0 Dependent
@tanstack/ai-codex 0.2.3 β†’ 1.0.0 Dependent
@tanstack/ai-elevenlabs 0.2.34 β†’ 1.0.0 Dependent
@tanstack/ai-fal 0.9.12 β†’ 1.0.0 Dependent
@tanstack/ai-gemini 0.20.1 β†’ 1.0.0 Dependent
@tanstack/ai-grok 0.14.9 β†’ 1.0.0 Dependent
@tanstack/ai-grok-build 0.2.3 β†’ 1.0.0 Dependent
@tanstack/ai-groq 0.5.3 β†’ 1.0.0 Dependent
@tanstack/ai-isolate-node 0.1.47 β†’ 1.0.0 Dependent
@tanstack/ai-isolate-quickjs 0.1.47 β†’ 1.0.0 Dependent
@tanstack/ai-mistral 0.2.3 β†’ 1.0.0 Dependent
@tanstack/ai-ollama 0.8.16 β†’ 1.0.0 Dependent
@tanstack/ai-openai 0.17.1 β†’ 1.0.0 Dependent
@tanstack/ai-opencode 0.2.3 β†’ 1.0.0 Dependent
@tanstack/ai-openrouter 0.15.10 β†’ 1.0.0 Dependent
@tanstack/ai-preact 0.11.1 β†’ 1.0.0 Dependent
@tanstack/ai-react-ui 0.8.15 β†’ 1.0.0 Dependent
@tanstack/ai-sandbox 0.2.4 β†’ 1.0.0 Dependent
@tanstack/ai-sandbox-cloudflare 0.2.4 β†’ 1.0.0 Dependent
@tanstack/ai-sandbox-daytona 0.2.0 β†’ 1.0.0 Dependent
@tanstack/ai-sandbox-docker 0.2.0 β†’ 1.0.0 Dependent
@tanstack/ai-sandbox-local-process 0.2.0 β†’ 1.0.0 Dependent
@tanstack/ai-sandbox-sprites 0.2.1 β†’ 1.0.0 Dependent
@tanstack/ai-sandbox-vercel 0.2.0 β†’ 1.0.0 Dependent
@tanstack/ai-solid-ui 0.7.14 β†’ 1.0.0 Dependent
@tanstack/openai-base 0.9.9 β†’ 1.0.0 Dependent

🟨 Minor bumps

Package Version Reason
@tanstack/ai 0.42.0 β†’ 0.43.0 Changeset
@tanstack/ai-client 0.22.1 β†’ 0.23.0 Changeset

🟩 Patch bumps

Package Version Reason
@tanstack/ai-devtools-core 0.4.24 β†’ 0.4.25 Dependent
@tanstack/ai-isolate-cloudflare 0.2.38 β†’ 0.2.39 Dependent
@tanstack/ai-mcp 0.2.5 β†’ 0.2.6 Dependent
@tanstack/ai-vue-ui 0.2.34 β†’ 0.2.35 Dependent
@tanstack/preact-ai-devtools 0.1.67 β†’ 0.1.68 Dependent
@tanstack/react-ai-devtools 0.2.67 β†’ 0.2.68 Dependent
@tanstack/solid-ai-devtools 0.2.67 β†’ 0.2.68 Dependent
ag-ui 0.0.2 β†’ 0.0.3 Dependent

@nx-cloud

nx-cloud Bot commented Jul 20, 2026

Copy link
Copy Markdown

View your CI Pipeline Execution β†— for commit 2225a78

Command Status Duration Result
nx run-many --targets=build --exclude=examples/... βœ… Succeeded 2m 12s View β†—

☁️ Nx Cloud last updated this comment at 2026-07-20 17:41:17 UTC

@coderabbitai

coderabbitai Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

πŸ“ Walkthrough

Walkthrough

A new resumable WebSocket transport adds server streaming and replay helpers, a client adapter with reconnect tracking, framework re-exports, React and end-to-end examples, comprehensive tests, and documentation covering protocol, lifecycle, hosting, and durability.

Changes

WebSocket transport

Layer / File(s) Summary
Server streaming, replay, and upgrade helpers
packages/ai/src/stream-to-websocket.ts, packages/ai/src/stream-to-response.ts, packages/ai/src/index.ts, packages/ai/tests/stream-to-websocket.test.ts
Adds WebSocket frame codecs, conversation-scoped streaming, per-turn abort handling, durable replay, heartbeat and idle timeout behavior, and WebSocketPair response wrappers.
Client adapter and reconnect tracking
packages/ai-client/src/connection-adapters.ts, packages/ai-client/tests/*, packages/ai-*/src/index.ts
Adds webSocket(), shared offset tracking, duplicate suppression, bounded reconnects, fatal stream errors, and framework package re-exports.
Examples and end-to-end coverage
examples/ts-react-chat/..., testing/e2e/...
Adds React WebSocket chat routing and hosting, a durable delivery WebSocket harness, and tests for streaming and offset-based resume.
Documentation and release metadata
.changeset/*, docs/chat/*, docs/resumable-streams/*, docs/config.json
Documents WebSocket selection, protocol frames, lifecycle, hosting, durability, reconnection, and release navigation metadata.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant useChat
  participant webSocket
  participant toWebSocketStream
  participant chat
  participant memoryStream
  useChat->>webSocket: send RunAgentInput
  webSocket->>toWebSocketStream: send run frame
  toWebSocketStream->>chat: invoke onRun with turn context
  chat->>memoryStream: persist durable chunks
  toWebSocketStream-->>webSocket: send chunk envelopes
  webSocket-->>useChat: yield streamed chunks
  webSocket->>toWebSocketStream: reconnect with runId and offset
  toWebSocketStream->>memoryStream: replay missing chunks
  memoryStream-->>useChat: deliver resumed chunks
Loading

Possibly related PRs

  • TanStack/ai#955: Both changes modify resumable-delivery reconnect tracking and shared durable streaming behavior.

Suggested reviewers: tombeckenham

πŸš₯ Pre-merge checks | βœ… 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 51.43% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
βœ… Passed checks (4 passed)
Check name Status Explanation
Title check βœ… Passed The title clearly and concisely summarizes the main change: adding a full-duplex WebSocket resumable stream transport.
Description check βœ… Passed The description thoroughly covers changes, testing, and limitations, though it does not follow the repository template sections exactly.
Linked Issues check βœ… Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check βœ… Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
πŸ“ Generate docstrings
  • Create stacked PR
  • Commit on current branch
πŸ§ͺ Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/websocket-transport

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

πŸ”§ ESLint

If the error stems from missing dependencies, add them to the package.json file. For unrecoverable errors (e.g., due to private dependencies), disable the tool in the CodeRabbit configuration.

ESLint install timed out. The project may have too many dependencies for the sandbox.


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❀️ Share

Comment @coderabbitai help to get the list of available commands.

@pkg-pr-new

pkg-pr-new Bot commented Jul 20, 2026

Copy link
Copy Markdown

Open in StackBlitz

@tanstack/ai

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai@969

@tanstack/ai-acp

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-acp@969

@tanstack/ai-angular

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-angular@969

@tanstack/ai-anthropic

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-anthropic@969

@tanstack/ai-bedrock

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-bedrock@969

@tanstack/ai-claude-code

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-claude-code@969

@tanstack/ai-client

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-client@969

@tanstack/ai-code-mode

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-code-mode@969

@tanstack/ai-code-mode-skills

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-code-mode-skills@969

@tanstack/ai-codex

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-codex@969

@tanstack/ai-devtools-core

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-devtools-core@969

@tanstack/ai-durable-stream

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-durable-stream@969

@tanstack/ai-elevenlabs

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-elevenlabs@969

@tanstack/ai-event-client

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-event-client@969

@tanstack/ai-fal

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-fal@969

@tanstack/ai-gemini

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-gemini@969

@tanstack/ai-grok

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-grok@969

@tanstack/ai-grok-build

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-grok-build@969

@tanstack/ai-groq

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-groq@969

@tanstack/ai-isolate-cloudflare

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-isolate-cloudflare@969

@tanstack/ai-isolate-node

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-isolate-node@969

@tanstack/ai-isolate-quickjs

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-isolate-quickjs@969

@tanstack/ai-mcp

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-mcp@969

@tanstack/ai-mistral

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-mistral@969

@tanstack/ai-ollama

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-ollama@969

@tanstack/ai-openai

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-openai@969

@tanstack/ai-opencode

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-opencode@969

@tanstack/ai-openrouter

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-openrouter@969

@tanstack/ai-preact

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-preact@969

@tanstack/ai-react

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-react@969

@tanstack/ai-react-ui

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-react-ui@969

@tanstack/ai-sandbox

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox@969

@tanstack/ai-sandbox-cloudflare

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-cloudflare@969

@tanstack/ai-sandbox-daytona

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-daytona@969

@tanstack/ai-sandbox-docker

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-docker@969

@tanstack/ai-sandbox-local-process

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-local-process@969

@tanstack/ai-sandbox-sprites

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-sprites@969

@tanstack/ai-sandbox-vercel

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-vercel@969

@tanstack/ai-solid

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-solid@969

@tanstack/ai-solid-ui

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-solid-ui@969

@tanstack/ai-svelte

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-svelte@969

@tanstack/ai-utils

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-utils@969

@tanstack/ai-vue

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-vue@969

@tanstack/ai-vue-ui

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-vue-ui@969

@tanstack/openai-base

npm i https://pkg.pr.new/TanStack/ai/@tanstack/openai-base@969

@tanstack/preact-ai-devtools

npm i https://pkg.pr.new/TanStack/ai/@tanstack/preact-ai-devtools@969

@tanstack/react-ai-devtools

npm i https://pkg.pr.new/TanStack/ai/@tanstack/react-ai-devtools@969

@tanstack/solid-ai-devtools

npm i https://pkg.pr.new/TanStack/ai/@tanstack/solid-ai-devtools@969

commit: 2225a78

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 5

🧹 Nitpick comments (2)
examples/ts-react-chat/src/lib/websocket-chat-plugin.ts (1)

40-46: πŸ“ Maintainability & Code Quality | πŸ”΅ Trivial | πŸ’€ Low value

Propagate the abort reason to the new controller.

When bridging an AbortSignal to a new AbortController, passing signal.reason ensures that any downstream cancellation logs or error handlers receive the correct context rather than a generic abort error.

  • examples/ts-react-chat/src/lib/websocket-chat-plugin.ts#L40-L46: pass signal.reason to controller.abort() in the plugin.
  • docs/resumable-streams/websockets.md#L56-L61: pass signal.reason to controller.abort() in the documentation sample.
πŸ€– Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/ts-react-chat/src/lib/websocket-chat-plugin.ts` around lines 40 -
46, Propagate the original abort reason when bridging signals: update
abortControllerFromSignal in
examples/ts-react-chat/src/lib/websocket-chat-plugin.ts lines 40-46 and the
corresponding abort listener in docs/resumable-streams/websockets.md lines 56-61
to pass signal.reason to controller.abort(), preserving the existing immediate
and event-driven cancellation behavior.
packages/ai/src/stream-to-websocket.ts (1)

204-211: πŸš€ Performance & Scalability | πŸ”΅ Trivial | βš–οΈ Poor tradeoff

No send backpressure β€” bufferedAmount is declared but never consulted.

Both pump loops call socket.send(...) on every chunk without awaiting drain. A fast onRun/durable replay against a slow client will queue frames unbounded in the socket's send buffer, growing memory. The WebSocketLike.bufferedAmount field (Line 20) appears intended for this but is unused; the resume pump at Line 265-266 has the same gap. Consider pausing iteration while bufferedAmount exceeds a high-water mark.

πŸ€– Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/ai/src/stream-to-websocket.ts` around lines 204 - 211, Update both
stream pump loops around source and init.onRun, plus the resume pump, to apply
send backpressure before calling socket.send. Pause iteration while
WebSocketLike.bufferedAmount exceeds a defined high-water threshold, then resume
once the buffer drains, preserving chunk order and existing frame encoding.
πŸ€– Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/chat/connection-adapters.md`:
- Around line 299-312: Add a concise server endpoint example to the WebSockets
section alongside the existing client webSocket usage, demonstrating the paired
toWebSocketStream or toWebSocketResponse integration. Keep the example
consistent with the documented server-side WebSocket APIs and retain the
existing link for protocol and hosting details.

In `@docs/resumable-streams/websockets.md`:
- Line 243: Remove the `as unknown as WebSocketLike` assertion from the
`socketLike` assignment in the WebSocket example; pass `ws` directly if it
satisfies the interface, otherwise use a type guard or property validation so
the sample type-checks without any `as` cast.

In `@packages/ai-client/src/connection-adapters.ts`:
- Around line 1690-1691: Update the ws.onmessage handler to guard
JSON.parse(String(event.data)) with try/catch. On malformed frames, call failAll
with a StreamReadError containing the caught error so stream consumers receive a
deterministic failure instead of hanging; preserve the existing handling for
successfully parsed messages.
- Around line 1721-1734: Update the ws.onclose handler in the connection adapter
to notify joinRun listeners when the socket closes even if currentSession is
undefined. Preserve the existing session reconnect and non-durable failure
behavior, while ensuring pending joinRun iterators are rejected or otherwise
terminated instead of remaining unresolved.

In `@packages/ai/src/stream-to-websocket.ts`:
- Around line 179-221: The handleInbound function must not overwrite an active
controller for the same runId. Before activeTurns.set(params.runId, turnAbort),
retrieve any existing controller, abort it or reject the duplicate run, then
store the new controller only when appropriate; preserve the existing ownership
check in the finally cleanup.

---

Nitpick comments:
In `@examples/ts-react-chat/src/lib/websocket-chat-plugin.ts`:
- Around line 40-46: Propagate the original abort reason when bridging signals:
update abortControllerFromSignal in
examples/ts-react-chat/src/lib/websocket-chat-plugin.ts lines 40-46 and the
corresponding abort listener in docs/resumable-streams/websockets.md lines 56-61
to pass signal.reason to controller.abort(), preserving the existing immediate
and event-driven cancellation behavior.

In `@packages/ai/src/stream-to-websocket.ts`:
- Around line 204-211: Update both stream pump loops around source and
init.onRun, plus the resume pump, to apply send backpressure before calling
socket.send. Pause iteration while WebSocketLike.bufferedAmount exceeds a
defined high-water threshold, then resume once the buffer drains, preserving
chunk order and existing frame encoding.
πŸͺ„ Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
βš™οΈ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 9e15cc23-dd08-43f6-a7e2-eebf70c11658

πŸ“₯ Commits

Reviewing files that changed from the base of the PR and between 7c7aa09 and 2225a78.

β›” Files ignored due to path filters (1)
  • pnpm-lock.yaml is excluded by !**/pnpm-lock.yaml
πŸ“’ Files selected for processing (30)
  • .changeset/websocket-transport.md
  • docs/chat/connection-adapters.md
  • docs/config.json
  • docs/resumable-streams/advanced.md
  • docs/resumable-streams/overview.md
  • docs/resumable-streams/websockets.md
  • examples/ts-react-chat/package.json
  • examples/ts-react-chat/src/components/Header.tsx
  • examples/ts-react-chat/src/lib/websocket-chat-plugin.ts
  • examples/ts-react-chat/src/routeTree.gen.ts
  • examples/ts-react-chat/src/routes/websocket-chat.tsx
  • examples/ts-react-chat/vite.config.ts
  • packages/ai-angular/src/index.ts
  • packages/ai-client/src/connection-adapters.ts
  • packages/ai-client/src/index.ts
  • packages/ai-client/tests/connection-adapters-resumable.test.ts
  • packages/ai-client/tests/connection-adapters-websocket.test.ts
  • packages/ai-react/src/index.ts
  • packages/ai-solid/src/index.ts
  • packages/ai-svelte/src/index.ts
  • packages/ai-vue/src/index.ts
  • packages/ai/src/index.ts
  • packages/ai/src/stream-to-response.ts
  • packages/ai/src/stream-to-websocket.ts
  • packages/ai/tests/stream-to-websocket.test.ts
  • testing/e2e/package.json
  • testing/e2e/src/lib/durable-delivery-ws-plugin.ts
  • testing/e2e/src/routes/api.durable-delivery.ts
  • testing/e2e/tests/websocket.spec.ts
  • testing/e2e/vite.config.ts

Comment on lines +299 to +312
## WebSockets

For a persistent, resumable WebSocket, use the built-in `webSocket()` adapter instead of hand-rolling a `SubscribeConnectionAdapter`. It opens one socket for the whole conversation, reconnects a dropped durable run automatically, and pairs with the server's `toWebSocketStream` / `toWebSocketResponse`:

```typescript
import { useChat, webSocket } from "@tanstack/ai-react";

const connection = webSocket("/api/chat-ws");

const { messages, sendMessage } = useChat({ connection });
```

See [WebSockets](../resumable-streams/websockets) for the server side, the wire protocol, reconnect details, and hosting on Node vs Cloudflare.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

πŸ“ Maintainability & Code Quality | 🟑 Minor | ⚑ Quick win

Add a server endpoint snippet for WebSockets.

As per coding guidelines, when a documentation page covers both server and client behavior, snippets for both halves (the server endpoint and the client consumption) must be included. The SSE section on this page includes a server snippet, but the new WebSockets section only provides the client webSocket snippet. Please add a brief example of the paired server endpoint (e.g., using toWebSocketStream or toWebSocketResponse).

πŸ€– Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/chat/connection-adapters.md` around lines 299 - 312, Add a concise
server endpoint example to the WebSockets section alongside the existing client
webSocket usage, demonstrating the paired toWebSocketStream or
toWebSocketResponse integration. Keep the example consistent with the documented
server-side WebSocket APIs and retain the existing link for protocol and hosting
details.

Source: Coding guidelines


wss.handleUpgrade(req, socket, head, (ws) => {
const request = new Request(url)
const socketLike = ws as unknown as WebSocketLike

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

πŸ“ Maintainability & Code Quality | 🟑 Minor | ⚑ Quick win

Remove the as type-assertion cast.

As per coding guidelines, do not use as type-assertion casts in documentation code samples. Examples must type-check without as SomeType; consider using a type guard or validating properties if the types do not natively align, or simply passing ws directly if it structurally satisfies the interface.

πŸ›  Proposed fix
-          const socketLike = ws as unknown as WebSocketLike
-
           if (url.searchParams.get('offset') !== null) {
-            resumeWebSocketStream(socketLike, { adapter: memoryStream(request) })
+            resumeWebSocketStream(ws, { adapter: memoryStream(request) })
           } else {
-            handleChatSocket(socketLike, request)
+            handleChatSocket(ws, request)
           }
πŸ“ Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
const socketLike = ws as unknown as WebSocketLike
if (url.searchParams.get('offset') !== null) {
resumeWebSocketStream(ws, { adapter: memoryStream(request) })
} else {
handleChatSocket(ws, request)
}
πŸ€– Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/resumable-streams/websockets.md` at line 243, Remove the `as unknown as
WebSocketLike` assertion from the `socketLike` assignment in the WebSocket
example; pass `ws` directly if it satisfies the interface, otherwise use a type
guard or property validation so the sample type-checks without any `as` cast.

Source: Coding guidelines

Comment on lines +1690 to +1691
ws.onmessage = (event: MessageEvent) => {
const parsed: unknown = JSON.parse(String(event.data))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟑 Minor | ⚑ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

file="packages/ai-client/src/connection-adapters.ts"

echo "== outline =="
ast-grep outline "$file" --view expanded || true

echo
echo "== relevant slice around lines 1660-1715 =="
sed -n '1660,1715p' "$file" | cat -n

echo
echo "== search for failAll and StreamReadError in file =="
rg -n "failAll|StreamReadError|onmessage|JSON\.parse" "$file"

Repository: TanStack/ai

Length of output: 8433


Guard malformed WebSocket frames before parsing

ws.onmessage parses event.data unguarded, so a bad frame throws out of the handler and leaves the stream without a deterministic failure path. Wrap this in try/catch and call failAll(new StreamReadError(error)) (or drop the frame explicitly) so consumers don’t hang on a corrupted message.

Suggested guard
     ws.onmessage = (event: MessageEvent) => {
-      const parsed: unknown = JSON.parse(String(event.data))
+      let parsed: unknown
+      try {
+        parsed = JSON.parse(String(event.data))
+      } catch (error) {
+        failAll(new StreamReadError(error))
+        return
+      }
πŸ“ Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
ws.onmessage = (event: MessageEvent) => {
const parsed: unknown = JSON.parse(String(event.data))
ws.onmessage = (event: MessageEvent) => {
let parsed: unknown
try {
parsed = JSON.parse(String(event.data))
} catch (error) {
failAll(new StreamReadError(error))
return
}
πŸ€– Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/ai-client/src/connection-adapters.ts` around lines 1690 - 1691,
Update the ws.onmessage handler to guard JSON.parse(String(event.data)) with
try/catch. On malformed frames, call failAll with a StreamReadError containing
the caught error so stream consumers receive a deterministic failure instead of
hanging; preserve the existing handling for successfully parsed messages.

Comment on lines +1721 to +1734
ws.onclose = () => {
const session = currentSession
if (!session || session.signal?.aborted || session.sawTerminal) return
const lastEventId = session.tracker.lastEventId
if (lastEventId === undefined) {
// Non-durable run (no offset ever observed) β€” nothing to resume
// from. Surface a hard failure rather than silently reconnecting
// forever against a server that never tags its events.
currentSession = undefined
failAll(new StreamReadError(new Error('WebSocket connection closed')))
return
}
void reconnect(session, lastEventId)
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟑 Minor | ⚑ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Map the file structure first
ast-grep outline packages/ai-client/src/connection-adapters.ts --view expanded || true

# Inspect the relevant region around the cited lines
sed -n '1600,1785p' packages/ai-client/src/connection-adapters.ts

# Find joinRun / send / subscribe / onclose usage
rg -n "joinRun|subscribe|send\(|onclose|currentSession|sawTerminal|lastEventId|reconnect" packages/ai-client/src/connection-adapters.ts

Repository: TanStack/ai

Length of output: 19758


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the runtime contract around the joinRun API and session lifecycle
sed -n '1,260p' packages/ai-client/src/connection-adapters.ts

# Look for tests or callers that clarify intended close behavior
rg -n "joinRun\\(|WebSocket connection closed|failAll\\(|abortSignal|StreamReadError|currentSession = undefined" packages -g '!**/dist/**'

Repository: TanStack/ai

Length of output: 34747


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Read the specific implementation in smaller chunks
cat -n packages/ai-client/src/connection-adapters.ts | sed -n '1,260p'

Repository: TanStack/ai

Length of output: 11228


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the subscribe() iterator and joinRun() implementation in the WebSocket adapter
sed -n '1770,1915p' packages/ai-client/src/connection-adapters.ts

# Inspect websocket tests around joinRun and close behavior
sed -n '1,260p' packages/ai-client/tests/connection-adapters-websocket.test.ts
sed -n '260,420p' packages/ai-client/tests/connection-adapters-websocket.test.ts

Repository: TanStack/ai

Length of output: 20365


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
text = Path("packages/ai-client/src/connection-adapters.ts").read_text()
start = text.index("    subscribe(abortSignal?: AbortSignal): AsyncIterable<StreamChunk> {")
end = text.index("export function stream(", start)
print(text[start:end])
PY

Repository: TanStack/ai

Length of output: 39994


Fail joinRun() listeners on socket close

ws.onclose returns early when currentSession is undefined, so a joinRun()-only socket drop can leave the iterator parked on its pending promise forever. Surface the close to join-run listeners too, even without an active send/session.

πŸ€– Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/ai-client/src/connection-adapters.ts` around lines 1721 - 1734,
Update the ws.onclose handler in the connection adapter to notify joinRun
listeners when the socket closes even if currentSession is undefined. Preserve
the existing session reconnect and non-durable failure behavior, while ensuring
pending joinRun iterators are rejected or otherwise terminated instead of
remaining unresolved.

Comment on lines +179 to +221
async function handleInbound(input: unknown): Promise<void> {
const params = await chatParamsFromRequestBody(input)
const turnAbort = new AbortController()
activeTurns.set(params.runId, turnAbort)
const ctx: WsRunContext = {
messages: params.messages,
threadId: params.threadId,
runId: params.runId,
forwardedProps: params.forwardedProps,
resumeOffset: null,
request: buildTurnRequest(request, params.runId, null),
signal: turnAbort.signal,
}
try {
if (init.durability) {
const adapter = init.durability(ctx)
const { source, getId } = durableStreamSource(
init.onRun(ctx),
adapter,
{
abortController: turnAbort,
...(init.batch === undefined ? {} : { batch: init.batch }),
logger,
},
)
for await (const chunk of source) {
socket.send(encodeWsFrame(chunk, getId(chunk)))
}
} else {
for await (const chunk of init.onRun(ctx)) {
socket.send(encodeWsFrame(chunk, undefined))
}
}
} finally {
// Only delete if this turn still owns the entry: a duplicate in-flight
// runId (e.g. a client resubmitting before the first turn finished)
// would otherwise let the OLDER turn's cleanup delete the NEWER turn's
// still-active controller (TOCTOU).
if (activeTurns.get(params.runId) === turnAbort) {
activeTurns.delete(params.runId)
}
}
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟑 Minor | ⚑ Quick win

🧩 Analysis chain

🏁 Script executed:

set -euo pipefail

# Map the relevant file structure and locate the symbols.
git ls-files packages/ai/src | sed -n '1,200p'
printf '\n--- outline ---\n'
ast-grep outline packages/ai/src/stream-to-websocket.ts --view expanded || true
printf '\n--- matches ---\n'
rg -n "activeTurns|runId|close handler|AbortController|handleInbound|onRun|durability" packages/ai/src/stream-to-websocket.ts packages/ai/src -S

Repository: TanStack/ai

Length of output: 28059


🏁 Script executed:

set -euo pipefail

# Read the relevant sections with line numbers.
sed -n '1,320p' packages/ai/src/stream-to-websocket.ts | cat -n

Repository: TanStack/ai

Length of output: 13582


🏁 Script executed:

set -euo pipefail

# Focused read of the file sections likely relevant to the duplicate-runId concern.
sed -n '120,260p' packages/ai/src/stream-to-websocket.ts | cat -n

Repository: TanStack/ai

Length of output: 5819


Do not overwrite an in-flight runId. A second run frame with the same runId replaces the existing AbortController in activeTurns, so the earlier turn becomes unreachable. Both the close handler and the abort control frame only act on the map entry, which means the orphaned turn can keep running until it finishes on its own. Abort or reject any existing controller before storing the new one.

πŸ€– Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/ai/src/stream-to-websocket.ts` around lines 179 - 221, The
handleInbound function must not overwrite an active controller for the same
runId. Before activeTurns.set(params.runId, turnAbort), retrieve any existing
controller, abort it or reject the duplicate run, then store the new controller
only when appropriate; preserve the existing ownership check in the finally
cleanup.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant