@ably/ai-transport 0.1.0 → 0.3.0

package/src/vercel/codec/reducer-state.ts

@@ -0,0 +1,169 @@
+/**
+ * Shared reducer state: the projection shape, its internal tracker types,
+ * `init`, and the message/tracker lookup helpers the per-concern fold modules
+ * build on. This module is the base of the reducer's import DAG — the fold
+ * modules depend on it; it depends on none of them.
+ */
+
+import type * as AI from 'ai';
+
+import type { CodecMessage } from '../../core/codec/index.js';
+
+// ---------------------------------------------------------------------------
+// Internal tracker state
+// ---------------------------------------------------------------------------
+
+/**
+ * Tracks an in-progress tool part within a UIMessage. Text and reasoning
+ * parts don't need this — we write to them directly via partIndex. Tool
+ * parts need an extra `inputText` buffer because deltas arrive as raw
+ * JSON fragments that must be accumulated before parsing.
+ */
+export interface ToolPartTracker {
+  /** Index in the message's parts array. */
+  partIndex: number;
+  /** Accumulated streaming input text (for JSON parsing on completion). */
+  inputText: string;
+}
+
+/** Per-codecMessageId tracking state for in-progress streams within a UIMessage. */
+export interface MessageTrackers {
+  /** Text stream id → partIndex. */
+  text: Map<string, number>;
+  /** Reasoning stream id → partIndex. */
+  reasoning: Map<string, number>;
+  /** Tool call id → tracker. */
+  tools: Map<string, ToolPartTracker>;
+}
+
+// ---------------------------------------------------------------------------
+// Projection
+// ---------------------------------------------------------------------------
+
+/**
+ * The per-Run state produced by the Vercel codec's reducer.
+ *
+ * The SDK reads only `messages` (via `Codec.getMessages`). The remaining
+ * fields are internal to the reducer; they happen to live on the
+ * projection because the projection is the only thing the reducer can
+ * carry from fold to fold (it has no instance state).
+ */
+export interface VercelProjection {
+  /**
+   * UIMessages produced or modified in this Run, in publication order,
+   * each paired with its codec-message-id. The reducer correlates strictly
+   * on `codecMessageId`; `message.id` is preserved verbatim from the source
+   * (the AI SDK stream's `start.messageId` for assistants, the caller's id
+   * for user messages) and is never used as an identity key.
+   */
+  messages: CodecMessage<AI.UIMessage>[];
+  /** Per-codecMessageId tracker state for streamed parts. Internal — do not access. */
+  trackers: Map<string, MessageTrackers>;
+  /**
+   * Tool-resolution events that arrived before any assistant in this
+   * projection had a matching `toolCallId`. Re-evaluated on every
+   * subsequent fold so that an out-of-order tool output is folded as
+   * soon as the corresponding assistant lands.
+   */
+  pendingToolResolutions: PendingToolResolution[];
+}
+
+/**
+ * A buffered tool resolution waiting for its assistant message to arrive.
+ * The reducer scans pending entries after every successful fold so an
+ * out-of-order tool output is promoted as soon as the matching assistant
+ * is added to the projection.
+ */
+export interface PendingToolResolution {
+  /** The codec-message-id of the assistant the resolution targets. */
+  targetCodecMessageId: string;
+  /** Tool call this resolution targets. */
+  toolCallId: string;
+  /** Variant of the tool-resolution used to transition the assistant's tool part. */
+  resolution:
+    | { kind: 'tool-result'; output: unknown }
+    | { kind: 'tool-result-error'; message: string }
+    | { kind: 'tool-approval-response'; approved: boolean; reason?: string };
+}
+
+/** A located `dynamic-tool` part with its owning message and tracker. */
+export interface OwnerLookup {
+  /** The message owning the tool part. */
+  message: AI.UIMessage;
+  /** The tracker pointing at the part's index. */
+  tracker: ToolPartTracker;
+  /** The resolved `dynamic-tool` part itself. */
+  part: AI.DynamicToolUIPart;
+}
+
+// ---------------------------------------------------------------------------
+// init
+// ---------------------------------------------------------------------------
+
+/**
+ * Build an empty initial projection.
+ * @returns A fresh VercelProjection with no messages and no tracker state.
+ */
+export const init = (): VercelProjection => ({
+  messages: [],
+  trackers: new Map(),
+  pendingToolResolutions: [],
+});
+
+// ---------------------------------------------------------------------------
+// Message + tracker helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Resolve the assistant message for a codec-message-id, creating an empty
+ * placeholder when none exists yet.
+ * @param state - Projection to read or extend.
+ * @param codecMessageId - The codec-message-id to resolve.
+ * @returns The existing or newly-seeded UIMessage for that id.
+ */
+export const ensureMessage = (state: VercelProjection, codecMessageId: string): AI.UIMessage => {
+  let entry = state.messages.find((e) => e.codecMessageId === codecMessageId);
+  if (!entry) {
+    // No source id seen yet — seed the domain `message.id` with the
+    // codec-message-id as a fallback. The `start` chunk overwrites it with
+    // the stream's `messageId` when the stream provides one.
+    entry = { codecMessageId, message: { id: codecMessageId, role: 'assistant', parts: [] } };
+    state.messages.push(entry);
+  }
+  return entry.message;
+};
+
+/**
+ * Resolve the stream trackers for a codec-message-id, creating empty maps
+ * when none exist yet.
+ * @param state - Projection to read or extend.
+ * @param messageId - The codec-message-id whose trackers to resolve.
+ * @returns The existing or newly-created tracker maps for that id.
+ */
+export const ensureTrackers = (state: VercelProjection, messageId: string): MessageTrackers => {
+  let trackers = state.trackers.get(messageId);
+  if (!trackers) {
+    trackers = { text: new Map(), reasoning: new Map(), tools: new Map() };
+    state.trackers.set(messageId, trackers);
+  }
+  return trackers;
+};
+
+/**
+ * Resolve the `dynamic-tool` part tracked for a toolCallId within a message.
+ * @param message - The message whose parts to read.
+ * @param trackers - The message's tracker maps.
+ * @param toolCallId - The tool call to resolve.
+ * @returns The tracker and part, or `undefined` if untracked or the part is not a dynamic-tool.
+ */
+export const getToolPart = (
+  message: AI.UIMessage,
+  trackers: MessageTrackers,
+  toolCallId: string,
+): { tracker: ToolPartTracker; part: AI.DynamicToolUIPart } | undefined => {
+  const tracker = trackers.tools.get(toolCallId);
+  if (!tracker) return undefined;
+  const part = message.parts[tracker.partIndex];
+  if (part?.type !== 'dynamic-tool') return undefined;
+  return { tracker, part };
+};

