@ably/ai-transport 0.1.0 → 0.2.0

package/src/vercel/codec/decoder.ts

@@ -1,26 +1,51 @@
 /**
- * Vercel AI SDK Decoder
+ * Vercel AI SDK Decoder.
  *
- * Maps Ably inbound messages to DecoderOutput<UIMessageChunk, UIMessage>[].
+ * Maps Ably inbound messages to {@link DecodedMessage} — a `{ inputs,
+ * outputs }` tagged result. The decoder routes by the wire `name`
+ * (`ai-input` vs `ai-output`) so the SDK never has to inspect direction:
+ * input-side messages produce `VercelInput` variants; output-side
+ * messages produce `VercelOutput` (`UIMessageChunk`) variants.
  *
- * Delegates action dispatch and serial tracking to the decoder core.
- * This file contains only the Vercel-specific event building, discrete
- * event decoding, and synthetic event emission.
+ * The `LifecycleTracker` is an internal helper used to pre-roll missing
+ * `start` / `start-step` chunks on mid-stream join (history compaction,
+ * rewind miss, partial page) so the reducer always sees a clean event
+ * sequence for streamed output.
  *
- * Domain-specific headers use the `x-domain-` prefix. Transport-level
- * headers use the `x-ably-` prefix.
+ * Receive-side dispatch reads the wire `name` first and then routes by
+ * the codec `type` header carrying the codec event type. Codec headers live
+ * under `extras.ai.codec` and transport headers under `extras.ai.transport`;
+ * both are read unprefixed from their respective tier.
  */
 
 import type * as Ably from 'ably';
 import type * as AI from 'ai';
 
-import { HEADER_DISCRETE, HEADER_ROLE, HEADER_TURN_ID } from '../../constants.js';
+import {
+  EVENT_AI_INPUT,
+  EVENT_AI_OUTPUT,
+  HEADER_CODEC_MESSAGE_ID,
+  HEADER_DISCRETE,
+  HEADER_ROLE,
+  HEADER_RUN_ID,
+} from '../../constants.js';
 import type { DecoderCore, DecoderCoreHooks, DecoderCoreOptions } from '../../core/codec/decoder.js';
-import { createDecoderCore, eventOutput } from '../../core/codec/decoder.js';
+import { createDecoderCore } from '../../core/codec/decoder.js';
 import type { LifecycleTracker } from '../../core/codec/lifecycle-tracker.js';
 import { createLifecycleTracker } from '../../core/codec/lifecycle-tracker.js';
-import type { DecoderOutput, MessagePayload, StreamDecoder, StreamTrackerState } from '../../core/codec/types.js';
+import type {
+  DecodedMessage,
+  Decoder,
+  MessagePayload,
+  StreamTrackerState,
+  UserMessage,
+} from '../../core/codec/types.js';
 import { type DomainHeaderReader, headerReader as rawHeaderReader, stripUndefined } from '../../utils.js';
+import type { VercelInput, VercelOutput } from './events.js';
+
+// Decoder-internal union — the codec emits inputs and outputs through the
+// same flat list from the underlying core and partitions on the way out.
+type AnyEvent = VercelInput | VercelOutput;
 
 // ---------------------------------------------------------------------------
 // Vercel-specific header reader (casts providerMetadata to AI.ProviderMetadata)
