@ably/ai-transport 0.1.0 → 0.3.0

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,35 +1,42 @@
 /**
- * Vercel AI SDK codec — maps UIMessageChunk events and UIMessage objects
- * to/from native Ably message primitives (publish, append, update, delete).
+ * Vercel AI SDK codec — `UIMessageCodec`.
+ *
+ * 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';
  *
  * const encoder = UIMessageCodec.createEncoder(writer, options);
  * const decoder = UIMessageCodec.createDecoder();
- * const accumulator = UIMessageCodec.createAccumulator();
+ * const projection = UIMessageCodec.init();
  * ```
  */
 
-import type * as AI from 'ai';
-
-import type { Codec } from '../../core/codec/types.js';
-import { createAccumulator } from './accumulator.js';
-import { createDecoder } from './decoder.js';
-import { createEncoder } from './encoder.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<UIMessageChunk, UIMessage>`.
- *
- * Provides factory methods for creating encoders, decoders, and accumulators
- * that map between Vercel's UIMessageChunk/UIMessage types and Ably's native
- * message primitives.
+ * Vercel AI SDK codec implementing
+ * `Codec<VercelInput, VercelOutput, VercelProjection, UIMessage>`. `VercelProjection`
+ * and `UIMessage` are inferred from the reducer.
  */
-export const UIMessageCodec: Codec<AI.UIMessageChunk, AI.UIMessage> = {
-  createEncoder,
-  createDecoder,
-  createAccumulator,
+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,
+});
 
-  isTerminal: (event: AI.UIMessageChunk): boolean =>
-    event.type === 'finish' || event.type === 'error' || event.type === 'abort',
-};
+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] } };
+    },
+  }),
+];

package/src/vercel/codec/outputs.ts

