feat: session adapter and per-message stream state

[thruflo]

This PR introduces a SessionAdapter abstraction and refactors the ChatClient to support durable sessions, where conversation state survives page reloads, network interruptions and reconnections.

The core idea: instead of the ChatClient owning the request-response lifecycle for each LLM call, it delegates to a SessionAdapter that decouples sending messages from receiving chunks. Chunks arrive through a persistent subscription rather than being returned from the send call. This inversion makes it possible for an external session backend to replay missed events, inject message snapshots, and manage the stream lifecycle independently.

What changes

1. Per-message stream state in StreamProcessor (@tanstack/ai)

The StreamProcessor previously tracked a single "current assistant message" with flat instance variables (totalTextContent, toolCalls, etc.). This assumed one message streams at a time.

Now it maintains a Map<string, MessageStreamState> keyed by messageId, with a Set<string> of active message IDs and a toolCallId → messageId routing map. This allows:

• Multiple messages streaming concurrently (interleaved TEXT_MESSAGE_START / content / TEXT_MESSAGE_END events)
• Event-driven message creation from TEXT_MESSAGE_START rather than external startAssistantMessage() calls
• Proper lifecycle tracking per message (each TEXT_MESSAGE_END fires onStreamEnd for that message)

Backward compatibility is preserved: startAssistantMessage() still works by setting a pendingManualMessageId that gets reconciled when TEXT_MESSAGE_START arrives, or operates standalone if it never does.

2. AG-UI type alignment (@tanstack/ai)

Three additive type changes to match the AG-UI protocol:

• TextMessageStartEvent.role widened from 'assistant' to 'user' | 'assistant' | 'system' | 'tool' — session backends can replay messages of any role
• ToolCallStartEvent.parentMessageId added — enables routing tool calls to the correct message without relying on "current" message state
• MessagesSnapshotEvent added — a first-class AG-UI event for hydrating the full conversation transcript on connect/reconnect (distinct from STATE_SNAPSHOT which carries arbitrary application state)

3. SessionAdapter interface and createDefaultSession (@tanstack/ai-client)

interface SessionAdapter {
  subscribe(signal?: AbortSignal): AsyncIterable<StreamChunk>
  send(messages: Array<UIMessage>, data?: Record<string, any>, signal?: AbortSignal): Promise<void>
}

subscribe() returns a long-lived async iterable of AG-UI events. send() dispatches messages — responses arrive through subscribe(), not as a return value.

createDefaultSession(connection) wraps the existing ConnectionAdapter in this interface using an async queue: send() calls connection.connect() and pushes chunks to the queue; subscribe() yields them. This means existing ConnectionAdapter users get the new architecture transparently.

4. ChatClient refactored to use SessionAdapter (@tanstack/ai-client)

The constructor resolves a SessionAdapter — either provided directly via a new session option, or by wrapping the connection option with createDefaultSession. A background subscription loop (consumeSubscription) reads from subscribe() and feeds chunks to the StreamProcessor.

streamResponse() now coordinates with this loop: it calls session.send() to push chunks into the subscription, then awaits a processingResolve promise that fires when onStreamEnd is triggered by the processor. A streamGeneration counter prevents stale stream cleanup from clobbering a newer stream (e.g. when reload() is called during an active stream).

5. Framework hooks thread session option (ai-react, ai-solid, ai-vue, ai-svelte, ai-preact)

Each framework's useChat / createChat hook now passes options.session to the ChatClient constructor. Since ChatClientOptions accepts either connection or session, this flows through the existing options type with no additional framework-level changes.

Why this approach

The key design decision is Unified SessionAdapter (the ChatClient always operates through a SessionAdapter). The alternative — having ChatClient conditionally use either a ConnectionAdapter or SessionAdapter — would fork every method that touches the stream lifecycle. By always going through SessionAdapter, the ChatClient has a single code path, and createDefaultSession makes the wrapping invisible to existing users.

The per-message state refactor is a prerequisite because durable session backends emit TEXT_MESSAGE_START events for each message (potentially replaying an entire conversation). The processor needs to create and track messages from these events rather than relying on the client to call startAssistantMessage() in advance.

Migration

Existing usage is unchanged:

