@ably/ai-transport 0.2.0 → 0.3.0

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;
+    }
+  }
+};

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

@@ -0,0 +1,86 @@
+/**
+ * Tool-input streaming folds: tool-input-start / -delta / -available / -error.
+ * Tool deltas arrive as raw JSON fragments accumulated in the tracker's
+ * `inputText` buffer and parsed on each delta.
+ */
+
+import type * as AI from 'ai';
+
+import { parseJson } from '../../utils.js';
+import { ensureMessage, ensureTrackers, getToolPart, type VercelProjection } from './reducer-state.js';
+import { toolBase } from './tool-transitions.js';
+
+/**
+ * Fold a tool-input streaming chunk into the projection.
+ * @param state - Projection to fold into.
+ * @param chunk - The tool-input start, delta, available, or error chunk.
+ * @param messageId - The target codec-message-id.
+ * @returns The same projection reference.
+ */
+export const foldToolInput = (
+  state: VercelProjection,
+  chunk: Extract<
+    AI.UIMessageChunk,
+    { type: 'tool-input-start' | 'tool-input-delta' | 'tool-input-available' | 'tool-input-error' }
+  >,
+  messageId: string,
+): VercelProjection => {
+  const message = ensureMessage(state, messageId);
+  const trackers = ensureTrackers(state, messageId);
+
+  switch (chunk.type) {
+    case 'tool-input-start': {
+      const partIndex = message.parts.length;
+      message.parts.push({ ...toolBase(chunk), state: 'input-streaming', input: undefined });
+      trackers.tools.set(chunk.toolCallId, { partIndex, inputText: '' });
+      return state;
+    }
+    case 'tool-input-delta': {
+      const tracker = trackers.tools.get(chunk.toolCallId);
+      if (!tracker) return state;
+      tracker.inputText += chunk.inputTextDelta;
+
+      const parsedInput = parseJson(tracker.inputText);
+
+      const found = getToolPart(message, trackers, chunk.toolCallId);
+      if (!found) return state;
+      message.parts[found.tracker.partIndex] = {
+        ...toolBase(found.part),
+        state: 'input-streaming',
+        input: parsedInput,
+      };
+      return state;
+    }
+    case 'tool-input-available': {
+      const found = getToolPart(message, trackers, chunk.toolCallId);
+      if (!found) return state;
+      message.parts[found.tracker.partIndex] = {
+        ...toolBase(found.part),
+        state: 'input-available',
+        input: chunk.input,
+      };
+      return state;
+    }
+    case 'tool-input-error': {
+      const found = getToolPart(message, trackers, chunk.toolCallId);
+      if (found) {
+        message.parts[found.tracker.partIndex] = {
+          ...toolBase(found.part),
+          state: 'output-error',
+          input: chunk.input,
+          errorText: chunk.errorText,
+        };
+      } else {
+        const partIndex = message.parts.length;
+        message.parts.push({
+          ...toolBase(chunk),
+          state: 'output-error',
+          input: chunk.input,
+          errorText: chunk.errorText,
+        });
+        trackers.tools.set(chunk.toolCallId, { partIndex, inputText: '' });
+      }
+      return state;
+    }
+  }
+};

package/src/vercel/codec/fold-tool-output.ts