package/src/vercel/codec/reducer.ts

@@ -0,0 +1,191 @@
+/**
+ * Vercel AI SDK reducer.
+ *
+ * Pure `(init, fold)` over the `VercelInput | VercelOutput` union. Folds
+ * input variants (user-message, tool-result, tool-result-error,
+ * tool-approval-response) and `UIMessageChunk` outputs into a
+ * VercelProjection holding `UIMessage[]` plus internal stream-tracker
+ * state.
+ *
+ * The reducer is stateless: every fold is `(state, event, meta) → state'`,
+ * with no instance state. Mutation in place is allowed — the projection
+ * is single-owner.
+ *
+ * The reducer does not dedup or reorder. The transport sequences events
+ * canonically — ascending by wire serial across messages, in decode order
+ * within a wire — and delivers each exactly once, so the reducer folds
+ * unconditionally. Last-writer-wins for events competing over the same
+ * logical state (e.g. two `tool-output-available` for one `toolCallId`)
+ * falls out of fold order: the highest-serial event folds last.
+ *
+ * Client-published tool resolutions (`ToolResult`, `ToolResultError`,
+ * `ToolApprovalResponse`) carry `codecMessageId` targeting the assistant
+ * they amend; the reducer applies the resolution onto that assistant's
+ * `dynamic-tool` part directly. If the assistant has not yet arrived in
+ * the projection (out-of-order delivery), the resolution is buffered in
+ * `pendingToolResolutions` and re-evaluated on each subsequent fold.
+ *
+ * This file is the reducer's public facade and dispatch: `init`,
+ * `getMessages`, `fold`, and the output-chunk router. The per-concern fold
+ * logic lives in the sibling `fold-*` modules over a shared `reducer-state`
+ * base; the import graph is an acyclic DAG rooted here.
+ */
+
+import type * as AI from 'ai';
+
+import type { CodecEvent, CodecMessage, ReducerMeta } from '../../core/codec/index.js';
+import type { VercelInput, VercelOutput } from './events.js';
+import { foldContentPart } from './fold-content.js';
+import { foldDataPart } from './fold-data.js';
+import {
+  foldClientToolResult,
+  foldClientToolResultError,
+  foldToolApprovalResponse,
+  foldUserMessage,
+  retryPendingResolutions,
+} from './fold-input.js';
+import { foldLifecycle } from './fold-lifecycle.js';
+import { foldTextOrReasoning } from './fold-text.js';
+import { foldToolInput } from './fold-tool-input.js';
+import { foldToolOutput } from './fold-tool-output.js';
+import type { VercelProjection } from './reducer-state.js';
+
+// ---------------------------------------------------------------------------
+// fold
+// ---------------------------------------------------------------------------
+
+/**
+ * Fold one input or output event into the projection. Mutates and returns
+ * `state`.
+ *
+ * The transport invokes `fold` exactly once per event, in canonical order,
+ * so the reducer folds unconditionally — no dedup or high-water-mark here.
+ * Competing events resolve by order (the highest-serial event folds last
+ * and wins). Orphan events (e.g. tool-output for an unknown toolCallId) are
+ * dropped silently inside the per-variant fold helpers.
+ * @param state - Projection to fold into (may be mutated in place).
+ * @param event - Input or output event to fold.
+ * @param meta - Transport-derived metadata (serial, optional messageId).
+ * @returns The same projection reference, possibly mutated.
+ */
+export const fold = (
+  state: VercelProjection,
+  event: CodecEvent<VercelInput, VercelOutput>,
+  meta: ReducerMeta,
+): VercelProjection => {
+  if (event.direction === 'input') {
+    const input = event.event;
+    switch (input.kind) {
+      case 'user-message': {
+        foldUserMessage(state, input.message, meta);
+        break;
+      }
+      case 'regenerate': {
+        // Regenerate input — wire-only signal. Carries no projection state;
+        // the agent reads `target` / `parent` from the wire headers via
+        // the input-event lookup path. No fold work to do here.
+        break;
+      }
+      case 'tool-result': {
+        foldClientToolResult(state, input);
+        break;
+      }
+      case 'tool-result-error': {
+        foldClientToolResultError(state, input);
+        break;
+      }
+      case 'tool-approval-response': {
+        foldToolApprovalResponse(state, input);
+        break;
+      }
+    }
+  } else {
+    foldChunk(state, event.event, meta);
+  }
+
+  // Re-evaluate pending tool resolutions in case the just-folded event
+  // produced the assistant they were waiting on. Cheap when the list is
+  // empty (the common case).
+  if (state.pendingToolResolutions.length > 0) {
+    retryPendingResolutions(state);
+  }
+
+  return state;
+};
+
+// ---------------------------------------------------------------------------
+// UIMessageChunk dispatch
+// ---------------------------------------------------------------------------
+
+const foldChunk = (state: VercelProjection, chunk: VercelOutput, meta: ReducerMeta): VercelProjection => {
+  const messageId = meta.messageId;
+  if (messageId === undefined) {
+    // Without a target codec-message-id, a chunk has nowhere to land. Drop.
+    return state;
+  }
+
+  switch (chunk.type) {
+    case 'start':
+    case 'start-step':
+    case 'finish-step':
+    case 'finish':
+    case 'abort':
+    case 'error':
+    case 'message-metadata': {
+      return foldLifecycle(state, chunk, messageId);
+    }
+
+    case 'text-start':
+    case 'text-delta':
+    case 'text-end':
+    case 'reasoning-start':
+    case 'reasoning-delta':
+    case 'reasoning-end': {
+      return foldTextOrReasoning(state, chunk, messageId);
+    }
+
+    case 'tool-input-start':
+    case 'tool-input-delta':
+    case 'tool-input-available':
+    case 'tool-input-error': {
+      return foldToolInput(state, chunk, messageId);
+    }
+
+    case 'tool-output-available':
+    case 'tool-output-error':
+    case 'tool-output-denied':
+    case 'tool-approval-request': {
+      return foldToolOutput(state, chunk, messageId);
+    }
+
+    case 'file':
+    case 'source-url':
+    case 'source-document': {
+      return foldContentPart(state, chunk, messageId);
+    }
+
+    default: {
+      if (chunk.type.startsWith('data-')) {
+        return foldDataPart(state, chunk, messageId);
+      }
+      return state;
+    }
+  }
+};
+
+// ---------------------------------------------------------------------------
+// getMessages
+// ---------------------------------------------------------------------------
+
+/**
+ * Extract the UIMessage list from a projection, each paired with its
+ * codec-message-id. Client-published tool resolutions amend existing
+ * assistants in place via `kind: 'tool-result'` etc. — they never
+ * materialise as their own UIMessage in the projection, so no filtering is
+ * needed here.
+ * @param projection - Projection produced by `init` + repeated `fold` calls.
+ * @returns The visible messages with their codec-message-ids, in publication order.
+ */
+export const getMessages = (projection: VercelProjection): CodecMessage<AI.UIMessage>[] => projection.messages;
+
+export { init, type VercelProjection } from './reducer-state.js';

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

