feat: Subscription Lifecycle Amendment

packages/typescript/ai-client/src/connection-adapters.ts

@@ -62,15 +62,9 @@ async function* readStreamLines(
   }
 }
 
-/**
- * Connection adapter interface - converts a connection into a stream of chunks
- */
-export interface ConnectionAdapter {
+export interface ConnectConnectionAdapter {
   /**
-   * Connect and return an async iterable of StreamChunks
-   * @param messages - The messages to send (UIMessages or ModelMessages)
-   * @param data - Additional data to send
-   * @param abortSignal - Optional abort signal for request cancellation
+   * Connect and return an async iterable of StreamChunks.
    */
   connect: (
     messages: Array<UIMessage> | Array<ModelMessage>,
@@ -79,6 +73,147 @@ export interface ConnectionAdapter {
   ) => AsyncIterable<StreamChunk>
 }
 
+export interface SubscribeConnectionAdapter {
+  /**
+   * Subscribe to stream chunks.
+   */
+  subscribe: (abortSignal?: AbortSignal) => AsyncIterable<StreamChunk>
+  /**
+   * Send a request; chunks arrive through subscribe().
+   */
+  send: (
+    messages: Array<UIMessage> | Array<ModelMessage>,
+    data?: Record<string, any>,
+    abortSignal?: AbortSignal,
+  ) => Promise<void>
+}
+
+/**
+ * Connection adapter union.
+ * Provide either `connect`, or `subscribe` + `send`.
+ */
+export type ConnectionAdapter =
+  | ConnectConnectionAdapter
+  | SubscribeConnectionAdapter
+
+/**
+ * Normalize a ConnectionAdapter to subscribe/send operations.
+ *
+ * If a connection provides native subscribe/send, that mode is used.
+ * Otherwise, connect() is wrapped using an async queue.
+ */
+export function normalizeConnectionAdapter(
+  connection: ConnectionAdapter | undefined,
+): SubscribeConnectionAdapter {
+  if (!connection) {
+    throw new Error('Connection adapter is required')
+  }
+
+  const hasConnect = 'connect' in connection
+  const hasSubscribe = 'subscribe' in connection
+  const hasSend = 'send' in connection
+
+  if (hasConnect && (hasSubscribe || hasSend)) {
+    throw new Error(
+      'Connection adapter must provide either connect or both subscribe and send, not both modes',
+    )
+  }
+
+  if (hasSubscribe && hasSend) {
+    return {
+      subscribe: connection.subscribe.bind(connection),
+      send: connection.send.bind(connection),
+    }
+  }
+
+  if (!hasConnect) {
+    throw new Error(
+      'Connection adapter must provide either connect or both subscribe and send',
+    )
+  }
+
+  // Legacy connect() wrapper
+  let activeBuffer: Array<StreamChunk> = []
+  let activeWaiters: Array<(chunk: StreamChunk | null) => void> = []
+
+  function push(chunk: StreamChunk): void {
+    const waiter = activeWaiters.shift()
+    if (waiter) {
+      waiter(chunk)
+    } else {
+      activeBuffer.push(chunk)
+    }
+  }
+
+  return {
+    subscribe(abortSignal?: AbortSignal): AsyncIterable<StreamChunk> {
+      // Transfer ownership to the latest subscriber so only one active
+      // subscribe() call receives chunks from the shared connect-wrapper queue.
+      const myBuffer: Array<StreamChunk> = activeBuffer.splice(0)
+      const myWaiters: Array<(chunk: StreamChunk | null) => void> = []
+      activeBuffer = myBuffer
+      activeWaiters = myWaiters
+
+      return (async function* () {
+        while (!abortSignal?.aborted) {
+          let chunk: StreamChunk | null
+          if (myBuffer.length > 0) {
+            chunk = myBuffer.shift()!
+          } else {
+            chunk = await new Promise<StreamChunk | null>((resolve) => {
+              const onAbort = () => resolve(null)
+              myWaiters.push((c) => {
+                abortSignal?.removeEventListener('abort', onAbort)
+                resolve(c)
+              })
+              abortSignal?.addEventListener('abort', onAbort, { once: true })
+            })
+          }
+          if (chunk !== null) yield chunk
+        }
+      })()
+    },
+    async send(messages, data, abortSignal) {
+      let hasTerminalEvent = false
+      try {
+        const stream = connection.connect(messages, data, abortSignal)
+        for await (const chunk of stream) {
+          if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') {
+            hasTerminalEvent = true
+          }
+          push(chunk)
+        }
+
+        // If the connect stream ended cleanly without a terminal event,
+        // synthesize RUN_FINISHED so request-scoped consumers can complete.
+        if (!abortSignal?.aborted && !hasTerminalEvent) {
+          push({
+            type: 'RUN_FINISHED',
+            runId: `run-${Date.now()}`,
+            model: 'connect-wrapper',
+            timestamp: Date.now(),
+            finishReason: 'stop',
+          })
+        }
+      } catch (err) {
+        if (!abortSignal?.aborted && !hasTerminalEvent) {
+          push({
+            type: 'RUN_ERROR',
+            timestamp: Date.now(),
+            error: {
+              message:
+                err instanceof Error
+                  ? err.message
+                  : 'Unknown error in connect()',
+            },
+          })
+        }
+        throw err
+      }
+    },
+  }
+}
+
 /**
  * Options for fetch-based connection adapters
  */
@@ -129,7 +264,7 @@ export function fetchServerSentEvents(
   options:
     | FetchConnectionOptions
     | (() => FetchConnectionOptions | Promise<FetchConnectionOptions>) = {},
-): ConnectionAdapter {
+): ConnectConnectionAdapter {
   return {
     async *connect(messages, data, abortSignal) {
       // Resolve URL and options if they are functions
@@ -228,7 +363,7 @@ export function fetchHttpStream(
   options:
     | FetchConnectionOptions
     | (() => FetchConnectionOptions | Promise<FetchConnectionOptions>) = {},
-): ConnectionAdapter {
+): ConnectConnectionAdapter {
   return {
     async *connect(messages, data, abortSignal) {
       // Resolve URL and options if they are functions
@@ -301,7 +436,7 @@ export function stream(
     messages: Array<UIMessage> | Array<ModelMessage>,
     data?: Record<string, any>,
   ) => AsyncIterable<StreamChunk>,
-): ConnectionAdapter {
+): ConnectConnectionAdapter {
   return {
     async *connect(messages, data) {
       // Pass messages as-is (UIMessages with parts preserved)
@@ -332,7 +467,7 @@ export function rpcStream(
     messages: Array<UIMessage> | Array<ModelMessage>,
     data?: Record<string, any>,
   ) => AsyncIterable<StreamChunk>,
-): ConnectionAdapter {
+): ConnectConnectionAdapter {
   return {
     async *connect(messages, data) {
       // Pass messages as-is (UIMessages with parts preserved)

packages/typescript/ai-client/src/index.ts

@@ -12,6 +12,7 @@ export type {
   ChatRequestBody,
   InferChatMessages,
   ChatClientState,
+  ConnectionStatus,
   // Multimodal content input type
   MultimodalContent,
 } from './types'
@@ -27,8 +28,10 @@ export {
   fetchHttpStream,
   stream,
   rpcStream,
+  type ConnectConnectionAdapter,
   type ConnectionAdapter,
   type FetchConnectionOptions,
+  type SubscribeConnectionAdapter,
 } from './connection-adapters'
 
 // Re-export message converters from @tanstack/ai

packages/typescript/ai-client/src/types.ts

@@ -36,6 +36,15 @@ export type ToolResultState =
  */
 export type ChatClientState = 'ready' | 'submitted' | 'streaming' | 'error'
 
+/**
+ * Connection lifecycle state for the subscription loop.
+ */
+export type ConnectionStatus =
+  | 'disconnected'
+  | 'connecting'
+  | 'connected'
+  | 'error'
+
 /**
  * Multimodal content input for sending messages with rich media.
  * Allows sending text, images, audio, video, and documents to the LLM.
@@ -178,8 +187,9 @@ export interface ChatClientOptions<
   TTools extends ReadonlyArray<AnyClientTool> = any,
 > {
   /**
-   * Connection adapter for streaming
-   * Use fetchServerSentEvents(), fetchHttpStream(), or stream() to create adapters
+   * Connection adapter for streaming.
+   * Supports mutually exclusive modes: request-response via `connect()`, or
+   * subscribe/send mode via `subscribe()` + `send()`.
    */
   connection: ConnectionAdapter
 
@@ -239,6 +249,25 @@ export interface ChatClientOptions<
    */
   onStatusChange?: (status: ChatClientState) => void
 
+  /**
+   * Callback when subscription lifecycle changes.
+   * This is independent from request lifecycle (`isLoading`, `status`).
+   */
+  onSubscriptionChange?: (isSubscribed: boolean) => void
+
+  /**
+   * Callback when connection lifecycle changes.
+   */
+  onConnectionStatusChange?: (status: ConnectionStatus) => void
+
+  /**
+   * Callback when session generation activity changes.
+   * Derived from stream run events (RUN_STARTED / RUN_FINISHED / RUN_ERROR).
+   * Unlike `onLoadingChange` (request-local), this reflects shared generation
+   * activity visible to all subscribers (e.g. across tabs/devices).
+   */
+  onSessionGeneratingChange?: (isGenerating: boolean) => void
+
   /**
    * Callback when a custom event is received from a server-side tool.
    * Custom events are emitted by tools using `context.emitCustomEvent()` during execution.