const { messages } = useChat({
  connection: fetchServerSentEvents('/api/chat'),
})

New session mode:

const { messages } = useChat({
  session: myDurableSessionAdapter,
})

For example:

import { createDurableSession } from '@durable-streams/sessions/tanstack-ai'

const { messages } = useChat({
  session: createDurableSession({
      proxyUrl: PROXY_URL,
      sessionId: '1234',
      sendUrl: '/api/chat',
      connectUrl: '/api/connect'
    }),
})

(Where connectUrl is an endpoint to authorize session access. If authorized, the session proxy will handle fetching / materializing message history from the session stream).

What's next

This PR establishes the transport and processing foundation. Follow-up work:

• @durable-streams/sessions/tanstack-ai — a concrete SessionAdapter implementation for durable stream backends
• expand the AG-UI support for STATE_SNAPSHOT / STATE_DELTA handling with a managed sessionState container on ChatClient
• onSessionStateChange and onCustomEvent callbacks for application-specific events
• Per-message isLoading tracking
• Server-side session mode support in TextEngine

[samwillis]

I had a good look through the code, it's sensibly structured, and follows the patterns from the repo.

I also asked Codex5.3 and Opus4.6 (with the later having the context of the Codex review) to do reviews - I'll past them below.

[samwillis]

Codex Review:

---

Findings

• High – streamResponse() can unblock on the wrong stream/message end event
  ChatClient uses a single global processingResolve, and resolves it on any onStreamEnd. In multi-message/interleaved sessions, the first ended message can unblock await processingComplete too early. It can also be resolved by a stale stream’s end event because there is no generation/message guard at resolve time.
  This can flip isLoading/status early and run continuation logic before the active response is actually finished.

onStreamEnd: (message: UIMessage) => {
  this.callbacksRef.current.onFinish(message)
  this.setStatus('ready')
  // Resolve the processing-complete promise so streamResponse can continue
  this.processingResolve?.()
  this.processingResolve = null
},

// Set up promise that resolves when onStreamEnd fires
const processingComplete = this.waitForProcessing()

// Send through session adapter (pushes chunks to subscription queue)
await this.session.send(messages, mergedBody, this.abortController.signal)

// Wait for subscription loop to finish processing all chunks
await processingComplete

• High – pending manual assistant message can be incorrectly reconciled to non-assistant TEXT_MESSAGE_START roles
  streamResponse() always pre-creates an assistant message via startAssistantMessage(). If the first streamed start event is a replayed user/system/tool message, reconciliation updates ID but does not update role/state role, so that message can be represented as assistant incorrectly.
  This is especially risky for durable/replay backends (which this PR explicitly enables).

// Generate stream ID and start assistant message
this.currentStreamId = this.generateUniqueId('stream')
const messageId = this.processor.startAssistantMessage()
this.currentMessageId = messageId

