@tanstack/ai-client 0.11.8 → 0.12.0

package/src/chat-client.ts

@@ -5,7 +5,10 @@ import {
   normalizeToUIMessage,
 } from '@tanstack/ai'
 import { DefaultChatClientEventEmitter } from './events'
-import { normalizeConnectionAdapter } from './connection-adapters'
+import {
+  fetcherToConnectionAdapter,
+  normalizeConnectionAdapter,
+} from './connection-adapters'
 import type {
   AnyClientTool,
   ContentPart,
@@ -20,6 +23,7 @@ import type { ChatClientEventEmitter } from './events'
 import type {
   ChatClientOptions,
   ChatClientState,
+  ChatFetcher,
   ConnectionStatus,
   MessagePart,
   MultimodalContent,
@@ -27,6 +31,21 @@ import type {
   UIMessage,
 } from './types'
 
+function resolveTransport(transport: {
+  connection?: ConnectionAdapter
+  fetcher?: ChatFetcher
+}): ConnectionAdapter {
+  const { connection, fetcher } = transport
+  if (connection && fetcher) {
+    throw new Error(
+      'ChatClient: pass either `connection` or `fetcher`, not both.',
+    )
+  }
+  if (connection) return connection
+  if (fetcher) return fetcherToConnectionAdapter(fetcher)
+  throw new Error('ChatClient: either `connection` or `fetcher` is required.')
+}
+
 export class ChatClient {
   private readonly processor: StreamProcessor
   private connection: SubscribeConnectionAdapter
@@ -96,7 +115,7 @@ export class ChatClient {
     // winning on key collision.
     this.bodyOption = options.body || {}
     this.forwardedPropsOption = options.forwardedProps || {}
-    this.connection = normalizeConnectionAdapter(options.connection)
+    this.connection = normalizeConnectionAdapter(resolveTransport(options))
     this.events = new DefaultChatClientEventEmitter(this.uniqueId)
 
     // Build client tools map
@@ -422,9 +441,9 @@ export class ChatClient {
         // both so a RUN_ERROR with a runId only clears that run, not every
         // active run in the session.
         const runId =
-          chunk.type === 'RUN_FINISHED'
+          'runId' in chunk && typeof chunk.runId === 'string'
             ? chunk.runId
-            : (chunk as { runId?: string }).runId
+            : undefined
         if (runId) {
           this.activeRunIds.delete(runId)
         } else if (chunk.type === 'RUN_ERROR') {
@@ -964,9 +983,13 @@ export class ChatClient {
    */
   private shouldAutoSend(): boolean {
     const messages = this.processor.getMessages()
-    const lastAssistant = messages.findLast((m) => m.role === 'assistant')
+    const lastAssistant = messages.findLast(
+      (m: UIMessage) => m.role === 'assistant',
+    )
     if (!lastAssistant) return false
-    const hasToolCalls = lastAssistant.parts.some((p) => p.type === 'tool-call')
+    const hasToolCalls = lastAssistant.parts.some(
+      (p: MessagePart) => p.type === 'tool-call',
+    )
     if (!hasToolCalls) return false
     return this.processor.areAllToolsComplete()
   }
@@ -1035,6 +1058,7 @@ export class ChatClient {
    */
   updateOptions(options: {
     connection?: ConnectionAdapter
+    fetcher?: ChatFetcher
     /** @deprecated Use `forwardedProps` instead. */
     body?: Record<string, any>
     forwardedProps?: Record<string, any>
@@ -1052,7 +1076,7 @@ export class ChatClient {
       context: { toolCallId?: string },
     ) => void
   }): void {
-    if (options.connection !== undefined) {
+    if (options.connection !== undefined || options.fetcher !== undefined) {
       const wasSubscribed = this.isSubscribed
 
       if (this.isLoading) {
@@ -1067,7 +1091,12 @@ export class ChatClient {
       this.resetSessionGenerating()
       this.setIsSubscribed(false)
       this.setConnectionStatus('disconnected')
-      this.connection = normalizeConnectionAdapter(options.connection)
+      this.connection = normalizeConnectionAdapter(
+        resolveTransport({
+          connection: options.connection,
+          fetcher: options.fetcher,
+        }),
+      )
 
       if (wasSubscribed) {
         this.subscribe()

package/src/connection-adapters.ts

@@ -1,10 +1,49 @@
 import { EventType, uiMessagesToWire } from '@tanstack/ai'
-import type { ModelMessage, StreamChunk, UIMessage } from '@tanstack/ai'
+import type {
+  ModelMessage,
+  RunErrorEvent,
+  RunFinishedEvent,
+  StreamChunk,
+  UIMessage,
+} from '@tanstack/ai'
+import type { ChatFetcher } from './types'
+
+/**
+ * Thrown when an SSE/HTTP stream ends with a non-empty unterminated buffer.
+ * Indicates the connection was cut mid-line (server crash, dropped TCP, proxy
+ * timeout) so the partial content cannot be safely parsed.
+ */
+export class StreamTruncatedError extends Error {
+  constructor() {
+    super(
+      'Stream ended with unterminated trailing data — connection was likely cut short.',
+    )
+    this.name = 'StreamTruncatedError'
+  }
+}
 
 function generateRunId(prefix: string): string {
   return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`
 }
 
+/**
+ * Asserts an id is present when synthesizing a terminal event. The chat
+ * client always supplies `runContext.threadId` / `runContext.runId`, so an
+ * absent id at this layer indicates the adapter was wired up by a caller
+ * that bypassed that contract — surface it rather than fabricating one.
+ */
+function requireSyntheticId(
+  value: string | undefined,
+  field: 'threadId' | 'runId',
+): string {
+  if (!value) {
+    throw new Error(
+      `Cannot synthesize terminal event: ${field} not supplied via runContext and not observed in the upstream stream.`,
+    )
+  }
+  return value
+}
+
 /**
  * Merge custom headers into request headers
  */
@@ -52,15 +91,83 @@ async function* readStreamLines(
       }
     }
 
-    // Process any remaining data in the buffer
-    if (buffer.trim()) {
-      yield buffer
+    // A non-empty trailing buffer means the connection was cut mid-line.
+    // Surface this as an error so the chat client transitions to 'error'
+    // state instead of silently presenting a partial stream as success.
+    // Skip when the consumer aborted — a user-initiated stop() interrupting
+    // mid-line is expected, not a truncation bug.
+    if (buffer.trim() && !abortSignal?.aborted) {
+      throw new StreamTruncatedError()
     }
   } finally {
     reader.releaseLock()
   }
 }
 
+/**
+ * Yield StreamChunks parsed from an SSE Response body.
+ *
+ * Accepts either `data: {...}` lines or bare JSON lines. Skips comments
+ * starting with `:` (proxies and CDNs inject these as keepalives) and the
+ * `event:` / `id:` / `retry:` SSE control fields. A `[DONE]` sentinel is
+ * treated as a terminal event: a synthesized RUN_FINISHED is yielded using
+ * the most recent upstream `threadId` / `runId`, ensuring the consumer sees
+ * a clean terminal event with real correlation ids.
+ *
+ * A JSON parse failure throws — the consumer surfaces it as an error.
+ */
+async function* responseToSSEChunks(
+  response: Response,
+  abortSignal?: AbortSignal,
+): AsyncGenerator<StreamChunk> {
+  if (!response.ok) {
+    throw new Error(
+      `HTTP error! status: ${response.status} ${response.statusText}`,
+    )
+  }
+  const reader = response.body?.getReader()
+  if (!reader) {
+    throw new Error('Response body is not readable')
+  }
+  let lastThreadId: string | undefined
+  let lastRunId: string | undefined
+  let lastModel: string | undefined
+  for await (const line of readStreamLines(reader, abortSignal)) {
+    if (
+      line.startsWith(':') ||
+      line.startsWith('event:') ||
+      line.startsWith('id:') ||
+      line.startsWith('retry:')
+    ) {
+      continue
+    }
+    const data = line.startsWith('data: ') ? line.slice(6) : line
+    if (data === '[DONE]') {
+      const synthetic: RunFinishedEvent = {
+        type: EventType.RUN_FINISHED,
+        threadId: lastThreadId ?? '',
+        runId: lastRunId ?? '',
+        model: lastModel ?? '',
+        timestamp: Date.now(),
+        finishReason: 'stop',
+      }
+      yield synthetic
+      return
+    }
+    const chunk = JSON.parse(data) as StreamChunk
+    if ('threadId' in chunk && typeof chunk.threadId === 'string') {
+      lastThreadId = chunk.threadId
+    }
+    if ('runId' in chunk && typeof chunk.runId === 'string') {
+      lastRunId = chunk.runId
+    }
+    if ('model' in chunk && typeof chunk.model === 'string') {
+      lastModel = chunk.model
+    }
+    yield chunk
+  }
+}
+
 /**
  * Per-send context provided by the chat client to the connection adapter.
  * The adapter combines this with serialized messages to build a full
@@ -196,6 +303,8 @@ export function normalizeConnectionAdapter(
     },
     async send(messages, data, abortSignal, runContext) {
       let hasTerminalEvent = false
+      let upstreamThreadId: string | undefined
+      let upstreamRunId: string | undefined
       try {
         const stream = connection.connect(
           messages,
@@ -204,6 +313,12 @@ export function normalizeConnectionAdapter(
           runContext,
         )
         for await (const chunk of stream) {
+          if ('threadId' in chunk && typeof chunk.threadId === 'string') {
+            upstreamThreadId = chunk.threadId
+          }
+          if ('runId' in chunk && typeof chunk.runId === 'string') {
+            upstreamRunId = chunk.runId
+          }
           if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') {
             hasTerminalEvent = true
           }
@@ -214,31 +329,40 @@ export function normalizeConnectionAdapter(
         // synthesize RUN_FINISHED so request-scoped consumers can complete.
         // Reuse the caller's threadId/runId so client-side activeRunIds tracking matches.
         if (!abortSignal?.aborted && !hasTerminalEvent) {
-          push({
+          const synthetic: RunFinishedEvent = {
             type: EventType.RUN_FINISHED,
-            threadId: runContext?.threadId ?? `thread-${Date.now()}`,
-            runId: runContext?.runId ?? `run-${Date.now()}`,
+            threadId: requireSyntheticId(
+              upstreamThreadId ?? runContext?.threadId,
+              'threadId',
+            ),
+            runId: requireSyntheticId(
+              upstreamRunId ?? runContext?.runId,
+              'runId',
+            ),
             model: 'connect-wrapper',
             timestamp: Date.now(),
             finishReason: 'stop',
-          })
+          }
+          push(synthetic)
         }
       } catch (err) {
         if (!abortSignal?.aborted && !hasTerminalEvent) {
-          push({
+          const message =
+            err instanceof Error ? err.message : 'Unknown error in connect()'
+          const synthetic: RunErrorEvent = {
             type: EventType.RUN_ERROR,
-            threadId: runContext?.threadId ?? `thread-${Date.now()}`,
-            runId: runContext?.runId ?? `run-${Date.now()}`,
+            threadId: requireSyntheticId(
+              upstreamThreadId ?? runContext?.threadId,
+              'threadId',
+            ),
+            runId: requireSyntheticId(
+              upstreamRunId ?? runContext?.runId,
+              'runId',
+            ),
             timestamp: Date.now(),
-            message:
-              err instanceof Error ? err.message : 'Unknown error in connect()',
-            error: {
-              message:
-                err instanceof Error
-                  ? err.message
-                  : 'Unknown error in connect()',
-            },
-          })
+            message,
+          }
+          push(synthetic)
         }
         throw err
       }
@@ -358,37 +482,7 @@ export function fetchServerSentEvents(
         ...(signal ? { signal } : {}),
       })
 
-      if (!response.ok) {
-        throw new Error(
-          `HTTP error! status: ${response.status} ${response.statusText}`,
-        )
-      }
-
-      // Parse Server-Sent Events format
-      const reader = response.body?.getReader()
-      if (!reader) {
-        throw new Error('Response body is not readable')
-      }
-
-      for await (const line of readStreamLines(reader, abortSignal)) {
-        // Handle Server-Sent Events format
-        const data = line.startsWith('data: ') ? line.slice(6) : line
-
-        if (data === '[DONE]') {
-          console.warn(
-            '[@tanstack/ai-client] Received [DONE] sentinel. This is deprecated — upgrade your @tanstack/ai server package. RUN_FINISHED is the stream terminator.',
-          )
-          continue
-        }
-
-        try {
-          const parsed: StreamChunk = JSON.parse(data)
-          yield parsed
-        } catch (parseError) {
-          // Skip non-JSON lines or malformed chunks
-          console.warn('Failed to parse SSE chunk:', data)
-        }
-      }
+      yield* responseToSSEChunks(response, abortSignal)
     },
   }
 }
@@ -507,12 +601,7 @@ export function fetchHttpStream(
       }
 
       for await (const line of readStreamLines(reader, abortSignal)) {
-        try {
-          const parsed: StreamChunk = JSON.parse(line)
-          yield parsed
-        } catch (parseError) {
-          console.warn('Failed to parse HTTP stream chunk:', line)
-        }
+        yield JSON.parse(line) as StreamChunk
       }
     },
   }
@@ -536,17 +625,92 @@ export function stream(
   streamFactory: (
     messages: Array<UIMessage> | Array<ModelMessage>,
     data?: Record<string, any>,
+    abortSignal?: AbortSignal,
   ) => AsyncIterable<StreamChunk>,
 ): ConnectConnectionAdapter {
   return {
-    async *connect(messages, data, _abortSignal, _runContext) {
+    async *connect(messages, data, abortSignal) {
       // Pass messages as-is (UIMessages with parts preserved)
       // Server-side chat() handles conversion to ModelMessages
-      yield* streamFactory(messages, data)
+      yield* streamFactory(messages, data, abortSignal)
+    },
+  }
+}
+
+/**
+ * Wrap a `ChatFetcher` as a `ConnectConnectionAdapter` so the chat client can
+ * consume it through the same `subscribe`/`send` plumbing used for SSE /
+ * HTTP-stream / RPC connections. May return either a `Response` (parsed as
+ * SSE) or an `AsyncIterable<StreamChunk>` (yielded directly).
+ *
+ * @internal
+ */
+export function fetcherToConnectionAdapter(
+  fetcher: ChatFetcher,
+): ConnectConnectionAdapter {
+  return {
+    async *connect(messages, data, abortSignal, runContext) {
+      if (!abortSignal) {
+        throw new Error(
+          'fetcherToConnectionAdapter requires an AbortSignal — the chat client always supplies one.',
+        )
+      }
+      if (!runContext) {
+        throw new Error(
+          'fetcherToConnectionAdapter requires a RunAgentInputContext — the chat client always supplies one.',
+        )
+      }
+      const uiMessages = messages as Array<UIMessage>
+      const result = await fetcher(
+        {
+          messages: uiMessages,
+          data,
+          threadId: runContext.threadId,
+          runId: runContext.runId,
+        },
+        { signal: abortSignal },
+      )
+      if (result instanceof Response) {
+        yield* responseToSSEChunks(result, abortSignal)
+      } else {
+        yield* abortableIterable(result, abortSignal)
+      }
     },
   }
 }
 
+/**
+ * Wrap an AsyncIterable so iteration aborts when `signal` fires. Without
+ * this, a fetcher that returns a generator ignoring its signal would leave
+ * the for-await loop hanging until the iterable naturally ends.
+ */
+async function* abortableIterable<T>(
+  iterable: AsyncIterable<T>,
+  signal: AbortSignal,
+): AsyncGenerator<T> {
+  if (signal.aborted) return
+  const iterator = iterable[Symbol.asyncIterator]()
+  const abortPromise = new Promise<{ done: true; value: undefined }>(
+    (resolve) => {
+      signal.addEventListener(
+        'abort',
+        () => resolve({ done: true, value: undefined }),
+        { once: true },
+      )
+    },
+  )
+  try {
+    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+    while (true) {
+      const result = await Promise.race([iterator.next(), abortPromise])
+      if (result.done) return
+      yield result.value
+    }
+  } finally {
+    await iterator.return?.()
+  }
+}
+
 /**
  * Create an RPC stream connection adapter (for RPC-based streaming like Cap'n Web RPC)
  *
@@ -567,13 +731,14 @@ export function rpcStream(
   rpcCall: (
     messages: Array<UIMessage> | Array<ModelMessage>,
     data?: Record<string, any>,
+    abortSignal?: AbortSignal,
   ) => AsyncIterable<StreamChunk>,
 ): ConnectConnectionAdapter {
   return {
-    async *connect(messages, data, _abortSignal, _runContext) {
+    async *connect(messages, data, abortSignal) {
       // Pass messages as-is (UIMessages with parts preserved)
       // Server-side chat() handles conversion to ModelMessages
-      yield* rpcCall(messages, data)
+      yield* rpcCall(messages, data, abortSignal)
     },
   }
 }

package/src/index.ts

@@ -17,7 +17,11 @@ export type {
   InferChatMessages,
   ChatClientState,
   ConnectionStatus,
-  // Multimodal content input type
+  ChatFetcher,
+  ChatFetcherInput,
+  ChatFetcherOptions,
+  ChatTransport,
+  DistributedOmit,
   MultimodalContent,
 } from './types'
 // Generation client types
@@ -58,6 +62,7 @@ export {
   fetchHttpStream,
   stream,
   rpcStream,
+  StreamTruncatedError,
   type ConnectConnectionAdapter,
   type ConnectionAdapter,
   type FetchConnectionOptions,

package/src/types.ts

@@ -16,6 +16,68 @@ import type { ConnectionAdapter } from './connection-adapters'
 
 export type { StructuredOutputPart } from '@tanstack/ai'
 
+/**
+ * `messages` is the full UIMessage history (not a delta). `data` is the
+ * merged body — `ChatClientOptions.body` plus any per-call data passed to
+ * `sendMessage(...)`. `threadId` / `runId` are the AG-UI correlation ids
+ * the chat client uses to track this turn — forward them to your server
+ * if it needs to correlate requests.
+ */
+export interface ChatFetcherInput {
+  messages: Array<UIMessage>
+  data?: Record<string, unknown>
+  threadId: string
+  runId: string
+}
+
+export interface ChatFetcherOptions {
+  /** Fires when `stop()` is called or the request is superseded. */
+  signal: AbortSignal
+}
+
+/**
+ * Direct function that performs a chat request. Mirrors
+ * `GenerationFetcher`. Returns either a `Response` (SSE body parsed by the
+ * chat client) or an `AsyncIterable<StreamChunk>` (yielded directly). May
+ * return the value synchronously, as a `Promise`, or as an async generator
+ * (`async function*`) — the chat client awaits whichever shape is returned.
+ *
+ * @example
+ * ```ts
+ * useChat({
+ *   fetcher: ({ messages }, { signal }) =>
+ *     chatFn({ data: { messages }, signal }),
+ * })
+ * ```
+ */
+export type ChatFetcher = (
+  input: ChatFetcherInput,
+  options: ChatFetcherOptions,
+) =>
+  | Response
+  | AsyncIterable<StreamChunk>
+  | Promise<Response | AsyncIterable<StreamChunk>>
+
+/**
+ * Distributive `Omit` — applies `Omit<O, K>` per branch of a union so
+ * discriminated unions survive omission. Plain `Omit` collapses unions
+ * into a single object shape, which would erase the `ChatTransport` XOR
+ * when framework hooks omit React-managed callbacks from
+ * `ChatClientOptions`.
+ */
+export type DistributedOmit<
+  TObject,
+  TKeys extends keyof any,
+> = TObject extends unknown ? Omit<TObject, TKeys> : never
+
+/**
+ * Discriminated union enforcing that exactly one of `connection` or
+ * `fetcher` is provided. Mirrors `GenerationTransport`.
+ */
+export type ChatTransport =
+  | { connection: ConnectionAdapter; fetcher?: never }
+  | { fetcher: ChatFetcher; connection?: never }
+
 /**
  * Tool call states - track the lifecycle of a tool call
  */
@@ -202,16 +264,13 @@ export interface UIMessage<
   createdAt?: Date
 }
 
-export interface ChatClientOptions<
+/**
+ * Options for `ChatClient`. Exactly one of `connection` or `fetcher` must be
+ * provided — the type-level XOR is enforced via `ChatTransport`.
+ */
+export type ChatClientOptions<
   TTools extends ReadonlyArray<AnyClientTool> = any,
-> {
-  /**
-   * Connection adapter for streaming.
-   * Supports mutually exclusive modes: request-response via `connect()`, or
-   * subscribe/send mode via `subscribe()` + `send()`.
-   */
-  connection: ConnectionAdapter
-
+> = {
   /**
    * Initial messages to populate the chat
    */
@@ -339,7 +398,7 @@ export interface ChatClientOptions<
      */
     chunkStrategy?: ChunkStrategy
   }
-}
+} & ChatTransport
 
 export interface ChatRequestBody {
   messages: Array<ModelMessage>