@@ -0,0 +1,207 @@
+/**
+ * Vercel output (`ai-output`) event descriptors — the single source of truth for
+ * encoding/decoding `UIMessageChunk` outputs. `defineCodec` injects the
+ * direction-scoped `{ event, stream }` builder; the generic drivers consume the
+ * returned array. Adding an ordinary output event is one entry here.
+ *
+ * Author-facing acceptance gate: this file contains **zero `as` casts**. The
+ * injected `event`/`stream` builders narrow each chunk member, so the
+ * `data`/`encode`/`decode` callbacks are fully typed.
+ */
+
+import * as Ably from 'ably';
+
+import type { OutputBuilder, OutputDescriptor } from '../../core/codec/index.js';
+import { boolField, jsonField, strField } from '../../core/codec/index.js';
+import { ErrorCode, errorInfoIs } from '../../errors.js';
+import { parseJsonOrString, stripUndefined } from '../../utils.js';
+import type { VercelOutput } from './events.js';
+import {
+  fDynamic,
+  fFinishReason,
+  fId,
+  fMediaType,
+  fMessageId,
+  fMeta,
+  fProviderExecuted,
+  fSourceId,
+  fTitle,
+  fToolCallId,
+  fToolName,
+} from './fields.js';
+import {
+  asString,
+  isAgentToolOutputErrorWireData,
+  isToolInputErrorWireData,
+  isToolOutputAvailableWireData,
+} from './wire-data.js';
+
+/**
+ * The Vercel codec's `ai-output` descriptors, built from the injected
+ * direction-scoped builder.
+ * @param builder - The `{ event, stream }` builder curried on `VercelOutput`.
+ * @param builder.event - Define a single discrete output event.
+ * @param builder.stream - Define a streamed output family (start / delta / end).
+ * @returns The output descriptor table the generic output drivers consume.
+ */
+export const outputs = ({ event, stream }: OutputBuilder<VercelOutput>): readonly OutputDescriptor<VercelOutput>[] => [
+  // --- streamed families -----------------------------------------------------
+
+  stream('text', {
+    start: 'text-start',
+    delta: 'text-delta',
+    end: 'text-end',
+    idField: 'id',
+    deltaField: 'delta',
+    fields: [fId, fMeta],
+  }),
+
+  stream('reasoning', {
+    start: 'reasoning-start',
+    delta: 'reasoning-delta',
+    end: 'reasoning-end',
+    idField: 'id',
+    deltaField: 'delta',
+    fields: [fId, fMeta],
+  }),
+
+  // tool-input streams; the close step is a close-or-discrete fallback, the end
+  // chunk reconstructs `input` from the accumulated text, and the family also
+  // decodes from a non-streamed discrete publish.
+  stream('tool-input', {
+    start: 'tool-input-start',
+    delta: 'tool-input-delta',
+    end: 'tool-input-available',
+    idField: 'toolCallId',
+    deltaField: 'inputTextDelta',
+    fields: [fToolCallId, fToolName, fDynamic, fTitle, fProviderExecuted, fMeta],
+    onEnd: async (c, core, { h, name }) => {
+      try {
+        await core.closeStream(c.toolCallId, {
+          name,
+          data: '',
+          codecHeaders: h(c, ['toolCallId', 'toolName', 'providerMetadata']),
+        });
+      } catch (error: unknown) {
+        // closeStream raises InvalidArgument when there is no active stream for
+        // this id; fall through to a discrete publish, rethrow anything else.
+        if (!(error instanceof Ably.ErrorInfo && errorInfoIs(error, ErrorCode.InvalidArgument))) {
+          throw error;
+        }
+        await core.publishDiscrete({ name, data: c.input, codecHeaders: h(c) });
+      }
+    },
+    decodeEnd: ({ streamId, accumulated, codecHeaders, closingCodecHeaders }) => [
+      stripUndefined({
+        type: 'tool-input-available' as const,
+        toolCallId: streamId,
+        toolName: fToolName.read(closingCodecHeaders) || fToolName.read(codecHeaders),
+        input: parseJsonOrString(accumulated),
+        providerMetadata: fMeta.read(closingCodecHeaders),
+      }),
+    ],
+    decodeDiscrete: ({ codecHeaders, data }) => [
+      stripUndefined({
+        type: 'tool-input-start' as const,
+        toolCallId: fToolCallId.read(codecHeaders),
+        toolName: fToolName.read(codecHeaders),
+        dynamic: fDynamic.read(codecHeaders),
+        title: fTitle.read(codecHeaders),
+        providerExecuted: fProviderExecuted.read(codecHeaders),
+        providerMetadata: fMeta.read(codecHeaders),
+      }),
+      stripUndefined({
+        type: 'tool-input-available' as const,
+        toolCallId: fToolCallId.read(codecHeaders),
+        toolName: fToolName.read(codecHeaders),
+        input: data,
+        providerMetadata: fMeta.read(codecHeaders),
+      }),
+    ],
+  }),
+
+  // --- discrete lifecycle events ---------------------------------------------
+
+  // `start` injects the encoder's configured messageId as a fallback, so it
+  // builds its headers through a hatch rather than a pure descriptor.
+  event('start', {
+    fields: [fMessageId, jsonField('messageMetadata')],
+    encode: async (c, core, { h, name, messageId, opts }) => {
+      await core.publishDiscrete(
+        { name, data: '', codecHeaders: h({ ...c, messageId: c.messageId ?? messageId }) },
+        opts,
+      );
+    },
+  }),
+  event('start-step'),
+  event('finish-step'),
+  event('finish', {
+    fields: [fFinishReason, jsonField('messageMetadata')],
+  }),
+  event('message-metadata', { fields: [jsonField('messageMetadata')] }),
+  event('error', {
+    data: { encode: (c) => c.errorText, decode: (data) => ({ errorText: asString(data) }) },
+  }),
+
+  // abort: an ordinary discrete output carrying its reason as wire data. The
+  // agent's own stream emits it on abort; run cancellation closes in-flight
+  // streams via the encoder's cancelStreams() and terminates via the transport
+  // ai-run-end event — this chunk is content, not the run terminator.
+  event('abort', {
+    data: {
+      encode: (c) => c.reason ?? '',
+      decode: (data) => (typeof data === 'string' && data ? { reason: data } : {}),
+    },
+  }),
+
+  // --- content parts ---------------------------------------------------------
+
+  event('file', {
+    fields: [fMediaType, fMeta],
+    data: { encode: (c) => c.url, decode: (data) => ({ url: asString(data) }) },
+  }),
+  event('source-url', {
+    fields: [fSourceId, fTitle, fMeta],
+    data: { encode: (c) => c.url, decode: (data) => ({ url: asString(data) }) },
+  }),
+  event('source-document', {
+    fields: [fSourceId, fMediaType, strField('title', ''), strField('filename'), fMeta],
+  }),
+
+  // --- tool lifecycle (discrete) ---------------------------------------------
+
+  event('tool-input-error', {
+    fields: [fToolCallId, fToolName, fDynamic, fTitle, fProviderExecuted, fMeta],
+    data: {
+      encode: (c) => ({ errorText: c.errorText, input: c.input }),
+      decode: (data) =>
+        isToolInputErrorWireData(data) ? { errorText: data.errorText ?? '', input: data.input } : { errorText: '' },
+    },
+  }),
+  event('tool-output-available', {
+    fields: [fToolCallId, fDynamic, fProviderExecuted, boolField('preliminary')],
+    data: {
+      encode: (c) => ({ output: c.output }),
+      decode: (data) => (isToolOutputAvailableWireData(data) ? { output: data.output } : {}),
+    },
+  }),
+  event('tool-output-error', {
+    fields: [fToolCallId, fDynamic, fProviderExecuted],
+    data: {
+      encode: (c) => ({ errorText: c.errorText }),
+      decode: (data) => ({ errorText: isAgentToolOutputErrorWireData(data) ? (data.errorText ?? '') : '' }),
+    },
+  }),
+  event('tool-approval-request', {
+    fields: [fToolCallId, strField('approvalId', '')],
+  }),
+  event('tool-output-denied', { fields: [fToolCallId] }),
+
+  // --- data-* wildcard -------------------------------------------------------
+
+  event('data-*', {
+    fields: [fId, boolField('transient')],
+    ephemeral: (c) => c.transient === true,
+    data: { encode: (c) => c.data, decode: (data) => ({ data }) },
+  }),
+];