@@ -0,0 +1,79 @@
+/**
+ * Agent-published tool-output transitions: tool-output-available /
+ * tool-output-error / tool-output-denied / tool-approval-request.
+ */
+
+import type * as AI from 'ai';
+
+import {
+  ensureMessage,
+  ensureTrackers,
+  getToolPart,
+  type OwnerLookup,
+  type VercelProjection,
+} from './reducer-state.js';
+import { transitionToolPart } from './tool-transitions.js';
+
+/**
+ * Locate the `dynamic-tool` part for a `toolCallId` anywhere in the projection.
+ * Agent-emitted second-pass tool outputs (after an approved tool runs) are
+ * stamped with a fresh codec-message-id that differs from the assistant holding
+ * the tool call, so they can't be found via `meta.messageId` — they fold onto
+ * whichever message holds the matching tool call (created in the first pass or
+ * by an approval response).
+ * @param state - Projection to scan.
+ * @param toolCallId - The tool call to locate.
+ * @returns The owning message, tracker, and part, or `undefined` if absent.
+ */
+const findToolPartOwner = (state: VercelProjection, toolCallId: string): OwnerLookup | undefined => {
+  for (const entry of state.messages) {
+    const trackers = state.trackers.get(entry.codecMessageId);
+    if (!trackers) continue;
+    const found = getToolPart(entry.message, trackers, toolCallId);
+    if (found) return { message: entry.message, tracker: found.tracker, part: found.part };
+  }
+  return undefined;
+};
+
+/**
+ * Fold an agent-published tool-output chunk into the projection.
+ * @param state - Projection to fold into.
+ * @param chunk - The tool-output-available/-error/-denied or tool-approval-request chunk.
+ * @param messageId - The target codec-message-id (used for the approval-request / denied paths).
+ * @returns The same projection reference.
+ */
+export const foldToolOutput = (
+  state: VercelProjection,
+  chunk: Extract<
+    AI.UIMessageChunk,
+    { type: 'tool-output-available' | 'tool-output-error' | 'tool-output-denied' | 'tool-approval-request' }
+  >,
+  messageId: string,
+): VercelProjection => {
+  // `tool-output-available` / `tool-output-error` after an approved tool runs
+  // are emitted by streamText's continuation pass under a fresh
+  // codec-message-id that differs from the assistant holding the tool call.
+  // Resolve the owning part by toolCallId across the whole projection so the
+  // output folds onto the original message. Deliberately do NOT materialise
+  // `messageId` first — that would leave a phantom empty message behind the
+  // fresh id. Drop on miss: a tool output with no matching tool call has no
+  // anchor to attach to.
+  if (chunk.type === 'tool-output-available' || chunk.type === 'tool-output-error') {
+    const owner = findToolPartOwner(state, chunk.toolCallId);
+    if (!owner) return state;
+    owner.message.parts[owner.tracker.partIndex] = transitionToolPart(owner.part, chunk);
+    return state;
+  }
+
+  // `tool-approval-request` (first pass) creates the part on the run's own
+  // message; `tool-output-denied` transitions that same part. Both key on the
+  // stamped messageId.
+  const message = ensureMessage(state, messageId);
+  const trackers = ensureTrackers(state, messageId);
+
+  const found = getToolPart(message, trackers, chunk.toolCallId);
+  if (!found) return state;
+
+  message.parts[found.tracker.partIndex] = transitionToolPart(found.part, chunk);
+  return state;
+};

package/src/vercel/codec/index.ts

@@ -1,10 +1,12 @@
 /**
- * Vercel AI SDK codec — implements
- * `Codec<VercelInput, VercelOutput, VercelProjection, UIMessage>`.
+ * Vercel AI SDK codec — `UIMessageCodec`.
  *
- * The codec is the reducer (extends `Reducer<VercelInput | VercelOutput,
- * VercelProjection>`) plus encoder/decoder factories and `getMessages`
- * for Tree population.
+ * Assembled by `defineCodec` from the codec's parts: the reducer
+ * (`init`/`fold`/`getMessages`), the declarative output and input descriptor
+ * tables (`outputs` / `inputs`, each a builder function `defineCodec` injects
+ * the direction-scoped builder into), and the decode lifecycle policy.
+ * `defineCodec` builds the generic encoder/decoder and merges the well-known
+ * input factories internally.
  *
  * ```ts
  * import { UIMessageCodec } from '@ably/ai-transport/vercel';
@@ -15,68 +17,26 @@
  * ```
  */
 
-import type * as AI from 'ai';
-
-import type { Codec } from '../../core/codec/types.js';
-import { createDecoder } from './decoder.js';
-import { createEncoder } from './encoder.js';
-import type {
-  VercelInput,
-  VercelOutput,
-  VercelToolApprovalResponsePayload,
-  VercelToolResultErrorPayload,
-  VercelToolResultPayload,
-} from './events.js';
-import { fold, getMessages, init, type VercelProjection } from './reducer.js';
+import { defineCodec } from '../../core/codec/index.js';
+import { createVercelDecodeLifecycle } from './decode-lifecycle.js';
+import type { VercelInput, VercelOutput } from './events.js';
+import { inputs } from './inputs.js';
+import { outputs } from './outputs.js';
+import { fold, getMessages, init } from './reducer.js';
 
 /**
  * Vercel AI SDK codec implementing
- * `Codec<VercelInput, VercelOutput, VercelProjection, UIMessage>`.
- *
- * Folds `VercelInput`s and `VercelOutput`s into a `VercelProjection`
- * carrying `UIMessage[]`. Encoder and decoder factories handle the wire
- * mapping for both directions.
+ * `Codec<VercelInput, VercelOutput, VercelProjection, UIMessage>`. `VercelProjection`
+ * and `UIMessage` are inferred from the reducer.
  */
