@ably/ai-transport 0.2.0 → 0.3.0

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 }) },
+  }),
+];

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