@ably/ai-transport 0.1.0 → 0.3.0

package/src/vercel/codec/events.ts

@@ -0,0 +1,85 @@
+/**
+ * Vercel codec input/output unions.
+ *
+ * The codec splits cleanly along the protocol's `ai-input` / `ai-output`
+ * wire seam:
+ *
+ * - **`VercelOutput`** = `AI.UIMessageChunk` — the AI SDK's streamed-output
+ *   domain model, published by the agent on `ai-output`.
+ * - **`VercelInput`** = a discriminated union of the SDK's well-known
+ *   input shapes — published by the client on `ai-input`. The Vercel
+ *   codec has no codec-local input variants today: every variant comes
+ *   from `@ably/ai-transport`'s well-known set.
+ */
+
+import type * as AI from 'ai';
+
+import type {
+  Regenerate,
+  ToolApprovalResponse,
+  ToolResult,
+  ToolResultError,
+  UserMessage,
+} from '../../core/codec/index.js';
+
+// ---------------------------------------------------------------------------
+// Domain payloads
+//
+// The core well-known tool variants are domain-independent: the Vercel
+// layer supplies the concrete payload shapes. Tool outputs are inherently
+// tool-defined, so `output` stays `unknown` — but confined here, never in
+// the core.
+// ---------------------------------------------------------------------------
+
+/** Vercel domain payload for a {@link ToolResult}. */
+export interface VercelToolResultPayload {
+  /** The tool call this result corresponds to. */
+  toolCallId: string;
+  /** The tool's output value. Tool-defined shape. */
+  output: unknown;
+}
+
+/** Vercel domain payload for a {@link ToolResultError}. */
+export interface VercelToolResultErrorPayload {
+  /** The tool call this error corresponds to. */
+  toolCallId: string;
+  /** Human-readable description of the failure. */
+  message: string;
+}
+
+/** Vercel domain payload for a {@link ToolApprovalResponse}. */
+export interface VercelToolApprovalResponsePayload {
+  /** The tool call this approval responds to. */
+  toolCallId: string;
+  /** Whether the user approved the tool execution. */
+  approved: boolean;
+  /** Optional human-readable reason (typically used on denial). */
+  reason?: string;
+}
+
+// ---------------------------------------------------------------------------
+// Unions
+// ---------------------------------------------------------------------------
+
+/**
+ * The Vercel codec's `TInput` — every record-shape a client publishes on
+ * the `ai-input` wire. Composed from the SDK's well-known input shapes,
+ * with the tool variants parameterized by the Vercel domain payloads above.
+ */
+export type VercelInput =
+  | UserMessage<AI.UIMessage>
+  | Regenerate
+  | ToolResult<VercelToolResultPayload>
+  | ToolResultError<VercelToolResultErrorPayload>
+  | ToolApprovalResponse<VercelToolApprovalResponsePayload>;
+
+/**
+ * The Vercel codec's `TOutput` — every record-shape the agent publishes
+ * on the `ai-output` wire. The Vercel codec passes the AI SDK's
+ * `UIMessageChunk` through unchanged.
+ */
+export type VercelOutput = AI.UIMessageChunk;
+
+// ---------------------------------------------------------------------------
+// Projection re-export
+// ---------------------------------------------------------------------------

package/src/vercel/codec/fields.ts