-const uiMessageCodecImpl = {
-  // Internal field - picked up by registerAgent via AdapterTagHolder cast. Spec: AIT-CT1a3, AIT-ST1a3.
-  adapterTag: 'vercel-ai-sdk-ui-message' as const,
-  init,
-  fold,
-  createEncoder,
-  createDecoder,
-  getMessages,
-  createUserMessage: (message: AI.UIMessage): VercelInput => ({ kind: 'user-message', message }),
-  createRegenerate: (target: string, parent: string): VercelInput => ({
-    kind: 'regenerate',
-    target,
-    parent,
-  }),
-  createToolResult: (codecMessageId: string, payload: VercelToolResultPayload): VercelInput => ({
-    kind: 'tool-result',
-    codecMessageId,
-    payload,
-  }),
-  createToolResultError: (codecMessageId: string, payload: VercelToolResultErrorPayload): VercelInput => ({
-    kind: 'tool-result-error',
-    codecMessageId,
-    payload,
-  }),
-  createToolApprovalResponse: (codecMessageId: string, payload: VercelToolApprovalResponsePayload): VercelInput => ({
-    kind: 'tool-approval-response',
-    codecMessageId,
-    payload,
-  }),
-};
-
-// Validate Codec conformance via `satisfies` on the variable (no excess-property
-// check, so the internal `adapterTag` is permitted) while keeping the concrete
-// type so the codec-specific factories (createToolResult, etc.) stay callable.
-export const UIMessageCodec = uiMessageCodecImpl satisfies Codec<
-  VercelInput,
-  VercelOutput,
-  VercelProjection,
-  AI.UIMessage
->;
+export const UIMessageCodec = defineCodec<VercelInput, VercelOutput>()({
+  // Spec: AIT-CT1a3, AIT-ST1a3 — registers this codec as an Ably agent.
+  adapterTag: 'vercel-ai-sdk-ui-message',
+  reducer: { init, fold, getMessages },
+  output: outputs,
+  input: inputs,
+  decodeLifecycle: createVercelDecodeLifecycle,
+});
 
 export type { VercelInput, VercelOutput } from './events.js';
 export { type VercelProjection } from './reducer.js';

package/src/vercel/codec/inputs.ts

