@ably/ai-transport 0.2.0 → 0.3.0

package/src/vercel/codec/reducer.ts

@@ -11,11 +11,12 @@
  * with no instance state. Mutation in place is allowed — the projection
  * is single-owner.
  *
- * Idempotency is **per conflict key**, not stream-wide: when two events
- * compete for the same logical state (e.g. two `tool-output-available`
- * for the same `toolCallId`), the higher-serial one wins and the other
- * is dropped. Unrelated events arrive freely in any order. See
- * `_conflictKeyOf` for the per-variant key derivation.
+ * 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
@@ -23,129 +24,31 @@
  * `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.
- */
-
-import type * as AI from 'ai';
-
-import type {
-  CodecMessage,
-  ReducerMeta,
-  ToolApprovalResponse,
-  ToolResult,
-  ToolResultError,
-} from '../../core/codec/types.js';
-import { stripUndefined } from '../../utils.js';
-import type {
-  VercelInput,
-  VercelOutput,
-  VercelToolApprovalResponsePayload,
-  VercelToolResultErrorPayload,
-  VercelToolResultPayload,
-} from './events.js';
-import { toolBase, transitionToolPart } from './tool-transitions.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.
- */
-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. */
-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).
+ * 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.
  */
-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-conflict-key high-water-marks. Maps a codec-derived conflict key
-   * (see `_conflictKeyOf`) to the highest `meta.serial` already folded for
-   * that key. Events whose serial is `<=` the stored value are dropped as
-   * duplicates of an already-incorporated operation. Events that have no
-   * conflict key (additive content, lifecycle markers) are folded
-   * unconditionally.
-   */
-  conflictSerials: Map<string, string>;
-  /** 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.
- */
-interface PendingToolResolution {
-  /** The codec-message-id of the assistant the resolution targets. */
-  targetCodecMessageId: string;
-  /** Tool call this resolution targets. */
-  toolCallId: string;
-  /** Serial of the wire message — used by the conflict-key check on promotion. */
-  serial: 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 };
-}
+import type * as AI from 'ai';
 
-// ---------------------------------------------------------------------------
-// init
-// ---------------------------------------------------------------------------
-
-/**
- * Build an empty initial projection.
- * @returns A fresh VercelProjection with no messages and no tracker state.
- */
-export const init = (): VercelProjection => ({
-  messages: [],
-  conflictSerials: new Map(),
-  trackers: new Map(),
-  pendingToolResolutions: [],
-});
+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
@@ -155,12 +58,11 @@ export const init = (): VercelProjection => ({
  * Fold one input or output event into the projection. Mutates and returns
  * `state`.
  *
- * Idempotency is per conflict key (see `_conflictKeyOf`): if the event has
- * a conflict key and the projection has already folded an event for that
- * key at a higher-or-equal serial, this call is a no-op. Events without a
- * conflict key (additive content, lifecycle markers) are folded
- * unconditionally. Orphan events (e.g. tool-output for an unknown
- * toolCallId) are dropped silently inside the per-variant fold helpers.
+ * 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).
@@ -168,24 +70,14 @@ export const init = (): VercelProjection => ({
  */
 export const fold = (
   state: VercelProjection,
-  event: VercelInput | VercelOutput,
+  event: CodecEvent<VercelInput, VercelOutput>,
   meta: ReducerMeta,
 ): VercelProjection => {
-  if (meta.serial) {
-    const key = _conflictKeyOf(event, meta);
-    if (key !== undefined) {
-      const seen = state.conflictSerials.get(key);
-      if (seen !== undefined && meta.serial <= seen) {
-        return state;
-      }
-      state.conflictSerials.set(key, meta.serial);
-    }
-  }
-
-  if (_isInput(event)) {
-    switch (event.kind) {
+  if (event.direction === 'input') {
+    const input = event.event;
+    switch (input.kind) {
       case 'user-message': {
-        _foldUserMessage(state, event.message, meta);
+        foldUserMessage(state, input.message, meta);
         break;
       }
       case 'regenerate': {
@@ -195,377 +87,37 @@ export const fold = (
         break;
       }
       case 'tool-result': {
-        _foldClientToolResult(state, event, meta);
+        foldClientToolResult(state, input);
         break;
       }
       case 'tool-result-error': {
-        _foldClientToolResultError(state, event, meta);
+        foldClientToolResultError(state, input);
         break;
       }
       case 'tool-approval-response': {
-        _foldToolApprovalResponse(state, event, meta);
+        foldToolApprovalResponse(state, input);
         break;
       }
     }
   } else {
-    _foldChunk(state, event, meta);
+    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);
+    retryPendingResolutions(state);
   }
 
   return state;
 };
 
-/**
- * Narrow the union to TInput vs TOutput by the discriminator field name.
- * VercelInput variants carry `kind`; VercelOutput variants carry `type`.
- * @param event - The event to narrow.
- * @returns True when the event is a VercelInput, false for VercelOutput.
- */
-const _isInput = (event: VercelInput | VercelOutput): event is VercelInput => 'kind' in event;
-
-// ---------------------------------------------------------------------------
-// Conflict-key derivation
-// ---------------------------------------------------------------------------
-
-/**
- * Derive a per-event conflict key, or `undefined` if the event doesn't
- * compete with any other event for shared state. Used by `fold` to scope
- * the high-water-mark check to genuine conflicts (e.g. two
- * `tool-output-available` for the same `toolCallId`) rather than to every
- * event in the stream.
- * @param event - The event being folded.
- * @param meta - Transport-derived metadata (used for events keyed by codec-message-id).
- * @returns The conflict key, or `undefined` if the event is additive / independent.
- */
-const _conflictKeyOf = (event: VercelInput | VercelOutput, meta: ReducerMeta): string | undefined => {
-  if (_isInput(event)) {
-    switch (event.kind) {
-      case 'user-message': {
-        // Dedup re-publishes of the same user message by its wire
-        // codec-message-id, never by the domain `message.id`. Without a
-        // codec-message-id there is nothing to correlate on, so the fold
-        // is left unconditional.
-        return meta.messageId === undefined ? undefined : `user-msg:${meta.messageId}`;
-      }
-      case 'tool-approval-response': {
-        return `tool-approval:${event.payload.toolCallId}`;
-      }
-      // Client tool results compete for the same final state of the tool
-      // call (against agent-side `tool-output-available`/`tool-output-error`
-      // chunks and against `tool-output-denied`/`tool-approval-request`).
-      // Highest serial wins. Shares the `tool-output:` namespace with the
-      // agent-side chunks below.
-      case 'tool-result':
-      case 'tool-result-error': {
-        return `tool-output:${event.payload.toolCallId}`;
-      }
-      case 'regenerate': {
-        return undefined;
-      }
-    }
-  }
-
-  switch (event.type) {
-    // Tool-input state machine, keyed by toolCallId.
-    case 'tool-input-start':
-    case 'tool-input-available':
-    case 'tool-input-error': {
-      return `${event.type}:${event.toolCallId}`;
-    }
-
-    // All "tool-output-ish" output variants compete for the same final
-    // state of the tool call. Shares the `tool-output:` namespace with
-    // the client-published input variants above.
-    case 'tool-output-available':
-    case 'tool-output-error':
-    case 'tool-output-denied':
-    case 'tool-approval-request': {
-      return `tool-output:${event.toolCallId}`;
-    }
-
-    // Per-stream start/end markers: duplicates would create phantom parts
-    // or wipe accumulated text. Keyed by (codec-message-id, stream-id).
-    case 'text-start':
-    case 'text-end':
-    case 'reasoning-start':
-    case 'reasoning-end': {
-      return `${event.type}:${meta.messageId ?? ''}:${event.id}`;
-    }
-
-    // Message-level markers, keyed by codec-message-id.
-    case 'finish':
-    case 'message-metadata': {
-      return `${event.type}:${meta.messageId ?? ''}`;
-    }
-
-    // Purely additive or independent — never dedup:
-    //   text-delta / reasoning-delta / tool-input-delta (additive content)
-    //   start / start-step / finish-step / abort / error (lifecycle)
-    //   file / source-url / source-document (independent attachments)
-    //   data-* (opaque to the reducer)
-    default: {
-      return undefined;
-    }
-  }
-};
-
 // ---------------------------------------------------------------------------
-// Input folds
+// UIMessageChunk dispatch
 // ---------------------------------------------------------------------------
 
-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 existingIdx = state.messages.findIndex((e) => e.codecMessageId === codecMessageId);
-  if (existingIdx === -1) {
-    state.messages.push({ codecMessageId, message });
-  } else {
-    state.messages[existingIdx] = { codecMessageId, message };
-  }
-  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).
- * @param meta - Transport-derived metadata.
- * @returns The same projection reference.
- */
-const _foldClientToolResult = (
-  state: VercelProjection,
-  event: ToolResult<VercelToolResultPayload>,
-  meta: ReducerMeta,
-): VercelProjection => {
-  const { toolCallId, output } = event.payload;
-  const owner = _findOwner(state, event.codecMessageId, toolCallId);
-  if (owner) {
-    owner.message.parts[owner.tracker.partIndex] = transitionToolPart(owner.part, {
-      type: 'tool-output-available',
-      toolCallId,
-      output,
-    });
-    return state;
-  }
-
-  state.pendingToolResolutions.push({
-    targetCodecMessageId: event.codecMessageId,
-    toolCallId,
-    serial: meta.serial,
-    resolution: { kind: 'tool-result', output },
-  });
-  return state;
-};
-
-/**
- * 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).
- * @param meta - Transport-derived metadata.
- * @returns The same projection reference.
- */
-const _foldClientToolResultError = (
-  state: VercelProjection,
-  event: ToolResultError<VercelToolResultErrorPayload>,
-  meta: ReducerMeta,
-): VercelProjection => {
-  const { toolCallId, message } = event.payload;
-  const owner = _findOwner(state, event.codecMessageId, toolCallId);
-  if (owner) {
-    owner.message.parts[owner.tracker.partIndex] = transitionToolPart(owner.part, {
-      type: 'tool-output-error',
-      toolCallId,
-      errorText: message,
-    });
-    return state;
-  }
-
-  state.pendingToolResolutions.push({
-    targetCodecMessageId: event.codecMessageId,
-    toolCallId,
-    serial: meta.serial,
-    resolution: { kind: 'tool-result-error', message },
-  });
-  return state;
-};
-
-/**
- * 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.
- * @param meta - Transport-derived metadata.
- * @returns The same projection reference.
- */
-const _foldToolApprovalResponse = (
-  state: VercelProjection,
-  event: ToolApprovalResponse<VercelToolApprovalResponsePayload>,
-  meta: ReducerMeta,
-): VercelProjection => {
-  const { toolCallId, approved, reason } = event.payload;
-  const owner = _findOwner(state, event.codecMessageId, toolCallId);
-  if (owner) {
-    owner.message.parts[owner.tracker.partIndex] = _approvalTransition(owner.part, approved, reason);
-    return state;
-  }
-
-  state.pendingToolResolutions.push({
-    targetCodecMessageId: event.codecMessageId,
-    toolCallId,
-    serial: meta.serial,
-    resolution: {
-      kind: 'tool-approval-response',
-      approved,
-      ...(reason === undefined ? {} : { reason }),
-    },
-  });
-  return state;
-};
-
-interface OwnerLookup {
-  message: AI.UIMessage;
-  tracker: ToolPartTracker;
-  part: AI.DynamicToolUIPart;
-}
-
-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 };
-};
-
-/**
- * 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;
-};
-
-/**
- * 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 }),
-  });
-};
-
-/**
- * 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.
- */
-const _retryPendingResolutions = (state: VercelProjection): void => {
-  const next: PendingToolResolution[] = [];
-  for (const pending of state.pendingToolResolutions) {
-    const owner = _findOwner(state, pending.targetCodecMessageId, pending.toolCallId);
-    if (!owner) {
-      next.push(pending);
-      continue;
-    }
-    switch (pending.resolution.kind) {
-      case 'tool-result': {
-        owner.message.parts[owner.tracker.partIndex] = transitionToolPart(owner.part, {
-          type: 'tool-output-available',
-          toolCallId: pending.toolCallId,
-          output: pending.resolution.output,
-        });
-        break;
-      }
-      case 'tool-result-error': {
-        owner.message.parts[owner.tracker.partIndex] = transitionToolPart(owner.part, {
-          type: 'tool-output-error',
-          toolCallId: pending.toolCallId,
-          errorText: pending.resolution.message,
-        });
-        break;
-      }
-      case 'tool-approval-response': {
-        owner.message.parts[owner.tracker.partIndex] = _approvalTransition(
-          owner.part,
-          pending.resolution.approved,
-          pending.resolution.reason,
-        );
-        break;
-      }
-    }
-  }
-  state.pendingToolResolutions = next;
-};
-
-// ---------------------------------------------------------------------------
-// UIMessageChunk fold
-// ---------------------------------------------------------------------------
-
-const _foldChunk = (state: VercelProjection, chunk: VercelOutput, meta: ReducerMeta): VercelProjection => {
+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.
@@ -580,7 +132,7 @@ const _foldChunk = (state: VercelProjection, chunk: VercelOutput, meta: ReducerM
     case 'abort':
     case 'error':
     case 'message-metadata': {
-      return _foldLifecycle(state, chunk, messageId);
+      return foldLifecycle(state, chunk, messageId);
     }
 
     case 'text-start':
@@ -589,378 +141,38 @@ const _foldChunk = (state: VercelProjection, chunk: VercelOutput, meta: ReducerM
     case 'reasoning-start':
     case 'reasoning-delta':
     case 'reasoning-end': {
-      return _foldTextOrReasoning(state, chunk, messageId);
+      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);
+      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);
+      return foldToolOutput(state, chunk, messageId);
     }
 
     case 'file':
     case 'source-url':
     case 'source-document': {
-      return _foldContentPart(state, chunk, messageId);
+      return foldContentPart(state, chunk, messageId);
     }
 
     default: {
       if (chunk.type.startsWith('data-')) {
-        return _foldDataPart(state, chunk, messageId);
+        return foldDataPart(state, chunk, messageId);
       }
       return state;
     }
   }
 };
 
-// ---------------------------------------------------------------------------
-// Message + tracker helpers
-// ---------------------------------------------------------------------------
-
-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;
-};
-
-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;
-};
-
-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 };
-};
-
-// ---------------------------------------------------------------------------
-// Lifecycle events
-// ---------------------------------------------------------------------------
-
-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': {
-      const message = state.messages.find((e) => e.codecMessageId === messageId)?.message;
-      if (message && chunk.messageMetadata !== undefined) {
-        message.metadata = 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': {
-      const message = state.messages.find((e) => e.codecMessageId === messageId)?.message;
-      if (message && chunk.messageMetadata !== undefined) {
-        message.metadata = chunk.messageMetadata;
-      }
-      return state;
-    }
-  }
-};
-
-// ---------------------------------------------------------------------------
-// Text and reasoning streaming
-// ---------------------------------------------------------------------------
-
-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;
-    }
-  }
-};
-
-// ---------------------------------------------------------------------------
-// Tool input streaming
-// ---------------------------------------------------------------------------
-
-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;
-
-      let parsedInput: unknown;
-      try {
-        // CAST: JSON.parse returns any; unknown is the safe trust-boundary type.
-        parsedInput = JSON.parse(tracker.inputText) as unknown;
-      } catch {
-        parsedInput = undefined;
-      }
-
-      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;
-    }
-  }
-};
-
-// ---------------------------------------------------------------------------
-// Tool output transitions (agent-published chunks)
-// ---------------------------------------------------------------------------
-
-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;
-};
-
-// ---------------------------------------------------------------------------
-// File / source content parts
-// ---------------------------------------------------------------------------
-
-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;
-    }
-  }
-};
-
-// ---------------------------------------------------------------------------
-// data-* parts
-// ---------------------------------------------------------------------------
-
-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;
-};
-
 // ---------------------------------------------------------------------------
 // getMessages
 // ---------------------------------------------------------------------------
@@ -975,3 +187,5 @@ const _foldDataPart = (
  * @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';