@@ -0,0 +1,58 @@
+/**
+ * Shared Vercel codec header-field bindings.
+ *
+ * Each field binds a codec header key to its value type once (see
+ * {@link HeaderField}); the output/input descriptors and escape hatches all
+ * read and write through these bindings, so a header key cannot drift between
+ * the encode and decode side. Domain field names live in the Vercel layer, not
+ * core, per the header-discipline rule.
+ */
+
+import type * as AI from 'ai';
+
+import { boolField, enumField, type HeaderField, jsonField, strField } from '../../core/codec/index.js';
+
+/** Stream / message id (text & reasoning streams). */
+export const fId = strField('id');
+/**
+ * Provider metadata envelope, typed to the AI SDK shape. Annotated explicitly:
+ * the inferred type resolves to the AI SDK's internal `SharedV3ProviderMetadata`
+ * alias, which isn't portably nameable across the package boundary.
+ */
+export const fMeta: HeaderField<AI.ProviderMetadata | undefined, 'providerMetadata'> = jsonField<
+  AI.ProviderMetadata,
+  'providerMetadata'
+>('providerMetadata');
+/** Tool call id — defaulted to total: an absent header reads as `''`. */
+export const fToolCallId = strField('toolCallId', '');
+/** Tool name — defaulted to total. */
+export const fToolName = strField('toolName', '');
+/** Whether the tool is a dynamic tool. */
+export const fDynamic = boolField('dynamic');
+/** Optional human-readable title. */
+export const fTitle = strField('title');
+/** Whether the provider executed the tool. */
+export const fProviderExecuted = boolField('providerExecuted');
+/** Media type for file / source-document parts — defaulted to total. */
+export const fMediaType = strField('mediaType', '');
+/** Source id for source-url / source-document parts — defaulted to total. */
+export const fSourceId = strField('sourceId', '');
+
+// --- input-side bindings (shared by the input descriptors' encode/decode) ---
+
+/** Domain message id (`message.id`) stamped on every user-message part — distinct from the wire codec-message-id transport header. */
+export const fMessageId = strField('messageId');
+/** Whether the user approved a tool execution — defaulted to total so an absent header reads `false`. */
+export const fApproved = boolField('approved', false);
+/** Optional human-readable reason on a tool-approval response. */
+export const fReason = strField('reason');
+
+/**
+ * Validated finish reason. Mirrors the AI SDK's `FinishReason` literals and
+ * falls back to `'stop'` for an absent or unrecognized value.
+ */
+export const fFinishReason = enumField(
+  'finishReason',
+  ['stop', 'length', 'content-filter', 'tool-calls', 'error', 'other'] as const,
+  'stop',
+);

package/src/vercel/codec/fold-content.ts

@@ -0,0 +1,54 @@
+/**
+ * File and source content-part folds: file / source-url / source-document.
+ * These are independent attachments — each appends a part, never dedups.
+ */
+
+import type * as AI from 'ai';
+
+import { stripUndefined } from '../../utils.js';
+import { ensureMessage, type VercelProjection } from './reducer-state.js';
+
+/**
+ * Fold a file or source content chunk into the projection.
+ * @param state - Projection to fold into.
+ * @param chunk - The file, source-url, or source-document chunk.
+ * @param messageId - The target codec-message-id.
+ * @returns The same projection reference.
+ */
+export const foldContentPart = (
+  state: VercelProjection,
+  chunk: Extract<AI.UIMessageChunk, { type: 'file' | 'source-url' | 'source-document' }>,
+  messageId: string,
+): VercelProjection => {
+  const message = ensureMessage(state, messageId);
+
+  switch (chunk.type) {
+    case 'file': {
+      message.parts.push({ type: 'file', mediaType: chunk.mediaType, url: chunk.url });
+      return state;
+    }
+    case 'source-url': {
+      message.parts.push(
+        stripUndefined({
+          type: 'source-url' as const,
+          sourceId: chunk.sourceId,
+          url: chunk.url,
+          title: chunk.title,
+        }),
+      );
+      return state;
+    }
+    case 'source-document': {
+      message.parts.push(
+        stripUndefined({
+          type: 'source-document' as const,
+          sourceId: chunk.sourceId,
+          mediaType: chunk.mediaType,
+          title: chunk.title,
+          filename: chunk.filename,
+        }),
+      );
+      return state;
+    }
+  }
+};

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