@@ -49,45 +74,31 @@ const headerReader = (headers: Record<string, string>): VercelHeaderReader => {
 // Wire format types (trust boundaries for JSON-parsed data)
 // ---------------------------------------------------------------------------
 
-/** Wire format for tool-input-error data payload. */
+/** Wire format for the agent-side `tool-input-error` chunk data payload. */
 interface ToolInputErrorWireData {
   errorText?: string;
   input?: unknown;
 }
 
-/** Wire format for tool-output-available data payload. */
+/** Wire format for the `tool-output-available` (agent) / `tool-result` (client) data payload. */
 interface ToolOutputAvailableWireData {
   output?: unknown;
 }
 
-/** Wire format for tool-output-error data payload. */
-interface ToolOutputErrorWireData {
+/** Wire format for the agent-side `tool-output-error` chunk data payload. */
+interface AgentToolOutputErrorWireData {
   errorText?: string;
 }
 
-// ---------------------------------------------------------------------------
-// Shared output type alias
-// ---------------------------------------------------------------------------
-
-type Out = DecoderOutput<AI.UIMessageChunk, AI.UIMessage>;
-
-/**
- * Bind eventOutput to the Vercel domain types.
- * @param chunk - The UIMessageChunk to wrap.
- * @returns A single-element decoder output array.
- */
-const event = (chunk: AI.UIMessageChunk): Out[] => eventOutput<AI.UIMessageChunk, AI.UIMessage>(chunk);
+/** Wire format for the client-side `tool-result-error` input data payload. */
+interface ClientToolResultErrorWireData {
+  message?: string;
+}
 
 // ---------------------------------------------------------------------------
 // JSON boundary helpers
 // ---------------------------------------------------------------------------
 
-/**
- * Validate a finish reason string against the FinishReason union.
- * @param value - The finish reason string from the wire, or undefined.
- * @param fallback - Default finish reason if the value is not recognized.
- * @returns The validated FinishReason.
- */
 const parseFinishReason = (value: string | undefined, fallback: AI.FinishReason): AI.FinishReason => {
   if (
     value === 'stop' ||
@@ -102,19 +113,8 @@ const parseFinishReason = (value: string | undefined, fallback: AI.FinishReason)
   return fallback;
 };
 
-/**
- * Type predicate for data-* message names.
- * @param name - The message name to check.
- * @returns True if the name starts with "data-".
- */
 const isDataEventName = (name: string): name is `data-${string}` => name.startsWith('data-');
 
-/**
- * Parse a string as JSON, returning the raw string if parsing fails
- * or undefined if empty.
- * @param value - The string to attempt JSON parsing on.
- * @returns The parsed value, the raw string on parse failure, or undefined if empty.
- */
 const parseJsonOrString = (value: string): unknown => {
   if (!value) return undefined;
   try {
@@ -126,12 +126,22 @@ const parseJsonOrString = (value: string): unknown => {
 };
 
 // ---------------------------------------------------------------------------
-// Streamed message event builders
+// Streamed message event builders (output-side)
 // ---------------------------------------------------------------------------
 
+/**
+ * Read the codec event type from a tracker's codec headers. The encoder
+ * stamps the codec `type` header on every `ai-output` publish; the value
+ * carries the AI-SDK chunk family (`text` / `reasoning` / `tool-input`)
+ * that the stream represents.
+ * @param tracker - The stream tracker carrying the persistent headers.
+ * @returns The codec event type, or the empty string when absent.
+ */
+const codecTypeOf = (tracker: StreamTrackerState): string => headerReader(tracker.codecHeaders).strOr('type', '');
+
 const buildStartChunk = (tracker: StreamTrackerState): AI.UIMessageChunk => {
-  const r = headerReader(tracker.headers);
-  switch (tracker.name) {
+  const r = headerReader(tracker.codecHeaders);
+  switch (codecTypeOf(tracker)) {
     case 'text': {
       return stripUndefined({
         type: 'text-start' as const,
@@ -164,7 +174,7 @@ const buildStartChunk = (tracker: StreamTrackerState): AI.UIMessageChunk => {
 };
 
 const buildDeltaChunk = (tracker: StreamTrackerState, delta: string): AI.UIMessageChunk => {
-  switch (tracker.name) {
+  switch (codecTypeOf(tracker)) {
     case 'text': {
       return { type: 'text-delta', id: tracker.streamId, delta };
     }
@@ -182,7 +192,7 @@ const buildDeltaChunk = (tracker: StreamTrackerState, delta: string): AI.UIMessa
 
 const buildEndChunk = (tracker: StreamTrackerState, closingHeaders: Record<string, string>): AI.UIMessageChunk => {
   const r = headerReader(closingHeaders);
-  switch (tracker.name) {
+  switch (codecTypeOf(tracker)) {
     case 'text': {
       return stripUndefined({
         type: 'text-end' as const,
@@ -201,7 +211,7 @@ const buildEndChunk = (tracker: StreamTrackerState, closingHeaders: Record<strin
       return stripUndefined({
         type: 'tool-input-available' as const,
         toolCallId: tracker.streamId,
-        toolName: r.strOr('toolName', headerReader(tracker.headers).strOr('toolName', '')),
+        toolName: r.strOr('toolName', headerReader(tracker.codecHeaders).strOr('toolName', '')),
         input: parseJsonOrString(tracker.accumulated),
         providerMetadata: r.providerMetadata(),
       });
@@ -213,7 +223,7 @@ const buildEndChunk = (tracker: StreamTrackerState, closingHeaders: Record<strin
 };
 
 // ---------------------------------------------------------------------------
-// Lifecycle tracker configuration (synthetic event phases)
+// Lifecycle tracker configuration (synthetic event phases on mid-stream join)
 // ---------------------------------------------------------------------------
 
 const createVercelLifecycleTracker = (): LifecycleTracker<AI.UIMessageChunk> =>
@@ -228,107 +238,108 @@ const createVercelLifecycleTracker = (): LifecycleTracker<AI.UIMessageChunk> =>
     },
   ]);
 
-/**
- * Run the lifecycle tracker and wrap results as DecoderOutput events.
- * @param lifecycle - The lifecycle tracker instance.
- * @param turnId - The turn scope ID.
- * @param context - Context passed through to phase build functions.
- * @returns Decoder outputs for any synthesized lifecycle events.
- */
-const ensurePhases = (
-  lifecycle: LifecycleTracker<AI.UIMessageChunk>,
-  turnId: string,
-  context: Record<string, string | undefined>,
-): Out[] => lifecycle.ensurePhases(turnId, context).map((e) => ({ kind: 'event', event: e }));
-
 // ---------------------------------------------------------------------------
-// Discrete event decoders (one function per event type)
+// Discrete output decoders (ai-output → UIMessageChunk)
 // ---------------------------------------------------------------------------
 
-const decodeStart = (r: VercelHeaderReader, turnId: string, lifecycle: LifecycleTracker<AI.UIMessageChunk>): Out[] => {
-  lifecycle.markEmitted(turnId, 'start');
-  return event(
+const decodeStart = (
+  r: VercelHeaderReader,
+  runId: string,
+  lifecycle: LifecycleTracker<AI.UIMessageChunk>,
+): AI.UIMessageChunk[] => {
+  lifecycle.markEmitted(runId, 'start');
+  return [
     stripUndefined({
       type: 'start' as const,
       messageId: r.str('messageId'),
       messageMetadata: r.json('messageMetadata'),
     }),
-  );
+  ];
 };
 
-const decodeStartStep = (turnId: string, lifecycle: LifecycleTracker<AI.UIMessageChunk>): Out[] => {
-  lifecycle.markEmitted(turnId, 'start-step');
-  return event({ type: 'start-step' });
+const decodeStartStep = (runId: string, lifecycle: LifecycleTracker<AI.UIMessageChunk>): AI.UIMessageChunk[] => {
+  lifecycle.markEmitted(runId, 'start-step');
+  return [{ type: 'start-step' }];
 };
 
-const decodeFinishStep = (turnId: string, lifecycle: LifecycleTracker<AI.UIMessageChunk>): Out[] => {
-  lifecycle.resetPhase(turnId, 'start-step');
-  return event({ type: 'finish-step' });
+const decodeFinishStep = (runId: string, lifecycle: LifecycleTracker<AI.UIMessageChunk>): AI.UIMessageChunk[] => {
+  lifecycle.resetPhase(runId, 'start-step');
+  return [{ type: 'finish-step' }];
 };
 
-const decodeFinish = (r: VercelHeaderReader, turnId: string, lifecycle: LifecycleTracker<AI.UIMessageChunk>): Out[] => {
-  lifecycle.clearScope(turnId);
-  return event(
+const decodeFinish = (
+  r: VercelHeaderReader,
+  runId: string,
+  lifecycle: LifecycleTracker<AI.UIMessageChunk>,
+): AI.UIMessageChunk[] => {
+  lifecycle.clearScope(runId);
+  return [
     stripUndefined({
       type: 'finish' as const,
       finishReason: parseFinishReason(r.str('finishReason'), 'stop'),
       messageMetadata: r.json('messageMetadata'),
     }),
-  );
+  ];
 };
 
-const decodeError = (data: unknown, turnId: string, lifecycle: LifecycleTracker<AI.UIMessageChunk>): Out[] => {
-  lifecycle.clearScope(turnId);
+const decodeError = (
+  data: unknown,
+  runId: string,
+  lifecycle: LifecycleTracker<AI.UIMessageChunk>,
+): AI.UIMessageChunk[] => {
+  lifecycle.clearScope(runId);
   const errorText = typeof data === 'string' ? data : '';
-  return event({ type: 'error', errorText });
+  return [{ type: 'error', errorText }];
 };
 
-const decodeAbort = (data: unknown, turnId: string, lifecycle: LifecycleTracker<AI.UIMessageChunk>): Out[] => {
-  lifecycle.clearScope(turnId);
+const decodeAbort = (
+  data: unknown,
+  runId: string,
+  lifecycle: LifecycleTracker<AI.UIMessageChunk>,
+): AI.UIMessageChunk[] => {
+  lifecycle.clearScope(runId);
   const reason = typeof data === 'string' && data ? data : undefined;
-  return event(stripUndefined({ type: 'abort' as const, reason }));
+  return [stripUndefined({ type: 'abort' as const, reason })];
 };
 
-const decodeMessageMetadata = (r: VercelHeaderReader): Out[] =>
-  event({ type: 'message-metadata', messageMetadata: r.json('messageMetadata') });
-
-const decodeFile = (r: VercelHeaderReader, data: unknown): Out[] =>
-  event(
-    stripUndefined({
-      type: 'file' as const,
-      url: typeof data === 'string' ? data : '',
-      mediaType: r.strOr('mediaType', ''),
-      providerMetadata: r.providerMetadata(),
-    }),
-  );
-
-const decodeSourceUrl = (r: VercelHeaderReader, data: unknown): Out[] =>
-  event(
-    stripUndefined({
-      type: 'source-url' as const,
-      sourceId: r.strOr('sourceId', ''),
-      url: typeof data === 'string' ? data : '',
-      title: r.str('title'),
-      providerMetadata: r.providerMetadata(),
-    }),
-  );
-
-const decodeSourceDocument = (r: VercelHeaderReader): Out[] =>
-  event(
-    stripUndefined({
-      type: 'source-document' as const,
-      sourceId: r.strOr('sourceId', ''),
-      mediaType: r.strOr('mediaType', ''),
-      title: r.strOr('title', ''),
-      filename: r.str('filename'),
-      providerMetadata: r.providerMetadata(),
-    }),
-  );
-
-const decodeToolInputError = (r: VercelHeaderReader, data: unknown): Out[] => {
+const decodeMessageMetadata = (r: VercelHeaderReader): AI.UIMessageChunk[] => [
+  { type: 'message-metadata', messageMetadata: r.json('messageMetadata') },
+];
+
+const decodeFile = (r: VercelHeaderReader, data: unknown): AI.UIMessageChunk[] => [
+  stripUndefined({
+    type: 'file' as const,
+    url: typeof data === 'string' ? data : '',
+    mediaType: r.strOr('mediaType', ''),
+    providerMetadata: r.providerMetadata(),
+  }),
+];
+
+const decodeSourceUrl = (r: VercelHeaderReader, data: unknown): AI.UIMessageChunk[] => [
+  stripUndefined({
+    type: 'source-url' as const,
+    sourceId: r.strOr('sourceId', ''),
+    url: typeof data === 'string' ? data : '',
+    title: r.str('title'),
+    providerMetadata: r.providerMetadata(),
+  }),
+];
+
+const decodeSourceDocument = (r: VercelHeaderReader): AI.UIMessageChunk[] => [
+  stripUndefined({
+    type: 'source-document' as const,
+    sourceId: r.strOr('sourceId', ''),
+    mediaType: r.strOr('mediaType', ''),
+    title: r.strOr('title', ''),
+    filename: r.str('filename'),
+    providerMetadata: r.providerMetadata(),
+  }),
+];
+
+const decodeToolInputError = (r: VercelHeaderReader, data: unknown): AI.UIMessageChunk[] => {
   // CAST: Trust boundary — encoder produced the expected object shape.
   const parsed = data as ToolInputErrorWireData | undefined;
-  return event(
+  return [
     stripUndefined({
       type: 'tool-input-error' as const,
       toolCallId: r.strOr('toolCallId', ''),
@@ -340,13 +351,13 @@ const decodeToolInputError = (r: VercelHeaderReader, data: unknown): Out[] => {
       providerExecuted: r.bool('providerExecuted'),
       providerMetadata: r.providerMetadata(),
     }),
-  );
+  ];
 };
 
-const decodeToolOutputAvailable = (r: VercelHeaderReader, data: unknown): Out[] => {
+const decodeAgentToolOutputAvailable = (r: VercelHeaderReader, data: unknown): AI.UIMessageChunk[] => {
   // CAST: Trust boundary — encoder produced the expected object shape.
   const parsed = data as ToolOutputAvailableWireData | undefined;
-  return event(
+  return [
     stripUndefined({
       type: 'tool-output-available' as const,
       toolCallId: r.strOr('toolCallId', ''),
@@ -355,13 +366,13 @@ const decodeToolOutputAvailable = (r: VercelHeaderReader, data: unknown): Out[]
       providerExecuted: r.bool('providerExecuted'),
       preliminary: r.bool('preliminary'),
     }),
-  );
+  ];
 };
 
-const decodeToolOutputError = (r: VercelHeaderReader, data: unknown): Out[] => {
+const decodeAgentToolOutputError = (r: VercelHeaderReader, data: unknown): AI.UIMessageChunk[] => {
   // CAST: Trust boundary — encoder produced the expected object shape.
-  const parsed = data as ToolOutputErrorWireData | undefined;
-  return event(
+  const parsed = data as AgentToolOutputErrorWireData | undefined;
+  return [
     stripUndefined({
       type: 'tool-output-error' as const,
       toolCallId: r.strOr('toolCallId', ''),
@@ -369,92 +380,80 @@ const decodeToolOutputError = (r: VercelHeaderReader, data: unknown): Out[] => {
       dynamic: r.bool('dynamic'),
       providerExecuted: r.bool('providerExecuted'),
     }),
-  );
+  ];
 };
 
-const decodeToolApprovalRequest = (r: VercelHeaderReader): Out[] =>
-  event({
+const decodeToolApprovalRequest = (r: VercelHeaderReader): AI.UIMessageChunk[] => [
+  {
     type: 'tool-approval-request',
     toolCallId: r.strOr('toolCallId', ''),
     approvalId: r.strOr('approvalId', ''),
-  });
+  },
+];
 
-const decodeToolOutputDenied = (r: VercelHeaderReader): Out[] =>
-  event({ type: 'tool-output-denied', toolCallId: r.strOr('toolCallId', '') });
+const decodeToolOutputDenied = (r: VercelHeaderReader): AI.UIMessageChunk[] => [
+  { type: 'tool-output-denied', toolCallId: r.strOr('toolCallId', '') },
+];
 
-const decodeDataEvent = (name: `data-${string}`, r: VercelHeaderReader, data: unknown): Out[] =>
-  event(
-    stripUndefined({
-      type: name,
-      data,
-      id: r.str('id'),
-      transient: r.bool('transient'),
-    }),
-  );
+const decodeDataEvent = (name: `data-${string}`, r: VercelHeaderReader, data: unknown): AI.UIMessageChunk[] => [
+  stripUndefined({
+    type: name,
+    data,
+    id: r.str('id'),
+    transient: r.bool('transient'),
+  }),
+];
 
 // ---------------------------------------------------------------------------
-// Non-streaming tool-input helper
+// Non-streaming tool-input helper (agent-side)
 // ---------------------------------------------------------------------------
 
 const decodeNonStreamingToolInput = (
   r: VercelHeaderReader,
   data: unknown,
-  turnId: string,
+  runId: string,
   lifecycle: LifecycleTracker<AI.UIMessageChunk>,
-): Out[] => {
-  const outputs = ensurePhases(lifecycle, turnId, { messageId: r.str('messageId') });
-
-  outputs.push(
-    {
-      kind: 'event',
-      event: stripUndefined({
-        type: 'tool-input-start' as const,
-        toolCallId: r.strOr('toolCallId', ''),
-        toolName: r.strOr('toolName', ''),
-        dynamic: r.bool('dynamic'),
-        title: r.str('title'),
-        providerExecuted: r.bool('providerExecuted'),
-        providerMetadata: r.providerMetadata(),
-      }),
-    },
-    {
-      kind: 'event',
-      event: stripUndefined({
-        type: 'tool-input-available' as const,
-        toolCallId: r.strOr('toolCallId', ''),
-        toolName: r.strOr('toolName', ''),
-        input: data,
-        providerMetadata: r.providerMetadata(),
-      }),
-    },
-  );
-
-  return outputs;
-};
+): AI.UIMessageChunk[] => [
+  ...lifecycle.ensurePhases(runId, { messageId: r.str('messageId') }),
+  stripUndefined({
+    type: 'tool-input-start' as const,
+    toolCallId: r.strOr('toolCallId', ''),
+    toolName: r.strOr('toolName', ''),
+    dynamic: r.bool('dynamic'),
+    title: r.str('title'),
+    providerExecuted: r.bool('providerExecuted'),
+    providerMetadata: r.providerMetadata(),
+  }),
+  stripUndefined({
+    type: 'tool-input-available' as const,
+    toolCallId: r.strOr('toolCallId', ''),
+    toolName: r.strOr('toolName', ''),
+    input: data,
+    providerMetadata: r.providerMetadata(),
+  }),
+];
 
 // ---------------------------------------------------------------------------
-// Discrete event dispatch
+// Input-side decoders (ai-input → VercelInput)
 // ---------------------------------------------------------------------------
 
 /**
- * Reconstruct a UIMessage from a discrete message part published by writeMessages.
- * The encoder splits each UIMessage into per-part Ably messages with a shared
- * x-domain-messageId. This function rebuilds a single-part UIMessage from one
- * such Ably message. The transport's tree upsert merges parts that share the
- * same x-ably-msg-id, so multi-part messages accumulate correctly over
- * successive decoder calls.
- * @param input - The discrete message payload to decode.
- * @returns A single-element array with the reconstructed UIMessage, or empty if unrecognized.
+ * Decode a single discrete message part (from the user-message multi-part
+ * wire format) into a {@link UserMessage} carrying a one-part
+ * UIMessage. The reducer's `_foldUserMessage` merges parts that share
+ * the same codec-message-id.
+ * @param input - The discrete message payload (name, data, headers).
+ * @returns A single `user-message` input, or an empty array when the part type is unrecognised.
  */
-const decodeDiscreteMessage = (input: MessagePayload): Out[] => {
-  const h = input.headers ?? {};
-  const r = headerReader(h);
-  const role = (h[HEADER_ROLE] ?? 'user') as AI.UIMessage['role'];
+const decodeDiscreteMessagePart = (input: MessagePayload): VercelInput[] => {
+  const r = headerReader(input.codecHeaders ?? {});
+  const role = (input.transportHeaders?.[HEADER_ROLE] ?? 'user') as AI.UIMessage['role'];
   const messageId = r.str('messageId') ?? '';
+  const codecType = r.strOr('type', '');
 
   let part: AI.UIMessage['parts'][number] | undefined;
 
-  switch (input.name) {
+  switch (codecType) {
     case 'text': {
       part = { type: 'text', text: typeof input.data === 'string' ? input.data : '' };
       break;
@@ -468,9 +467,8 @@ const decodeDiscreteMessage = (input: MessagePayload): Out[] => {
       break;
     }
     default: {
-      if (isDataEventName(input.name)) {
-        // CAST: data-* part type matches the DataUIPart shape.
-        part = stripUndefined({ type: input.name, id: r.str('id'), data: input.data }) as AI.UIMessage['parts'][number];
+      if (isDataEventName(codecType)) {
+        part = stripUndefined({ type: codecType, id: r.str('id'), data: input.data });
       }
       break;
     }
@@ -479,76 +477,102 @@ const decodeDiscreteMessage = (input: MessagePayload): Out[] => {
   if (!part) return [];
 
   const message: AI.UIMessage = { id: messageId, role, parts: [part] };
-  return [{ kind: 'message', message }];
+  const userMessage: UserMessage<AI.UIMessage> = { kind: 'user-message', message };
+  return [userMessage];
 };
 
-/**
- * Whether a message name represents a discrete message part (written by writeMessages)
- * rather than a streaming lifecycle event. Distinguished by the `x-ably-discrete` header
- * which {@link publishDiscreteBatch} sets on batch-published message payloads. Lifecycle
- * events published via {@link publishDiscrete} (including streaming `data-*` chunks)
- * do not carry this header.
- * @param name - The Ably message name to check.
- * @param headers - The Ably message headers to inspect for discrete marker presence.
- * @returns True if this is a discrete message part, false if it's a lifecycle event.
- */
-const isDiscreteMessagePart = (name: string, headers: Record<string, string>): boolean =>
-  (name === 'text' || name === 'file' || isDataEventName(name)) && HEADER_DISCRETE in headers;
-
-const decodeDiscretePayload = (input: MessagePayload, lifecycle: LifecycleTracker<AI.UIMessageChunk>): Out[] => {
-  const h = input.headers ?? {};
-  const r = headerReader(h);
-  const turnId = h[HEADER_TURN_ID] ?? '';
-
-  // Discrete message parts from writeMessages (user messages, history entries).
-  // Distinguished from lifecycle events by the presence of x-ably-discrete.
-  if (isDiscreteMessagePart(input.name, h)) {
-    return decodeDiscreteMessage(input);
-  }
+const isDiscreteMessagePart = (codecType: string, headers: Record<string, string>): boolean =>
+  (codecType === 'text' || codecType === 'file' || isDataEventName(codecType)) && HEADER_DISCRETE in headers;
 
-  if (input.name === 'tool-input') {
-    return decodeNonStreamingToolInput(r, input.data, turnId, lifecycle);
-  }
+const decodeClientToolResult = (codecMessageId: string, r: VercelHeaderReader, data: unknown): VercelInput[] => {
+  // CAST: Trust boundary — encoder produced the expected object shape.
+  const parsed = data as ToolOutputAvailableWireData | undefined;
+  return [
+    {
+      kind: 'tool-result',
+      codecMessageId,
+      payload: { toolCallId: r.strOr('toolCallId', ''), output: parsed?.output },
+    },
+  ];
+};
+
+const decodeClientToolResultError = (codecMessageId: string, r: VercelHeaderReader, data: unknown): VercelInput[] => {
+  // CAST: Trust boundary — encoder produced the expected object shape.
+  const parsed = data as ClientToolResultErrorWireData | undefined;
+  return [
+    {
+      kind: 'tool-result-error',
+      codecMessageId,
+      payload: { toolCallId: r.strOr('toolCallId', ''), message: parsed?.message ?? '' },
+    },
+  ];
+};
+
+const decodeClientToolApprovalResponse = (codecMessageId: string, r: VercelHeaderReader): VercelInput[] => [
+  {
+    kind: 'tool-approval-response',
+    codecMessageId,
+    payload: stripUndefined({
+      toolCallId: r.strOr('toolCallId', ''),
+      approved: r.bool('approved') ?? false,
+      reason: r.str('reason'),
+    }),
+  },
+];
+
+// ---------------------------------------------------------------------------
+// Discrete payload dispatch
+// ---------------------------------------------------------------------------
 
-  switch (input.name) {
+const decodeAiOutputPayload = (
+  codecType: string,
+  r: VercelHeaderReader,
+  data: unknown,
+  runId: string,
+  lifecycle: LifecycleTracker<AI.UIMessageChunk>,
+): AnyEvent[] => {
+  switch (codecType) {
     case 'start': {
-      return decodeStart(r, turnId, lifecycle);
+      return decodeStart(r, runId, lifecycle);
     }
     case 'start-step': {
-      return decodeStartStep(turnId, lifecycle);
+      return decodeStartStep(runId, lifecycle);
     }
     case 'finish-step': {
-      return decodeFinishStep(turnId, lifecycle);
+      return decodeFinishStep(runId, lifecycle);
     }
     case 'finish': {
-      return decodeFinish(r, turnId, lifecycle);
+      return decodeFinish(r, runId, lifecycle);
     }
     case 'error': {
-      return decodeError(input.data, turnId, lifecycle);
+      return decodeError(data, runId, lifecycle);
     }
     case 'abort': {
-      return decodeAbort(input.data, turnId, lifecycle);
+      return decodeAbort(data, runId, lifecycle);
     }
     case 'message-metadata': {
       return decodeMessageMetadata(r);
     }
     case 'file': {
-      return decodeFile(r, input.data);
+      return decodeFile(r, data);
     }
     case 'source-url': {
-      return decodeSourceUrl(r, input.data);
+      return decodeSourceUrl(r, data);
     }
     case 'source-document': {
       return decodeSourceDocument(r);
     }
+    case 'tool-input': {
+      return decodeNonStreamingToolInput(r, data, runId, lifecycle);
+    }
     case 'tool-input-error': {
-      return decodeToolInputError(r, input.data);
+      return decodeToolInputError(r, data);
     }
     case 'tool-output-available': {
-      return decodeToolOutputAvailable(r, input.data);
+      return decodeAgentToolOutputAvailable(r, data);
     }
     case 'tool-output-error': {
-      return decodeToolOutputError(r, input.data);
+      return decodeAgentToolOutputError(r, data);
     }
     case 'tool-approval-request': {
       return decodeToolApprovalRequest(r);
@@ -557,50 +581,105 @@ const decodeDiscretePayload = (input: MessagePayload, lifecycle: LifecycleTracke
       return decodeToolOutputDenied(r);
     }
     default: {
-      return isDataEventName(input.name) ? decodeDataEvent(input.name, r, input.data) : [];
+      return isDataEventName(codecType) ? decodeDataEvent(codecType, r, data) : [];
+    }
+  }
+};
+
+const decodeAiInputPayload = (codecType: string, input: MessagePayload, r: VercelHeaderReader): AnyEvent[] => {
+  // Multi-part user-message parts (text / file / data-*) carry discrete
+  // because they ride publishDiscreteBatch; the receive-side fans them back
+  // out into a UserMessage.
+  if (isDiscreteMessagePart(codecType, input.transportHeaders ?? {})) {
+    return decodeDiscreteMessagePart(input);
+  }
+
+  const codecMessageId = input.transportHeaders?.[HEADER_CODEC_MESSAGE_ID] ?? '';
+
+  switch (codecType) {
+    case 'tool-result': {
+      return decodeClientToolResult(codecMessageId, r, input.data);
+    }
+    case 'tool-result-error': {
+      return decodeClientToolResultError(codecMessageId, r, input.data);
+    }
+    case 'tool-approval-response': {
+      return decodeClientToolApprovalResponse(codecMessageId, r);
+    }
+    case 'regenerate': {
+      // Wire-only signal — carries `parent` / `msg-regenerate` on transport
+      // headers, no domain payload. The agent's input-event lookup reads
+      // transport headers directly from the inbound Ably message; no
+      // projection fold is needed here.
+      return [];
+    }
+    default: {
+      return [];
     }
   }
 };
 
+const decodeDiscretePayload = (input: MessagePayload, lifecycle: LifecycleTracker<AI.UIMessageChunk>): AnyEvent[] => {
+  const r = headerReader(input.codecHeaders ?? {});
+  const runId = input.transportHeaders?.[HEADER_RUN_ID] ?? '';
+  const codecType = r.strOr('type', '');
+
+  if (input.name === EVENT_AI_INPUT) {
+    return decodeAiInputPayload(codecType, input, r);
+  }
+
+  if (input.name === EVENT_AI_OUTPUT) {
+    return decodeAiOutputPayload(codecType, r, input.data, runId, lifecycle);
+  }
+
+  return [];
+};
+
 // ---------------------------------------------------------------------------
 // Decoder core hooks
 // ---------------------------------------------------------------------------
 
-const createHooks = (
-  lifecycle: LifecycleTracker<AI.UIMessageChunk>,
-): DecoderCoreHooks<AI.UIMessageChunk, AI.UIMessage> => ({
-  buildStartEvents: (tracker: StreamTrackerState): Out[] => {
-    const turnId = tracker.headers[HEADER_TURN_ID] ?? '';
-    const messageId = headerReader(tracker.headers).str('messageId');
-    const outputs = ensurePhases(lifecycle, turnId, { messageId });
-    outputs.push({ kind: 'event', event: buildStartChunk(tracker) });
-    return outputs;
+const createHooks = (lifecycle: LifecycleTracker<AI.UIMessageChunk>): DecoderCoreHooks<AnyEvent> => ({
+  buildStartEvents: (tracker: StreamTrackerState): AnyEvent[] => {
+    const runId = tracker.transportHeaders[HEADER_RUN_ID] ?? '';
+    const messageId = headerReader(tracker.codecHeaders).str('messageId');
+    return [...lifecycle.ensurePhases(runId, { messageId }), buildStartChunk(tracker)];
   },
 
-  buildDeltaEvents: (tracker: StreamTrackerState, delta: string): Out[] => event(buildDeltaChunk(tracker, delta)),
+  buildDeltaEvents: (tracker: StreamTrackerState, delta: string): AnyEvent[] => [buildDeltaChunk(tracker, delta)],
 
-  buildEndEvents: (tracker: StreamTrackerState, closingHeaders: Record<string, string>): Out[] =>
-    event(buildEndChunk(tracker, closingHeaders)),
+  buildEndEvents: (tracker: StreamTrackerState, closingHeaders: Record<string, string>): AnyEvent[] => [
+    buildEndChunk(tracker, closingHeaders),
+  ],
 
-  decodeDiscrete: (payload: MessagePayload): Out[] => decodeDiscretePayload(payload, lifecycle),
+  decodeDiscrete: (payload: MessagePayload): AnyEvent[] => decodeDiscretePayload(payload, lifecycle),
 });
 
 // ---------------------------------------------------------------------------
 // Default implementation
 // ---------------------------------------------------------------------------
 
-class DefaultUIMessageDecoder implements StreamDecoder<AI.UIMessageChunk, AI.UIMessage> {
-  private readonly _core: DecoderCore<AI.UIMessageChunk, AI.UIMessage>;
+const isInput = (event: AnyEvent): event is VercelInput => 'kind' in event;
+
+class DefaultUIMessageDecoder implements Decoder<VercelInput, VercelOutput> {
+  private readonly _core: DecoderCore<AnyEvent>;
 
   constructor(options: DecoderCoreOptions = {}) {
-    this._core = createDecoderCore<AI.UIMessageChunk, AI.UIMessage>(
-      createHooks(createVercelLifecycleTracker()),
-      options,
-    );
+    this._core = createDecoderCore<AnyEvent>(createHooks(createVercelLifecycleTracker()), options);
   }
 
-  decode(message: Ably.InboundMessage): Out[] {
-    return this._core.decode(message);
+  decode(message: Ably.InboundMessage): DecodedMessage<VercelInput, VercelOutput> {
+    const events = this._core.decode(message);
+    const inputs: VercelInput[] = [];
+    const outputs: VercelOutput[] = [];
+    for (const event of events) {
+      if (isInput(event)) {
+        inputs.push(event);
+      } else {
+        outputs.push(event);
+      }
+    }
+    return { inputs, outputs };
   }
 }
 
@@ -609,10 +688,9 @@ class DefaultUIMessageDecoder implements StreamDecoder<AI.UIMessageChunk, AI.UIM
 // ---------------------------------------------------------------------------
 
 /**
- * Create a Vercel AI SDK decoder that maps Ably messages to UIMessageChunk
- * events and UIMessage objects via the decoder core.
+ * Create a Vercel AI SDK decoder that maps Ably messages to {@link DecodedMessage}.
  * @param options - Decoder configuration (callbacks, logger).
- * @returns A {@link StreamDecoder} for UIMessageChunk/UIMessage.
+ * @returns A {@link Decoder} typed in both directions for the Vercel codec.
  */
-export const createDecoder = (options: DecoderCoreOptions = {}): StreamDecoder<AI.UIMessageChunk, AI.UIMessage> =>
+export const createDecoder = (options: DecoderCoreOptions = {}): Decoder<VercelInput, VercelOutput> =>
   new DefaultUIMessageDecoder(options);