@ably/ai-transport 0.0.1 → 0.1.0

package/src/vercel/codec/accumulator.ts

@@ -17,6 +17,7 @@ import type * as AI from 'ai';
 
 import type { DecoderOutput, MessageAccumulator } from '../../core/codec/types.js';
 import { stripUndefined } from '../../utils.js';
+import { toolBase, transitionToolPart } from './tool-transitions.js';
 
 // ---------------------------------------------------------------------------
 // Internal types
@@ -38,15 +39,6 @@ interface ToolPartTracker {
   inputText: string;
 }
 
-/** Fields shared by all DynamicToolUIPart state variants. */
-interface ToolBaseFields {
-  type: 'dynamic-tool';
-  toolName: string;
-  toolCallId: string;
-  title?: string;
-  providerExecuted?: boolean;
-}
-
 /** Bundled per-message state for an in-progress message. */
 interface ActiveMessageState {
   message: AI.UIMessage;
@@ -56,34 +48,6 @@ interface ActiveMessageState {
   streamStatus: Map<string, StreamStatus>;
 }
 
-// ---------------------------------------------------------------------------
-// Tool base helper
-// ---------------------------------------------------------------------------
-
-/**
- * Extract the state-independent base fields for a DynamicToolUIPart.
- * Works with both chunks (tool-input-start, etc.) and existing parts.
- * @param source - Any object containing the required tool identity fields.
- * @param source.toolCallId - The tool call identifier.
- * @param source.toolName - The tool name.
- * @param source.title - Optional display title.
- * @param source.providerExecuted - Whether the provider executed the tool.
- * @returns Base fields shared across all DynamicToolUIPart state variants.
- */
-const toolBase = (source: {
-  toolCallId: string;
-  toolName: string;
-  title?: string;
-  providerExecuted?: boolean;
-}): ToolBaseFields =>
-  stripUndefined({
-    type: 'dynamic-tool' as const,
-    toolCallId: source.toolCallId,
-    toolName: source.toolName,
-    title: source.title,
-    providerExecuted: source.providerExecuted,
-  });
-
 // ---------------------------------------------------------------------------
 // DeltaStreamTracker — manages text or reasoning stream accumulation
 // ---------------------------------------------------------------------------
@@ -172,6 +136,68 @@ class DefaultUIMessageAccumulator implements MessageAccumulator<AI.UIMessageChun
     }
   }
 
+  initMessage(messageId: string, message: AI.UIMessage): void {
+    const existing = this._activeMessages.get(messageId);
+
+    if (existing) {
+      // Already active — sync with the externally updated message.
+      // Replace the message and rebuild tool trackers so the accumulator
+      // reflects updates (e.g. cross-turn amendments applied to the tree)
+      // that happened outside the streaming flow.
+      const cloned = structuredClone(message);
+      const listIdx = this._messageList.indexOf(existing.message);
+      existing.message = cloned;
+      if (listIdx !== -1) {
+        this._messageList[listIdx] = cloned;
+      }
+      existing.toolTrackers = {};
+      for (let i = 0; i < cloned.parts.length; i++) {
+        const part = cloned.parts[i];
+        if (part?.type === 'dynamic-tool') {
+          existing.toolTrackers[part.toolCallId] = { partIndex: i, inputText: '' };
+          existing.streamStatus.set(part.toolCallId, 'finished');
+        }
+      }
+      return;
+    }
+
+    // Not active — create tracking state from the existing message.
+    const cloned = structuredClone(message);
+    const toolTrackers: Record<string, ToolPartTracker> = {};
+    const streamStatus = new Map<string, StreamStatus>();
+
+    for (let i = 0; i < cloned.parts.length; i++) {
+      const part = cloned.parts[i];
+      if (part?.type === 'dynamic-tool') {
+        toolTrackers[part.toolCallId] = { partIndex: i, inputText: '' };
+        streamStatus.set(part.toolCallId, 'finished');
+      }
+    }
+
+    const state: ActiveMessageState = {
+      message: cloned,
+      textStreams: new DeltaStreamTracker('text'),
+      reasoningStreams: new DeltaStreamTracker('reasoning'),
+      toolTrackers,
+      streamStatus,
+    };
+
+    this._activeMessages.set(messageId, state);
+
+    // If this message is already in the list (completed previously),
+    // replace in-place. Otherwise push as a new entry.
+    const existingIdx = this._messageList.findIndex((m) => m.id === message.id);
+    if (existingIdx === -1) {
+      this._messageList.push(state.message);
+    } else {
+      this._messageList[existingIdx] = state.message;
+    }
+  }
+
+  completeMessage(messageId: string): void {
+    this._activeMessages.delete(messageId);
+  }
+
   // -------------------------------------------------------------------------
   // Shared helpers
   // -------------------------------------------------------------------------