@@ -0,0 +1,46 @@
+/**
+ * data-* part folds. Transient data parts are dropped; persistent ones are
+ * appended, or replaced in place when a matching `id` is already present.
+ */
+
+import type * as AI from 'ai';
+
+import { stripUndefined } from '../../utils.js';
+import { ensureMessage, type VercelProjection } from './reducer-state.js';
+
+/**
+ * Fold a `data-*` chunk into the projection.
+ * @param state - Projection to fold into.
+ * @param chunk - The data-* chunk.
+ * @param messageId - The target codec-message-id.
+ * @returns The same projection reference.
+ */
+export const foldDataPart = (
+  state: VercelProjection,
+  chunk: Extract<AI.UIMessageChunk, { type: `data-${string}` }>,
+  messageId: string,
+): VercelProjection => {
+  if (chunk.transient) return state;
+
+  const message = ensureMessage(state, messageId);
+
+  // CAST: chunk.type is `data-${string}` which satisfies DataUIPart, but
+  // TypeScript cannot verify the template literal matches a specific
+  // UIMessagePart variant at the type level.
+  const dataPart = stripUndefined({
+    type: chunk.type,
+    id: chunk.id,
+    data: chunk.data,
+  }) as AI.UIMessage['parts'][number];
+
+  if (chunk.id !== undefined) {
+    const idx = message.parts.findIndex((p) => p.type === chunk.type && 'id' in p && p.id === chunk.id);
+    if (idx !== -1) {
+      message.parts[idx] = dataPart;
+      return state;
+    }
+  }
+
+  message.parts.push(dataPart);
+  return state;
+};

package/src/vercel/codec/fold-input.ts