// Case 1: A manual message was created via startAssistantMessage()
if (this.pendingManualMessageId) {
  const pendingId = this.pendingManualMessageId
  this.pendingManualMessageId = null

  if (pendingId !== messageId) {
    // Update the message's ID in the messages array
    this.messages = this.messages.map((msg) =>
      msg.id === pendingId ? { ...msg, id: messageId } : msg,
    )
    // ... state key move only
  }

  // Ensure state exists
  if (!this.messageStates.has(messageId)) {
    this.createMessageState(messageId, role)
    this.activeMessageIds.add(messageId)
  }

• Medium – MESSAGES_SNAPSHOT does not reset transient stream state maps
  Snapshot replaces messages, but leaves messageStates, activeMessageIds, and toolCallToMessage untouched. After reconnect/hydration, stale active/tool mappings can affect subsequent tool routing/finalization behavior.

// Per-message stream state
private messageStates: Map<string, MessageStreamState> = new Map()
private activeMessageIds: Set<string> = new Set()
private toolCallToMessage: Map<string, string> = new Map()
private pendingManualMessageId: string | null = null

private handleMessagesSnapshotEvent(
  chunk: Extract<StreamChunk, { type: 'MESSAGES_SNAPSHOT' }>,
): void {
  this.messages = [...chunk.messages]
  this.emitMessagesChange()
}

• Medium – createDefaultSession.subscribe() can accumulate abort listeners over long-lived sessions
  A new abort listener is attached each time the subscriber waits, but not removed when a chunk arrives normally. In durable long-lived subscriptions this can build up listener/memory overhead until abort.

chunk = await new Promise<StreamChunk | null>((resolve) => {
  waiters.push(resolve)
  signal?.addEventListener('abort', () => resolve(null), {
    once: true,
  })
})

Testing gaps

• No ChatClient tests currently exercise session mode behavior (only constructor validation): packages/typescript/ai-client/tests/chat-client.test.ts.
• Missing tests for:
  • interleaved multi-message stream completion semantics in ChatClient,
  • replayed non-assistant TEXT_MESSAGE_START with pending manual message,
  • MESSAGES_SNAPSHOT arriving mid-session with prior active state.

Brief take

The architectural direction is strong, but there are a few stream-lifecycle edge cases that could cause incorrect completion/loading behavior or wrong message-role reconciliation in durable/replay scenarios.

[samwillis]

Opus review:

---

Review

Overall, this is a well-structured PR. The architectural decision to unify around SessionAdapter is sound -- it avoids forking every stream-lifecycle method into two code paths and makes the upgrade invisible to existing ConnectionAdapter users. The per-message state refactor in StreamProcessor is a clean prerequisite for durable sessions. The test coverage on the new StreamProcessor behavior (interleaved messages, TEXT_MESSAGE_START/END lifecycle, MESSAGES_SNAPSHOT, startAssistantMessage dedup) is solid.

I have a few findings of my own, and I'll comment on the other review's findings where I agree or disagree.

---

1. Double onStreamEnd when RUN_FINISHED precedes TEXT_MESSAGE_END -- High

This is new to this PR and the most concrete bug I found. The previous reviewer flagged the processingResolve being globally resolved by "any" onStreamEnd, which is related but I think the sharper framing is:

If a server sends RUN_FINISHED before TEXT_MESSAGE_END (which is a valid ordering), both fire onStreamEnd:

• RUN_FINISHED calls finalizeStream(), which iterates activeMessageIds, sets isComplete = true, clears activeMessageIds, then fires onStreamEnd for the last assistant message.
• TEXT_MESSAGE_END then arrives. handleTextMessageEndEvent looks up the state (which exists), finds the message, and fires onStreamEnd again -- it doesn't check state.isComplete.

In ChatClient, onStreamEnd fires onFinish(message) and setStatus('ready'). The first call resolves processingResolve (harmless second call since it's nulled). But onFinish fires twice, which is a user-visible bug for callback consumers.

The reverse order (TEXT_MESSAGE_END before RUN_FINISHED) is safe because finalizeStream() only iterates remaining activeMessageIds, which is already empty.

Fix: Add an early return at the top of handleTextMessageEndEvent:

if (state.isComplete) return

2. createDefaultSession: dead waiters cause lost chunks after stop/restart -- High

This is a new finding not in the previous review. The async queue in createDefaultSession has a subtle bug around the interaction between subscribe(), push(), and abort:

When the subscription loop is idle (waiting for chunks), a waiter resolve function sits in the waiters array. When the signal aborts, the abort listener calls resolve(null), and the subscribe() generator exits, clearing the buffer with buffer.length = 0. But the dead waiter is never removed from the waiters array.

On the next subscription cycle, when send() pushes a chunk, push() calls waiters.shift() -- which returns the dead waiter. Calling it is a no-op (promise already resolved). The chunk is silently dropped.

This happens every time stop() is called while the subscription is idle between messages. The first chunk of the next stream is lost.

Fix: Clear waiters alongside buffer when the subscription exits:

// At end of subscribe generator:
buffer.length = 0
waiters.length = 0

Or, more defensively, have push() skip dead waiters:

function push(chunk: StreamChunk): void {
  // Skip resolved waiters
  while (waiters.length > 0) {
    const waiter = waiters.shift()!
    // ...need a way to track if resolved
  }
}

The simplest approach is clearing waiters when subscribe exits.

3. MESSAGES_SNAPSHOT doesn't reset transient stream state -- Medium

I agree with the previous reviewer here. handleMessagesSnapshotEvent replaces this.messages but leaves messageStates, activeMessageIds, toolCallToMessage, and pendingManualMessageId stale. If a snapshot arrives while messages are actively streaming (e.g., reconnection during a stream), the stale maps could route tool calls or text content to messages that no longer exist.

private handleMessagesSnapshotEvent(
  chunk: Extract<StreamChunk, { type: 'MESSAGES_SNAPSHOT' }>,
): void {
  this.messages = [...chunk.messages]
  this.emitMessagesChange()
}

Fix: Reset the transient state maps when a snapshot arrives:

private handleMessagesSnapshotEvent(
  chunk: Extract<StreamChunk, { type: 'MESSAGES_SNAPSHOT' }>,
): void {
  this.messages = [...chunk.messages]
  this.messageStates.clear()
  this.activeMessageIds.clear()
  this.toolCallToMessage.clear()
  this.pendingManualMessageId = null
  this.emitMessagesChange()
}

For the primary use case (snapshot at subscription start, before any streaming) this isn't a problem because the maps are empty. But it's a correctness issue that will bite when durable sessions reconnect mid-stream.

4. Pending manual message reconciliation with non-assistant roles -- Medium

I agree with the previous reviewer, but I'd lower the severity to Medium for now. In handleTextMessageStartEvent Case 1, when reconciling a pending manual message with a TEXT_MESSAGE_START event, the code updates the message ID but not the role:

this.messages = this.messages.map((msg) =>
  msg.id === pendingId ? { ...msg, id: messageId } : msg,
)

If a durable session replays a user message before the assistant response, and that user message is the first TEXT_MESSAGE_START, it would be reconciled with the pre-created assistant message without updating role. The MessageStreamState.role would be correct (from the event), but the UIMessage.role in the messages array would still be 'assistant'.

In practice, for the described durable session flow (MESSAGES_SNAPSHOT hydrates history, then assistant TEXT_MESSAGE_START arrives), this shouldn't happen because the snapshot doesn't consume pendingManualMessageId. But it's a latent bug for other session implementations.

Fix: Update role during reconciliation:

msg.id === pendingId ? { ...msg, id: messageId, role: uiRole } : msg,

Where uiRole maps 'tool' to 'assistant' as in Case 3.

5. Abort listener accumulation in createDefaultSession -- Low

I agree with the previous reviewer that this exists, but I'd rate it Low. Each iteration of the subscribe() while loop that enters the await new Promise(...) path adds an abort listener. When a chunk arrives (resolving the waiter normally), the listener stays registered on the signal. Over thousands of chunks in a long-lived session, this accumulates listeners.

The { once: true } means each fires at most once, and they're all on the same AbortSignal which gets GC'd when the subscription ends. So the leak is bounded by the subscription lifetime and is more of a performance concern than a correctness issue.

Fix: Track and remove the listener when the waiter is resolved normally:

chunk = await new Promise<StreamChunk | null>((resolve) => {
  const onAbort = () => resolve(null)
  waiters.push((c) => {
    signal?.removeEventListener('abort', onAbort)
    resolve(c)
  })
  signal?.addEventListener('abort', onAbort, { once: true })
})

---

On the previous review's "processingResolve resolves on wrong message" finding

The previous reviewer flagged that processingResolve could be resolved by any onStreamEnd including from interleaved messages. I partially agree -- this is a real concern for durable sessions that might stream multiple messages. For the createDefaultSession path, the flow is tightly coupled: send() pushes all chunks, onStreamEnd fires once (from either TEXT_MESSAGE_END or finalizeStream via RUN_FINISHED), and processingResolve resolves. The coordination works because there's exactly one assistant message per streamResponse() call.

For durable sessions though, if subscribe() delivers events for multiple messages (e.g., replaying history as individual TEXT_MESSAGE_START/END pairs), each TEXT_MESSAGE_END fires onStreamEnd, and the first one would prematurely resolve processingResolve. This is real but it's a design question for durable session implementations: do they replay history as individual message events, or as a MESSAGES_SNAPSHOT? The PR description suggests MESSAGES_SNAPSHOT, which wouldn't trigger onStreamEnd.

I'd note this as a known limitation to document rather than a bug to fix now -- durable session adapters should use MESSAGES_SNAPSHOT for history hydration, not individual TEXT_MESSAGE_START/END events.

---

Testing gaps

The test coverage is good for StreamProcessor and createDefaultSession individually. I'd note these gaps:

1. No test for RUN_FINISHED + TEXT_MESSAGE_END in the same stream -- this would catch the double onStreamEnd bug (finding feat: session adapter and per-message stream state #1).
2. No ChatClient integration tests with session option -- the existing tests all use connection. Even one test exercising the session path end-to-end would be valuable.
3. No test for stop-then-resume cycle through createDefaultSession -- this would catch the dead waiter bug (finding simplify durable session support #2).

---

Minor observations

• The session!: SessionAdapter definite assignment assertion is fine given the constructor always assigns or throws, but a private session: SessionAdapter with a dummy initial value or making it a constructor parameter would be cleaner.
• The streamGeneration counter for superseded stream protection is a nice pattern.
• The framework hook changes are purely mechanical (threading session through) and look correct across all five frameworks.

[thruflo]

Plan: Address PR Review Findings for Resumable Sessions

Fix 1: Guard against double onStreamEnd — High

Problem: If RUN_FINISHED arrives before TEXT_MESSAGE_END, both fire onStreamEnd. finalizeStream() (called by RUN_FINISHED) sets state.isComplete = true and fires onStreamEnd. Then handleTextMessageEndEvent doesn't check isComplete and fires onStreamEnd again. This causes the onFinish callback to fire twice.

File: processor.ts, handleTextMessageEndEvent (~line 650)

Change: Add early return after existing if (!state) return:

if (state.isComplete) return

isComplete is per-message state in MessageStreamState, keyed by messageId in the messageStates Map — safe for interleaved concurrent messages.

New test in stream-processor.test.ts: RUN_FINISHED then TEXT_MESSAGE_END in same stream — verify onStreamEnd fires exactly once.

Fix 2: Clear dead waiters on subscribe exit — High

Problem: In createDefaultSession, the abort listener resolves the waiter promise with null, but the dead resolve function remains in the shared waiters array. On the next subscription cycle, push() shifts the dead waiter, calls it (no-op — promise already resolved), and the chunk is silently dropped.

File: session-adapter.ts, subscribe generator cleanup (~line 76)

Change: Add after existing buffer.length = 0:

waiters.length = 0

New test in session-adapter.test.ts: Stop-then-resume cycle — abort subscription while waiting, start new subscription, send chunks, verify first chunk is not lost.

Fix 3: Reset transient state on MESSAGES_SNAPSHOT — Medium

Problem: handleMessagesSnapshotEvent replaces this.messages but leaves messageStates, activeMessageIds, toolCallToMessage, and pendingManualMessageId stale. On reconnection mid-stream, stale maps could route content to messages that no longer exist.

File: processor.ts, handleMessagesSnapshotEvent (~line 674)

Change: Add 4 lines before this.emitMessagesChange():

this.messageStates.clear()
this.activeMessageIds.clear()
this.toolCallToMessage.clear()
this.pendingManualMessageId = null

New test in stream-processor.test.ts: MESSAGES_SNAPSHOT arriving after active streaming state — verify stale state is reset and subsequent TEXT_MESSAGE_START/TEXT_MESSAGE_CONTENT events are processed correctly.

Fix 4: Remove startAssistantMessage() from streamResponse() — Medium

Problem: streamResponse() optimistically creates an assistant message via startAssistantMessage() before the stream starts. This causes:

1. Reconciliation complexity when TEXT_MESSAGE_START arrives with a different ID/role (Finding 4 from reviews)
2. Orphan empty assistant messages on error (unreported bug — catch block doesn't clean up)
3. resetStreamState() on every send, which would destroy state of concurrent in-progress responses in the session paradigm

Approach: Remove startAssistantMessage() from streamResponse(). Let stream events create the assistant message naturally via TEXT_MESSAGE_START (Case 3 in handleTextMessageStartEvent) or ensureAssistantMessage() (backward compat auto-creation from TEXT_MESSAGE_CONTENT). startAssistantMessage() remains on StreamProcessor as a public method for the process() batch API.

File: chat-client.ts

Changes in streamResponse() (~lines 447-461):

Remove:

const messageId = this.processor.startAssistantMessage()
this.currentMessageId = messageId

const assistantMessage: UIMessage = {
  id: messageId,
  role: 'assistant',
  parts: [],
  createdAt: new Date(),
}
this.events.messageAppended(
  assistantMessage,
  this.currentStreamId || undefined,
)

Replace with:

this.currentMessageId = null

Keep this.currentStreamId = this.generateUniqueId('stream') (still needed for devtools).

Changes in onStreamStart handler (~lines 105-107):

Update to capture currentMessageId and emit messageAppended when the stream creates the message:

onStreamStart: () => {
  this.setStatus('streaming')
  const messages = this.processor.getMessages()
  const lastAssistant = messages.findLast((m: UIMessage) => m.role === 'assistant')
  if (lastAssistant) {
    this.currentMessageId = lastAssistant.id
    this.events.messageAppended(lastAssistant, this.currentStreamId || undefined)
  }
},

Changes in consumeSubscription() (~line 271):

After this.processor.processChunk(chunk), add RUN_FINISHED handling to resolve processingResolve. This covers the edge case where no message is created (empty stream) and onStreamEnd never fires:

this.processor.processChunk(chunk)
// RUN_FINISHED signals run completion — resolve processing
// (redundant if onStreamEnd already resolved it, harmless)
if (chunk.type === 'RUN_FINISHED') {
  this.processingResolve?.()
  this.processingResolve = null
}

Fix 5: Clean up abort listeners on normal resolution — Low

Problem: Each iteration through the promise-wait path in subscribe() adds an abort listener. When the waiter resolves normally (chunk arrives), the listener stays registered. Over long-lived subscriptions, listeners accumulate.

File: session-adapter.ts, subscribe promise-wait path (~lines 65-70)

Change: Wrap the waiter to removeEventListener when resolved normally:

chunk = await new Promise<StreamChunk | null>((resolve) => {
  const onAbort = () => resolve(null)
  waiters.push((c) => {
    signal?.removeEventListener('abort', onAbort)
    resolve(c)
  })
  signal?.addEventListener('abort', onAbort, { once: true })
})

No new test — existing abort tests provide coverage.

[thruflo]

All fixes implemented and verified. Here's a summary of the changes:

Summary

Fix 1: Guard against double onStreamEnd (High)

processor.ts:651 — Added if (state.isComplete) return in handleTextMessageEndEvent to prevent firing onStreamEnd twice when RUN_FINISHED arrives before TEXT_MESSAGE_END.

Fix 2: Clear dead waiters on subscribe exit (High)

session-adapter.ts:77 — Added waiters.length = 0 in subscribe cleanup to prevent dead resolve functions from silently dropping chunks on the next subscription cycle.

Fix 3: Reset transient state on MESSAGES_SNAPSHOT (Medium)

processor.ts:678-681 — Clear messageStates, activeMessageIds, toolCallToMessage, and pendingManualMessageId when a snapshot replaces messages, preventing stale state
from routing content to non-existent messages on reconnection.

Fix 4: Remove startAssistantMessage() from streamResponse() (Medium)

chat-client.ts — Three changes:

1. Removed the optimistic startAssistantMessage() call and synthetic messageAppended event. Let stream events create the assistant message naturally via
   TEXT_MESSAGE_START or ensureAssistantMessage.
2. Updated onStreamStart handler to capture currentMessageId and emit messageAppended when the processor creates the message.
3. Added RUN_FINISHED handling in consumeSubscription() to resolve processingResolve for empty streams where onStreamEnd never fires.

Bonus fix in processor.ts:905-910 — Made handleStepFinishedEvent use ensureAssistantMessage() instead of silently returning when no active assistant exists, matching the
backward-compat pattern used by handleTextMessageContentEvent.

Fix 5: Clean up abort listeners on normal resolution (Low)

session-adapter.ts:65-72 — Wrapped the waiter callback to removeEventListener when resolved normally, preventing abort listener accumulation on long-lived subscriptions.

New Tests

• stream-processor.test.ts: "double onStreamEnd guard" and "MESSAGES_SNAPSHOT resets transient state"
• session-adapter.test.ts: "should not lose chunks after stop-then-resume subscription cycle"

[thruflo]

Superseded by TanStack/ai PR