@@ -503,48 +529,7 @@ class DefaultUIMessageAccumulator implements MessageAccumulator<AI.UIMessageChun
     const found = this._getToolPart(chunk.toolCallId, state);
     if (!found) return;
 
-    switch (chunk.type) {
-      case 'tool-output-available': {
-        state.message.parts[found.tracker.partIndex] = stripUndefined({
-          ...toolBase(found.part),
-          state: 'output-available' as const,
-          input: found.part.input,
-          output: chunk.output,
-          preliminary: chunk.preliminary,
-        });
-        break;
-      }
-
-      case 'tool-output-error': {
-        state.message.parts[found.tracker.partIndex] = {
-          ...toolBase(found.part),
-          state: 'output-error',
-          input: found.part.input,
-          errorText: chunk.errorText,
-        };
-        break;
-      }
-
-      case 'tool-output-denied': {
-        state.message.parts[found.tracker.partIndex] = {
-          ...toolBase(found.part),
-          state: 'output-denied',
-          input: found.part.input,
-          approval: { id: '', approved: false },
-        };
-        break;
-      }
-
-      case 'tool-approval-request': {
-        state.message.parts[found.tracker.partIndex] = {
-          ...toolBase(found.part),
-          state: 'approval-requested',
-          input: found.part.input,
-          approval: { id: chunk.approvalId },
-        };
-        break;
-      }
-    }
+    state.message.parts[found.tracker.partIndex] = transitionToolPart(found.part, chunk);
   }
 
   // -------------------------------------------------------------------------

package/src/vercel/codec/decoder.ts

@@ -14,7 +14,7 @@
 import type * as Ably from 'ably';
 import type * as AI from 'ai';
 
-import { HEADER_ROLE, HEADER_TURN_ID } from '../../constants.js';
+import { HEADER_DISCRETE, HEADER_ROLE, HEADER_TURN_ID } from '../../constants.js';
 import type { DecoderCore, DecoderCoreHooks, DecoderCoreOptions } from '../../core/codec/decoder.js';
 import { createDecoderCore, eventOutput } from '../../core/codec/decoder.js';
 import type { LifecycleTracker } from '../../core/codec/lifecycle-tracker.js';
@@ -277,7 +277,8 @@ const decodeFinish = (r: VercelHeaderReader, turnId: string, lifecycle: Lifecycl
   );
 };
 