@@ -0,0 +1,255 @@
+/**
+ * Client-published input folds and the pending-resolution buffering.
+ *
+ * Tool resolutions (`ToolResult`, `ToolResultError`, `ToolApprovalResponse`)
+ * carry a `codecMessageId` targeting the assistant they amend. When that
+ * assistant (or its tool part) has not yet arrived, the resolution is buffered
+ * in `pendingToolResolutions` and {@link retryPendingResolutions} re-evaluates
+ * it after every subsequent fold.
+ */
+
+import type * as AI from 'ai';
+
+import type { ReducerMeta, ToolApprovalResponse, ToolResult, ToolResultError } from '../../core/codec/index.js';
+import type {
+  VercelToolApprovalResponsePayload,
+  VercelToolResultErrorPayload,
+  VercelToolResultPayload,
+} from './events.js';
+import {
+  ensureTrackers,
+  getToolPart,
+  type OwnerLookup,
+  type PendingToolResolution,
+  type VercelProjection,
+} from './reducer-state.js';
+import { toolBase, transitionToolPart } from './tool-transitions.js';
+
+/**
+ * Fold a user message into the projection, correlating on the wire
+ * codec-message-id (the caller's `message.id` is preserved verbatim). A
+ * multi-part user message fans out into one wire event per part, all sharing
+ * the codec-message-id — folding appends the incoming parts to the existing
+ * entry, reassembling the message part by part. The transport delivers each
+ * wire exactly once (its per-message version high-water-mark drops replays),
+ * so the merge sees every part once and stays consistent.
+ *
+ * Optimistic (serial-less) seeds need no special handling here: the transport
+ * refolds the node from its log when the echo's serial arrives, rebuilding the
+ * projection from a fresh `init` so the seed never coexists with its echo.
+ * @param state - Projection to fold into.
+ * @param message - The user message (or one decoded part of it) to add or merge.
+ * @param meta - Transport-derived metadata carrying the codec-message-id.
+ * @returns The same projection reference.
+ */
+export const foldUserMessage = (
+  state: VercelProjection,
+  message: AI.UIMessage,
+  meta: ReducerMeta,
+): VercelProjection => {
+  // Correlate the projection entry on the wire codec-message-id; the
+  // caller-supplied `message.id` is preserved verbatim and surfaced to the
+  // application unchanged. Without a codec-message-id the message has no
+  // identity to key on, so it is appended as a fresh entry.
+  const codecMessageId = meta.messageId;
+  if (codecMessageId === undefined) {
+    state.messages.push({ codecMessageId: message.id, message });
+    return state;
+  }
+  const existing = state.messages.find((e) => e.codecMessageId === codecMessageId);
+  if (existing === undefined) {
+    state.messages.push({ codecMessageId, message });
+  } else {
+    // Merge by codec-message-id: keep the existing envelope (id and role are
+    // stamped identically on every part of one message) and append the
+    // incoming parts in fold order — wire serials preserve publish order.
+    existing.message.parts.push(...message.parts);
+  }
+  return state;
+};
+
+/**
+ * Fold a client-published `ToolResult`. The input carries
+ * `codecMessageId` pointing at the assistant whose `dynamic-tool` part
+ * holds the matching `toolCallId`. If the assistant and its matching
+ * `dynamic-tool` part are both present, fold directly; otherwise pend
+ * until that tool part arrives.
+ * @param state - Projection to fold into.
+ * @param event - The tool-result input (codecMessageId + domain payload).
+ * @returns The same projection reference.
+ */
+export const foldClientToolResult = (
+  state: VercelProjection,
+  event: ToolResult<VercelToolResultPayload>,
+): VercelProjection => {
+  const { toolCallId, output } = event.payload;
+  return resolveOrPend(state, event.codecMessageId, toolCallId, { kind: 'tool-result', output });
+};
+
+/**
+ * Fold a client-published `ToolResultError`. Mirrors
+ * {@link foldClientToolResult} but with the error transition.
+ * @param state - Projection to fold into.
+ * @param event - The tool-result-error input (codecMessageId + domain payload).
+ * @returns The same projection reference.
+ */
+export const foldClientToolResultError = (
+  state: VercelProjection,
+  event: ToolResultError<VercelToolResultErrorPayload>,
+): VercelProjection => {
+  const { toolCallId, message } = event.payload;
+  return resolveOrPend(state, event.codecMessageId, toolCallId, { kind: 'tool-result-error', message });
+};
+
+/**
+ * Fold a client-published `ToolApprovalResponse`. The input carries
+ * `codecMessageId` pointing at the assistant whose `dynamic-tool` part
+ * holds the matching `toolCallId`. Approval → `approval-responded`;
+ * denial → `output-denied` via {@link transitionToolPart}.
+ * @param state - Projection to fold into.
+ * @param event - The approval-response input.
+ * @returns The same projection reference.
+ */
+export const foldToolApprovalResponse = (
+  state: VercelProjection,
+  event: ToolApprovalResponse<VercelToolApprovalResponsePayload>,
+): VercelProjection => {
+  const { toolCallId, approved, reason } = event.payload;
+  return resolveOrPend(state, event.codecMessageId, toolCallId, {
+    kind: 'tool-approval-response',
+    approved,
+    ...(reason === undefined ? {} : { reason }),
+  });
+};
+
+/**
+ * Apply a resolution when its tool part is present, otherwise buffer it in
+ * `pendingToolResolutions` for {@link retryPendingResolutions}.
+ * @param state - Projection to fold into.
+ * @param codecMessageId - The assistant the resolution targets.
+ * @param toolCallId - The tool call being resolved.
+ * @param resolution - The resolution variant to apply or buffer.
+ * @returns The same projection reference.
+ */
+const resolveOrPend = (
+  state: VercelProjection,
+  codecMessageId: string,
+  toolCallId: string,
+  resolution: PendingToolResolution['resolution'],
+): VercelProjection => {
+  const owner = findOwner(state, codecMessageId, toolCallId);
+  if (owner) {
+    applyResolution(owner, toolCallId, resolution);
+  } else {
+    state.pendingToolResolutions.push({ targetCodecMessageId: codecMessageId, toolCallId, resolution });
+  }
+  return state;
+};
+
+/**
+ * Apply one tool resolution onto its located `dynamic-tool` part, replacing
+ * the part with the transitioned shape — the single application point shared
+ * by the direct folds and {@link retryPendingResolutions}.
+ * @param owner - The located owner (message + tracker + part).
+ * @param toolCallId - The tool call being resolved.
+ * @param resolution - The resolution variant to apply.
+ */
+const applyResolution = (
+  owner: OwnerLookup,
+  toolCallId: string,
+  resolution: PendingToolResolution['resolution'],
+): void => {
+  switch (resolution.kind) {
+    case 'tool-result': {
+      owner.message.parts[owner.tracker.partIndex] = transitionToolPart(owner.part, {
+        type: 'tool-output-available',
+        toolCallId,
+        output: resolution.output,
+      });
+      break;
+    }
+    case 'tool-result-error': {
+      owner.message.parts[owner.tracker.partIndex] = transitionToolPart(owner.part, {
+        type: 'tool-output-error',
+        toolCallId,
+        errorText: resolution.message,
+      });
+      break;
+    }
+    case 'tool-approval-response': {
+      owner.message.parts[owner.tracker.partIndex] = approvalTransition(
+        owner.part,
+        resolution.approved,
+        resolution.reason,
+      );
+      break;
+    }
+  }
+};
+
+/**
+ * Re-attempt every pending tool resolution against the current projection.
+ * Successfully promoted entries are removed from the pending list. Cheap:
+ * bounded by the number of pending entries.
+ * @param state - Projection to walk and mutate.
+ */
+export const retryPendingResolutions = (state: VercelProjection): void => {
+  const next: VercelProjection['pendingToolResolutions'] = [];
+  for (const pending of state.pendingToolResolutions) {
+    const owner = findOwner(state, pending.targetCodecMessageId, pending.toolCallId);
+    if (!owner) {
+      next.push(pending);
+      continue;
+    }
+    applyResolution(owner, pending.toolCallId, pending.resolution);
+  }
+  state.pendingToolResolutions = next;
+};
+
+const findOwner = (state: VercelProjection, codecMessageId: string, toolCallId: string): OwnerLookup | undefined => {
+  const entry = state.messages.find((e) => e.codecMessageId === codecMessageId);
+  if (!entry) return undefined;
+  const trackers = ensureTrackers(state, codecMessageId);
+  const found = getToolPart(entry.message, trackers, toolCallId);
+  if (!found) return undefined;
+  return { message: entry.message, tracker: found.tracker, part: found.part };
+};
+
+/**
+ * Build the next `dynamic-tool` part shape for an approval response.
+ *
+ * For `approved=true`, transition to `approval-responded` so the AI SDK's
+ * multi-step loop will auto-run the tool on the next step.
+ * `transitionToolPart` has no shape for this transition, so we synthesize
+ * the part directly.
+ *
+ * For `approved=false`, delegate to `transitionToolPart` with a synthetic
+ * `tool-output-denied` chunk so denial mirrors the chunk-driven path.
+ * @param part - The existing `dynamic-tool` part being transitioned.
+ * @param approved - Whether the user approved the tool execution.
+ * @param reason - Optional human-readable reason.
+ * @returns The replacement `dynamic-tool` part.
+ */
+const approvalTransition = (
+  part: AI.DynamicToolUIPart,
+  approved: boolean,
+  reason: string | undefined,
+): AI.DynamicToolUIPart => {
+  if (approved) {
+    return {
+      ...toolBase(part),
+      state: 'approval-responded',
+      input: 'input' in part ? part.input : undefined,
+      approval: {
+        id: 'approval' in part && part.approval ? part.approval.id : '',
+        approved: true,
+        ...(reason === undefined ? {} : { reason }),
+      },
+    };
+  }
+  return transitionToolPart(part, {
+    type: 'tool-output-denied',
+    toolCallId: part.toolCallId,
+    ...(reason === undefined ? {} : { reason }),
+  });
+};