@@ -0,0 +1,116 @@
+/**
+ * Vercel input (`ai-input`) descriptors — the single source of truth for the
+ * `VercelInput` wire mapping, the user-message fan-out included.
+ *
+ * `defineCodec` injects the direction-scoped `{ event, batch }` builder; the
+ * generic input drivers consume the returned array. The tool inputs are single
+ * `event`s lensed onto their nested `payload`; `regenerate` is a wire-only
+ * signal; the multi-part user message is a `batch` that fans each
+ * `UIMessage` part out into one wire event (reassembled by the reducer).
+ *
+ * Author-facing acceptance gate: the injected `event`/`batch` builders narrow
+ * each member, so every `data` / `fields` / `parts` / `assemble` callback is
+ * fully typed. The file's single `as` cast is the wire trust boundary on the
+ * inbound role header (see `assemble`).
+ */
+
+import type * as AI from 'ai';
+
+import { HEADER_ROLE } from '../../constants.js';
+import type { InputBuilder, InputDescriptor } from '../../core/codec/index.js';
+import type { VercelInput } from './events.js';
+import { fApproved, fId, fMediaType, fMessageId, fReason, fToolCallId } from './fields.js';
+import { asString, isClientToolResultErrorWireData, isToolOutputAvailableWireData } from './wire-data.js';
+
+/** Fallback for a message with no encodable parts (see the `user-message` batch). */
+const EMPTY_MESSAGE_PARTS: AI.UIMessage['parts'] = [{ type: 'text', text: '' }];
+
+/**
+ * Part types the `user-message` batch's `parts` sub-table can encode — must
+ * stay in step with that table. Parts outside this set (e.g. `step-start`,
+ * tool parts) have no wire mapping; `explode` filters them so the batch always
+ * yields at least one encodable part and the message round-trips.
+ * @param part - The UIMessage part to test.
+ * @returns Whether the part has a wire mapping in the batch's part table.
+ */
+const isEncodablePart = (part: AI.UIMessage['parts'][number]): boolean =>
+  part.type === 'text' || part.type === 'file' || part.type.startsWith('data-');
+
+/**
+ * The Vercel codec's `ai-input` descriptors, built from the injected
+ * direction-scoped builder.
+ * @param builder - The `{ event, batch }` builder curried on `VercelInput`.
+ * @param builder.event - Define a single-event input (payload-nested, or `wireOnly`).
+ * @param builder.batch - Define a multi-part (batch) input that fans out into one wire event per part.
+ * @returns The input descriptor table the generic input drivers consume.
+ */
+export const inputs = ({ event, batch }: InputBuilder<VercelInput>): readonly InputDescriptor<VercelInput>[] => [
+  // --- tool inputs: nested payload, codec-message-id-addressed ----------------
+
+  event('tool-result', {
+    fields: [fToolCallId],
+    data: {
+      encode: (p) => ({ output: p.output }),
+      // Malformed wire data decodes to undefined, which the rebuild seam strips
+      // — the folded payload then has no `output` key (reads as undefined).
+      decode: (d) => ({ output: isToolOutputAvailableWireData(d) ? d.output : undefined }),
+    },
+  }),
+  event('tool-result-error', {
+    fields: [fToolCallId],
+    data: {
+      encode: (p) => ({ message: p.message }),
+      decode: (d) => ({ message: isClientToolResultErrorWireData(d) ? (d.message ?? '') : '' }),
+    },
+  }),
+  event('tool-approval-response', { fields: [fToolCallId, fApproved, fReason] }),
+
+  // --- wire-only signal -------------------------------------------------------
+
+  // `regenerate` carries no domain payload; `parent` / `target` ride the
+  // transport headers built by the client-session and read by the agent's
+  // input-event lookup, so it stamps only the `kind` header and decodes to [].
+  event('regenerate', { wireOnly: true }),
+
+  // --- multi-part client message ----------------------------------------------
+
+  // The user message fans out into one wire event per part, all sharing the
+  // `user-message` kind and codec-message-id, each carrying its `partType`. The
+  // message id (a codec header) and role (a transport header) are per-message,
+  // stamped on every part so the decode side can rebuild the envelope from any
+  // one; the reducer merges parts sharing a codec-message-id.
+  batch('user-message', {
+    // A message with no encodable parts (empty, or only unmapped types like
+    // step-start) still publishes one empty text part, so the codec-message-id
+    // and role survive and it round-trips to a one-part message. The driver's
+    // bare-headers fallback cannot round-trip (it carries no partType), so the
+    // ≥1-encodable-part guarantee lives here.
+    explode: (input) => {
+      const encodable = input.message.parts.filter((part) => isEncodablePart(part));
+      return encodable.length > 0 ? encodable : EMPTY_MESSAGE_PARTS;
+    },
+    partTypeOf: (part) => part.type,
+    parts: (p) => [
+      p('text', { data: { encode: (x) => x.text, decode: (d) => ({ text: asString(d) }) } }),
+      p('file', {
+        fields: [fMediaType],
+        data: { encode: (x) => x.url, decode: (d) => ({ url: asString(d) }) },
+      }),
+      p('data-*', {
+        fields: [fId],
+        data: { encode: (x) => x.data, decode: (d) => ({ data: d }) },
+      }),
+    ],
+    messageHeaders: (input) => {
+      const codecHeaders: Record<string, string> = {};
+      fMessageId.write(codecHeaders, input.message.id);
+      return { codecHeaders, transportHeaders: { [HEADER_ROLE]: input.message.role } };
+    },
+    assemble: (part, { codecHeaders, transportHeaders }) => {
+      // CAST: HEADER_ROLE is wire data; the role string is trusted as a UIMessage role.
+      const role = (transportHeaders[HEADER_ROLE] ?? 'user') as AI.UIMessage['role'];
+      const id = fMessageId.read(codecHeaders) ?? '';
+      return { message: { id, role, parts: [part] } };
+    },
+  }),
+];