-const decodeError = (data: unknown): Out[] => {
+const decodeError = (data: unknown, turnId: string, lifecycle: LifecycleTracker<AI.UIMessageChunk>): Out[] => {
+  lifecycle.clearScope(turnId);
   const errorText = typeof data === 'string' ? data : '';
   return event({ type: 'error', errorText });
 };
@@ -483,14 +484,16 @@ const decodeDiscreteMessage = (input: MessagePayload): Out[] => {
 
 /**
  * Whether a message name represents a discrete message part (written by writeMessages)
- * rather than a streaming lifecycle event. Discrete message parts carry x-ably-role
- * and encode a single UIMessage part each.
+ * rather than a streaming lifecycle event. Distinguished by the `x-ably-discrete` header
+ * which {@link publishDiscreteBatch} sets on batch-published message payloads. Lifecycle
+ * events published via {@link publishDiscrete} (including streaming `data-*` chunks)
+ * do not carry this header.
  * @param name - The Ably message name to check.
- * @param headers - The Ably message headers to inspect for role presence.
+ * @param headers - The Ably message headers to inspect for discrete marker presence.
  * @returns True if this is a discrete message part, false if it's a lifecycle event.
  */
 const isDiscreteMessagePart = (name: string, headers: Record<string, string>): boolean =>
-  (name === 'text' || name === 'file' || isDataEventName(name)) && HEADER_ROLE in headers;
+  (name === 'text' || name === 'file' || isDataEventName(name)) && HEADER_DISCRETE in headers;
 
 const decodeDiscretePayload = (input: MessagePayload, lifecycle: LifecycleTracker<AI.UIMessageChunk>): Out[] => {
   const h = input.headers ?? {};
@@ -498,7 +501,7 @@ const decodeDiscretePayload = (input: MessagePayload, lifecycle: LifecycleTracke
   const turnId = h[HEADER_TURN_ID] ?? '';
 
   // Discrete message parts from writeMessages (user messages, history entries).
-  // Distinguished from lifecycle events by the presence of x-ably-role.
+  // Distinguished from lifecycle events by the presence of x-ably-discrete.
   if (isDiscreteMessagePart(input.name, h)) {
     return decodeDiscreteMessage(input);
   }
@@ -521,7 +524,7 @@ const decodeDiscretePayload = (input: MessagePayload, lifecycle: LifecycleTracke
       return decodeFinish(r, turnId, lifecycle);
     }
     case 'error': {
-      return decodeError(input.data);
+      return decodeError(input.data, turnId, lifecycle);
     }
     case 'abort': {
       return decodeAbort(input.data, turnId, lifecycle);

package/src/vercel/codec/encoder.ts

@@ -40,16 +40,78 @@ import type { ChannelWriter, MessagePayload, StreamEncoder, WriteOptions } from
 import { ErrorCode, errorInfoIs } from '../../errors.js';
 import { headerWriter } from '../../utils.js';
 
+// ---------------------------------------------------------------------------
+// Discrete event payload builder
+// ---------------------------------------------------------------------------
+
+/**
+ * Build a MessagePayload for discrete (non-streaming) event types.
+ * Used by both `writeEvent` and `appendEvent` for tool output events,
+ * content parts, and data-* custom chunks.
+ * @param chunk - The UI message chunk to encode.
+ * @returns The message payload for publishing to the channel.
+ */
+const buildDiscretePayload = (chunk: AI.UIMessageChunk): MessagePayload => {
+  switch (chunk.type) {
+    case 'tool-output-available': {
+      const h = headerWriter()
+        .str('toolCallId', chunk.toolCallId)
+        .bool('dynamic', chunk.dynamic)
+        .bool('providerExecuted', chunk.providerExecuted)
+        .bool('preliminary', chunk.preliminary)
+        .build();
+      return { name: 'tool-output-available', data: { output: chunk.output }, headers: h };
+    }
+
+    case 'tool-output-error': {
+      const h = headerWriter()
+        .str('toolCallId', chunk.toolCallId)
+        .bool('dynamic', chunk.dynamic)
+        .bool('providerExecuted', chunk.providerExecuted)
+        .build();
+      return { name: 'tool-output-error', data: { errorText: chunk.errorText }, headers: h };
+    }
+
+    case 'tool-approval-request': {
+      const h = headerWriter().str('toolCallId', chunk.toolCallId).str('approvalId', chunk.approvalId).build();
+      return { name: 'tool-approval-request', data: '', headers: h };
+    }
+
+    case 'tool-output-denied': {
+      const h = headerWriter().str('toolCallId', chunk.toolCallId).build();
+      return { name: 'tool-output-denied', data: '', headers: h };
+    }
+
+    default: {
+      if (chunk.type.startsWith('data-')) {
+        // CAST: data-* chunks always have id, transient, and data fields per AI SDK types.
+        // TypeScript can't narrow the template literal union in a default case.
+        const dataChunk = chunk as Extract<AI.UIMessageChunk, { type: `data-${string}` }>;
+        const h = headerWriter().str('id', dataChunk.id).bool('transient', dataChunk.transient).build();
+        const ephemeral = dataChunk.transient === true;
+        return { name: chunk.type, data: dataChunk.data, headers: h, ephemeral };
+      }
+      throw new Ably.ErrorInfo(
+        `unable to write event; unsupported chunk type '${chunk.type}'`,
+        ErrorCode.InvalidArgument,
+        400,
+      );
+    }
+  }
+};
+
 // ---------------------------------------------------------------------------
 // Default implementation
 // ---------------------------------------------------------------------------
 
 class DefaultUIMessageEncoder implements StreamEncoder<AI.UIMessageChunk, AI.UIMessage> {
   private readonly _core: EncoderCore;
+  private readonly _messageId: string | undefined;
   private _aborted = false;
 
   constructor(writer: ChannelWriter, options: EncoderCoreOptions = {}) {
     this._core = createEncoderCore(writer, options);
+    this._messageId = options.messageId;
   }
 
   async appendEvent(chunk: AI.UIMessageChunk, perWrite?: WriteOptions): Promise<void> {
@@ -144,7 +206,7 @@ class DefaultUIMessageEncoder implements StreamEncoder<AI.UIMessageChunk, AI.UIM
 
       case 'start': {
         const h = headerWriter()
-          .str('messageId', chunk.messageId)
+          .str('messageId', chunk.messageId ?? this._messageId)
           .json('messageMetadata', chunk.messageMetadata)
           .build();
         await this._core.publishDiscrete({ name: 'start', data: '', headers: h }, perWrite);
@@ -204,44 +266,11 @@ class DefaultUIMessageEncoder implements StreamEncoder<AI.UIMessageChunk, AI.UIM
         break;
       }
 
-      case 'tool-output-available': {
-        const h = headerWriter()
-          .str('toolCallId', chunk.toolCallId)
-          .bool('dynamic', chunk.dynamic)
-          .bool('providerExecuted', chunk.providerExecuted)
-          .bool('preliminary', chunk.preliminary)
-          .build();
-        await this._core.publishDiscrete({
-          name: 'tool-output-available',
-          data: { output: chunk.output },
-          headers: h,
-        });
-        break;
-      }
-
-      case 'tool-output-error': {
-        const h = headerWriter()
-          .str('toolCallId', chunk.toolCallId)
-          .bool('dynamic', chunk.dynamic)
-          .bool('providerExecuted', chunk.providerExecuted)
-          .build();
-        await this._core.publishDiscrete({
-          name: 'tool-output-error',
-          data: { errorText: chunk.errorText },
-          headers: h,
-        });
-        break;
-      }
-
-      case 'tool-approval-request': {
-        const h = headerWriter().str('toolCallId', chunk.toolCallId).str('approvalId', chunk.approvalId).build();
-        await this._core.publishDiscrete({ name: 'tool-approval-request', data: '', headers: h }, perWrite);
-        break;
-      }
-
+      case 'tool-output-available':
+      case 'tool-output-error':
+      case 'tool-approval-request':
       case 'tool-output-denied': {
-        const h = headerWriter().str('toolCallId', chunk.toolCallId).build();
-        await this._core.publishDiscrete({ name: 'tool-output-denied', data: '', headers: h }, perWrite);
+        await this._core.publishDiscrete(buildDiscretePayload(chunk), perWrite);
         break;
       }
 
@@ -298,22 +327,7 @@ class DefaultUIMessageEncoder implements StreamEncoder<AI.UIMessageChunk, AI.UIM
   }
 
   async writeEvent(chunk: AI.UIMessageChunk, perWrite?: WriteOptions): Promise<Ably.PublishResult> {
-    if (!chunk.type.startsWith('data-')) {
-      throw new Ably.ErrorInfo(
-        `unable to write event; only data-* chunk types are supported, got '${chunk.type}'`,
-        ErrorCode.InvalidArgument,
-        400,
-      );
-    }
-    const h = headerWriter()
-      .str('id', 'id' in chunk ? chunk.id : undefined)
-      .bool('transient', 'transient' in chunk ? chunk.transient : undefined)
-      .build();
-    const ephemeral = 'transient' in chunk && chunk.transient === true;
-    return this._core.publishDiscrete(
-      { name: chunk.type, data: 'data' in chunk ? chunk.data : undefined, headers: h, ephemeral },
-      perWrite,
-    );
+    return this._core.publishDiscrete(buildDiscretePayload(chunk), perWrite);
   }
 
   async writeMessages(messages: AI.UIMessage[], perWrite?: WriteOptions): Promise<Ably.PublishResult> {

package/src/vercel/codec/index.ts

@@ -30,8 +30,6 @@ export const UIMessageCodec: Codec<AI.UIMessageChunk, AI.UIMessage> = {
   createDecoder,
   createAccumulator,
 
-  getMessageKey: (message: AI.UIMessage): string => message.id,
-
   isTerminal: (event: AI.UIMessageChunk): boolean =>
     event.type === 'finish' || event.type === 'error' || event.type === 'abort',
 };

package/src/vercel/codec/tool-transitions.ts

@@ -0,0 +1,122 @@
+/**
+ * Shared tool part transition logic for the Vercel AI SDK codec.
+ *
+ * Extracted from the accumulator so the tool output state transition logic
+ * lives in one place, reusable by the accumulator and any other callers.
+ */
+
+import type * as AI from 'ai';
+
+import { stripUndefined } from '../../utils.js';
+
+// ---------------------------------------------------------------------------
+// Tool output chunk type guard
+// ---------------------------------------------------------------------------
+
+/** The set of UIMessageChunk types that represent tool output transitions. */
+export type ToolOutputChunk = Extract<
+  AI.UIMessageChunk,
+  { type: 'tool-output-available' | 'tool-output-error' | 'tool-output-denied' | 'tool-approval-request' }
+>;
+
+/**
+ * Whether a UIMessageChunk is a tool output transition event.
+ * @param chunk - The chunk to test.
+ * @returns True if the chunk is a tool output transition type.
+ */
+export const isToolOutputChunk = (chunk: AI.UIMessageChunk): chunk is ToolOutputChunk =>
+  chunk.type === 'tool-output-available' ||
+  chunk.type === 'tool-output-error' ||
+  chunk.type === 'tool-output-denied' ||
+  chunk.type === 'tool-approval-request';
+
+// ---------------------------------------------------------------------------
+// Tool base helper
+// ---------------------------------------------------------------------------
+
+/** Fields shared by all DynamicToolUIPart state variants. */
+interface ToolBaseFields {
+  type: 'dynamic-tool';
+  toolName: string;
+  toolCallId: string;
+  title?: string;
+  providerExecuted?: boolean;
+}
+
+/**
+ * Extract the state-independent base fields for a DynamicToolUIPart.
+ * Works with both chunks (tool-input-start, etc.) and existing parts.
+ * @param source - Any object containing the required tool identity fields.
+ * @param source.toolCallId - The tool call identifier.
+ * @param source.toolName - The tool name.
+ * @param source.title - Optional display title.
+ * @param source.providerExecuted - Whether the provider executed the tool.
+ * @returns Base fields shared across all DynamicToolUIPart state variants.
+ */
+export const toolBase = (source: {
+  toolCallId: string;
+  toolName: string;
+  title?: string;
+  providerExecuted?: boolean;
+}): ToolBaseFields =>
+  stripUndefined({
+    type: 'dynamic-tool' as const,
+    toolCallId: source.toolCallId,
+    toolName: source.toolName,
+    title: source.title,
+    providerExecuted: source.providerExecuted,
+  });
+
+// ---------------------------------------------------------------------------
+// Tool part transition
+// ---------------------------------------------------------------------------
+
+/**
+ * Transition a DynamicToolUIPart to a new state based on a tool output chunk.
+ * Pure function — does not mutate the input part.
+ * @param part - The existing tool part to transition.
+ * @param chunk - The tool output chunk describing the transition.
+ * @returns A new DynamicToolUIPart in the target state.
+ */
+export const transitionToolPart = (part: AI.DynamicToolUIPart, chunk: ToolOutputChunk): AI.DynamicToolUIPart => {
+  const base = toolBase(part);
+
+  switch (chunk.type) {
+    case 'tool-output-available': {
+      return stripUndefined({
+        ...base,
+        state: 'output-available' as const,
+        input: part.input,
+        output: chunk.output,
+        preliminary: chunk.preliminary,
+      });
+    }
+
+    case 'tool-output-error': {
+      return {
+        ...base,
+        state: 'output-error',
+        input: part.input,
+        errorText: chunk.errorText,
+      };
+    }
+
+    case 'tool-output-denied': {
+      return {
+        ...base,
+        state: 'output-denied',
+        input: part.input,
+        approval: { id: '', approved: false },
+      };
+    }
+
+    case 'tool-approval-request': {
+      return {
+        ...base,
+        state: 'approval-requested',
+        input: part.input,
+        approval: { id: chunk.approvalId },
+      };
+    }
+  }
+};

package/src/vercel/index.ts

@@ -10,3 +10,20 @@ export type {
   VercelServerTransportOptions,
 } from './transport/index.js';
 export { createChatTransport, createClientTransport, createServerTransport } from './transport/index.js';
+
+// Server-side tool result merge helper
+export { applyToolEventsToHistory } from './tool-events.js';
+
+// Server-side tool approval helpers
+export type {
+  PrepareApprovalTurnOptions,
+  PrepareApprovalTurnResult,
+  StreamResponseWithApprovalRedirectOptions,
+  ToolApprovalDecision,
+} from './tool-approvals.js';
+export {
+  applyToolApprovalsToHistory,
+  extractApprovalDecisionsFromHistory,
+  prepareApprovalTurn,
+  streamResponseWithApprovalRedirect,
+} from './tool-approvals.js';

package/src/vercel/react/contexts/chat-transport-context.ts

@@ -0,0 +1,40 @@
+import type * as Ably from 'ably';
+import type * as AI from 'ai';
+import { createContext } from 'react';
+
+import type { ClientTransport } from '../../../core/transport/types.js';
+import type { ChatTransport } from '../../transport/chat-transport.js';
+
+/**
+ * A single entry in the chat transport registry, holding both the
+ * underlying {@link ClientTransport} and the {@link ChatTransport} wrapping it.
+ */
+export interface ChatTransportSlot {
+  /** The underlying client transport used to create the chat transport. */
+  readonly transport: ClientTransport<AI.UIMessageChunk, AI.UIMessage>;
+  /** Construction error from the underlying {@link ClientTransport}, or `undefined` on success. */
+  readonly transportError: Ably.ErrorInfo | undefined;
+  /** The chat transport adapter for use with Vercel's useChat hook. */
+  readonly chatTransport: ChatTransport;
+}
+
+/**
+ * The shape of the single {@link ChatTransportContext} value.
+ * Combines the nearest slot with the full registry in one context object.
+ */
+export interface ChatTransportContextValue {
+  /** The slot from the nearest {@link ChatTransportProvider} in the tree. */
+  readonly nearest: ChatTransportSlot | undefined;
+  /** All registered slots, keyed by channelName. */
+  readonly providers: Readonly<Record<string, ChatTransportSlot>>;
+}
+
+/**
+ * Context that carries both the nearest {@link ChatTransportSlot} and the full registry of
+ * registered slots keyed by channelName. Populated by {@link ChatTransportProvider};
+ * read by {@link useChatTransport}.
+ */
+export const ChatTransportContext = createContext<ChatTransportContextValue>({
+  nearest: undefined,
+  providers: {},
+});