package/src/vercel/codec/fold-lifecycle.ts

@@ -0,0 +1,85 @@
+/**
+ * Lifecycle chunk folds: start, start-step, finish-step, finish, abort,
+ * error, message-metadata.
+ */
+
+import type * as AI from 'ai';
+
+import { ensureMessage, type VercelProjection } from './reducer-state.js';
+
+/**
+ * Set a message's metadata from a chunk when both the message exists and the
+ * chunk carries metadata. Shared by the `finish` and `message-metadata` cases,
+ * which apply it identically. The `start` case is not routed through here — it
+ * creates the message via `ensureMessage` first.
+ * @param state - Projection holding the message.
+ * @param messageId - The target codec-message-id.
+ * @param metadata - The chunk's `messageMetadata`, or undefined to leave it unchanged.
+ */
+const applyMessageMetadata = (state: VercelProjection, messageId: string, metadata: AI.UIMessage['metadata']): void => {
+  if (metadata === undefined) return;
+  const message = state.messages.find((e) => e.codecMessageId === messageId)?.message;
+  if (message) message.metadata = metadata;
+};
+
+/**
+ * Fold a message-lifecycle chunk into the projection.
+ * @param state - Projection to fold into.
+ * @param chunk - The lifecycle chunk.
+ * @param messageId - The target codec-message-id.
+ * @returns The same projection reference.
+ */
+export const foldLifecycle = (
+  state: VercelProjection,
+  chunk: Extract<
+    AI.UIMessageChunk,
+    { type: 'start' | 'start-step' | 'finish-step' | 'finish' | 'abort' | 'error' | 'message-metadata' }
+  >,
+  messageId: string,
+): VercelProjection => {
+  switch (chunk.type) {
+    case 'start': {
+      // The projection entry is keyed on the wire codec-message-id
+      // (`messageId`); every subsequent chunk for this message correlates on
+      // that, independent of `message.id`. So we faithfully reproduce the
+      // stream's own `messageId` on the reconstructed `UIMessage.id` (the
+      // value surfaced to the application) without risk of orphaning later
+      // chunks. When the stream omits it, the codec-message-id seeded by
+      // `ensureMessage` stands as the fallback id.
+      const message = ensureMessage(state, messageId);
+      if (chunk.messageId !== undefined) message.id = chunk.messageId;
+      if (chunk.messageMetadata !== undefined) message.metadata = chunk.messageMetadata;
+      return state;
+    }
+    case 'start-step': {
+      const message = ensureMessage(state, messageId);
+      message.parts.push({ type: 'step-start' });
+      return state;
+    }
+    case 'finish-step': {
+      // Reset text/reasoning stream trackers so a follow-up step can start
+      // new parts with potentially-reused stream ids.
+      const trackers = state.trackers.get(messageId);
+      if (trackers) {
+        trackers.text.clear();
+        trackers.reasoning.clear();
+      }
+      return state;
+    }
+    case 'finish': {
+      applyMessageMetadata(state, messageId, chunk.messageMetadata);
+      // Tracker state retained — late events still resolvable; cleanup happens at Run end.
+      return state;
+    }
+    case 'abort':
+    case 'error': {
+      // No state mutation — run termination is observed via the wire run-end
+      // event, not the projection.
+      return state;
+    }
+    case 'message-metadata': {
+      applyMessageMetadata(state, messageId, chunk.messageMetadata);
+      return state;
+    }
+  }
+};