@@ -1,8 +1,8 @@
 /**
  * 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.
+ * Keeps the tool output state transition logic in one place, reusable by the
+ * Vercel codec reducer and any other callers.
  */
 
 import type * as AI from 'ai';
@@ -10,7 +10,7 @@ import type * as AI from 'ai';
 import { stripUndefined } from '../../utils.js';
 
 // ---------------------------------------------------------------------------
-// Tool output chunk type guard
+// Tool output chunk type
 // ---------------------------------------------------------------------------
 
 /** The set of UIMessageChunk types that represent tool output transitions. */
@@ -19,17 +19,6 @@ export type ToolOutputChunk = Extract<
   { 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
 // ---------------------------------------------------------------------------

package/src/vercel/codec/wire-data.ts

@@ -0,0 +1,64 @@
+/**
+ * Wire-data shapes and runtime guards for the tool payloads whose `data`
+ * envelope is JSON-parsed from the network (a trust boundary). The guards
+ * validate the typed envelope fields; tool-defined `output`/`input` stay
+ * unconstrained. Shared by the output and input descriptor tables.
+ */
+
+/** Wire format for the agent-side `tool-input-error` chunk data payload. */
+export interface ToolInputErrorWireData {
+  errorText?: string;
+  input?: unknown;
+}
+
+/** Wire format for the `tool-output-available` (agent) / `tool-result` (client) data payload. */
+export interface ToolOutputAvailableWireData {
+  output?: unknown;
+}
+
+/** Wire format for the agent-side `tool-output-error` chunk data payload. */
+export interface AgentToolOutputErrorWireData {
+  errorText?: string;
+}
+
+/** Wire format for the client-side `tool-result-error` input data payload. */
+export interface ClientToolResultErrorWireData {
+  message?: string;
+}
+
+// Narrow JSON-parsed wire data to a record. The encoder is expected to publish
+// an object for these payloads, but a malformed publish could carry a primitive
+// or null — callers fall back to field defaults when these guards reject.
+const isRecord = (data: unknown): data is Record<string, unknown> => typeof data === 'object' && data !== null;
+
+// Validate that `data` is a record whose named field is absent or a string. The
+// optional-string check for the typed error fields below lives here once so the
+// guards can't drift. No `as` needed: `isRecord` narrows `data` to a record, so
+// string-key indexing is well-typed.
+const isRecordWithOptionalString = (data: unknown, key: string): boolean =>
+  isRecord(data) && (data[key] === undefined || typeof data[key] === 'string');
+
+// Validates the typed `errorText` field; `input` is tool-defined and
+// intentionally left unconstrained.
+/**
+ * Coerce wire `data` to a string, falling back to `''` for any non-string
+ * payload — the defensive read for descriptors whose data is plain text.
+ * @param data - The inbound wire data.
+ * @returns The string payload, or `''` when the data is not a string.
+ */
+export const asString = (data: unknown): string => (typeof data === 'string' ? data : '');
+
+export const isToolInputErrorWireData = (data: unknown): data is ToolInputErrorWireData =>
+  isRecordWithOptionalString(data, 'errorText');
+
+// The sole field `output` is tool-defined and intentionally unconstrained, so
+// this asserts only that the payload is an object envelope.
+export const isToolOutputAvailableWireData = (data: unknown): data is ToolOutputAvailableWireData => isRecord(data);
+
+// Validates the typed `errorText` field.
+export const isAgentToolOutputErrorWireData = (data: unknown): data is AgentToolOutputErrorWireData =>
+  isRecordWithOptionalString(data, 'errorText');
+
+// Validates the typed `message` field.
+export const isClientToolResultErrorWireData = (data: unknown): data is ClientToolResultErrorWireData =>
+  isRecordWithOptionalString(data, 'message');

package/src/vercel/index.ts

@@ -1,4 +1,5 @@
 // Vercel AI SDK codec
+export type { VercelInput, VercelOutput, VercelProjection } from './codec/index.js';
 export { UIMessageCodec } from './codec/index.js';
 
 // Vercel AI SDK transport wrappers (pre-bound to UIMessageCodec)
@@ -6,24 +7,11 @@ export type {
   ChatTransport,
   ChatTransportOptions,
   SendMessagesRequestContext,
-  VercelClientTransportOptions,
-  VercelServerTransportOptions,
+  VercelAgentSessionOptions,
+  VercelClientSessionOptions,
 } from './transport/index.js';
-export { createChatTransport, createClientTransport, createServerTransport } from './transport/index.js';
+export { createAgentSession, createChatTransport, createClientSession } 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';
+// Vercel-shaped helpers
+export type { VercelRunOutcome } from './run-end-reason.js';
+export { vercelRunOutcome } from './run-end-reason.js';

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

@@ -2,18 +2,19 @@ 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 { ClientSession } from '../../../core/transport/types.js';
+import type { VercelInput, VercelOutput, VercelProjection } from '../../codec/index.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.
+ * underlying {@link ClientSession} 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 underlying client session used to create the chat transport. */
+  readonly session: ClientSession<VercelInput, VercelOutput, VercelProjection, AI.UIMessage>;
+  /** Construction error from the underlying {@link ClientSession}, or `undefined` on success. */
+  readonly sessionError?: Ably.ErrorInfo | undefined;
   /** The chat transport adapter for use with Vercel's useChat hook. */
   readonly chatTransport: ChatTransport;
 }
@@ -22,7 +23,7 @@ export interface ChatTransportSlot {
  * The shape of the single {@link ChatTransportContext} value.
  * Combines the nearest slot with the full registry in one context object.
  */
-export interface ChatTransportContextValue {
+interface ChatTransportContextValue {
   /** The slot from the nearest {@link ChatTransportProvider} in the tree. */
   readonly nearest: ChatTransportSlot | undefined;
   /** All registered slots, keyed by channelName. */