package/src/vercel/codec/fold-text.ts

@@ -0,0 +1,55 @@
+/**
+ * Text and reasoning streaming folds: the {start, delta, end} lifecycle for
+ * both `text-*` and `reasoning-*` chunks, which share the same shape.
+ */
+
+import type * as AI from 'ai';
+
+import { ensureMessage, ensureTrackers, type VercelProjection } from './reducer-state.js';
+
+/**
+ * Fold a text or reasoning streaming chunk into the projection.
+ * @param state - Projection to fold into.
+ * @param chunk - The text/reasoning start, delta, or end chunk.
+ * @param messageId - The target codec-message-id.
+ * @returns The same projection reference.
+ */
+export const foldTextOrReasoning = (
+  state: VercelProjection,
+  chunk: Extract<
+    AI.UIMessageChunk,
+    { type: 'text-start' | 'text-delta' | 'text-end' | 'reasoning-start' | 'reasoning-delta' | 'reasoning-end' }
+  >,
+  messageId: string,
+): VercelProjection => {
+  const message = ensureMessage(state, messageId);
+  const trackers = ensureTrackers(state, messageId);
+
+  const isText = chunk.type.startsWith('text-');
+  const partType = isText ? 'text' : 'reasoning';
+  const activeMap = isText ? trackers.text : trackers.reasoning;
+
+  switch (chunk.type) {
+    case 'text-start':
+    case 'reasoning-start': {
+      activeMap.set(chunk.id, message.parts.length);
+      message.parts.push({ type: partType, text: '' });
+      return state;
+    }
+    case 'text-delta':
+    case 'reasoning-delta': {
+      const idx = activeMap.get(chunk.id);
+      if (idx === undefined) return state;
+      const part = message.parts[idx];
+      if (part?.type === partType) {
+        part.text += chunk.delta;
+      }
+      return state;
+    }
+    case 'text-end':
+    case 'reasoning-end': {
+      activeMap.delete(chunk.id);
+      return state;
+    }
+  }
+};