@ably/ai-transport 0.1.0 → 0.2.0

package/src/vercel/codec/encoder.ts

@@ -1,137 +1,137 @@
 /**
- * Vercel AI SDK Encoder
+ * Vercel AI SDK encoder.
  *
- * Maps UIMessageChunk events and complete UIMessage objects to Ably channel
- * operations (publish, appendMessage, updateMessage).
+ * Two publish methods enforce direction at the call site:
  *
- * Delegates the message append lifecycle (publish, append, close, abort,
- * flush/recover) to the encoder core. This file contains only the
- * Vercel-specific event-to-operation mapping.
+ * - {@link DefaultUIMessageEncoder.publishInput} encodes a `VercelInput`
+ *   variant and publishes it on the `ai-input` wire.
+ * - {@link DefaultUIMessageEncoder.publishOutput} encodes a `VercelOutput`
+ *   (`AI.UIMessageChunk`) and publishes it on the `ai-output` wire,
+ *   driving the underlying stream-tracker for streamed chunks
+ *   (text / reasoning / tool-input) and falling back to discrete
+ *   publishes for everything else.
  *
- * Domain-specific headers use the `x-domain-` prefix to distinguish them
- * from transport-level `x-ably-` headers.
- *
- * ## Core operations and domain headers
- *
- * Each UIMessageChunk maps to exactly one encoder core operation. Domain
- * headers are passed to every operation that accepts them — the core handles
- * merging, persistence, and deduplication:
- *
- * - **`startStream`**: Opens a message stream. Domain headers become
- *   "persistent headers" — the core repeats them on every subsequent append.
- * - **`appendStream`**: Appends a text delta. Data only, no headers parameter.
- *   The core automatically carries persistent headers from start.
- * - **`closeStream`**: Closes the stream. Pass all domain headers from the
- *   chunk — the core merges them on top of persistent headers, so changed
- *   values (e.g. updated providerMetadata) are picked up and unchanged
- *   values are harmlessly deduplicated.
- * - **`publishDiscrete`**: Publishes a standalone message. All domain headers
- *   for the chunk are passed directly.
+ * The codec event's own discriminator (`kind` for inputs, `type` for
+ * outputs) is carried in the codec tier's `type` header so the
+ * decoder can dispatch. Stream-tracker state lives inside the encoder
+ * core; only the output direction (text / reasoning / tool-input chunks)
+ * drives it — inputs are always published as discrete messages.
  */
 
 import * as Ably from 'ably';
 import type * as AI from 'ai';
 import { isDataUIPart } from 'ai';
 
-import { HEADER_STATUS } from '../../constants.js';
+import { EVENT_AI_INPUT, EVENT_AI_OUTPUT, HEADER_ROLE, HEADER_STATUS } from '../../constants.js';
 import type { EncoderCore, EncoderCoreOptions } from '../../core/codec/encoder.js';
 import { createEncoderCore } from '../../core/codec/encoder.js';
-import type { ChannelWriter, MessagePayload, StreamEncoder, WriteOptions } from '../../core/codec/types.js';
+import type {
+  ChannelWriter,
+  Encoder,
+  MessagePayload,
+  ToolApprovalResponse,
+  ToolResult,
+  ToolResultError,
+  UserMessage,
+  WriteOptions,
+} from '../../core/codec/types.js';
 import { ErrorCode, errorInfoIs } from '../../errors.js';
 import { headerWriter } from '../../utils.js';
+import type {
+  VercelInput,
+  VercelOutput,
+  VercelToolApprovalResponsePayload,
+  VercelToolResultErrorPayload,
+  VercelToolResultPayload,
+} from './events.js';
 
 // ---------------------------------------------------------------------------
-// Discrete event payload builder
+// Default implementation
 // ---------------------------------------------------------------------------
 
-/**
- * Build a MessagePayload for discrete (non-streaming) event types.
- * Used by both `writeEvent` and `appendEvent` for tool output events,
- * content parts, and data-* custom chunks.
- * @param chunk - The UI message chunk to encode.
- * @returns The message payload for publishing to the channel.
- */
-const buildDiscretePayload = (chunk: AI.UIMessageChunk): MessagePayload => {
-  switch (chunk.type) {
-    case 'tool-output-available': {
-      const h = headerWriter()
-        .str('toolCallId', chunk.toolCallId)
-        .bool('dynamic', chunk.dynamic)
-        .bool('providerExecuted', chunk.providerExecuted)
-        .bool('preliminary', chunk.preliminary)
-        .build();
-      return { name: 'tool-output-available', data: { output: chunk.output }, headers: h };
-    }
-
-    case 'tool-output-error': {
-      const h = headerWriter()
-        .str('toolCallId', chunk.toolCallId)
-        .bool('dynamic', chunk.dynamic)
-        .bool('providerExecuted', chunk.providerExecuted)
-        .build();
-      return { name: 'tool-output-error', data: { errorText: chunk.errorText }, headers: h };
-    }
-
-    case 'tool-approval-request': {
-      const h = headerWriter().str('toolCallId', chunk.toolCallId).str('approvalId', chunk.approvalId).build();
-      return { name: 'tool-approval-request', data: '', headers: h };
-    }
+class DefaultUIMessageEncoder implements Encoder<VercelInput, VercelOutput> {
+  private readonly _core: EncoderCore;
+  private readonly _messageId: string | undefined;
+  private _cancelled = false;
 
-    case 'tool-output-denied': {
-      const h = headerWriter().str('toolCallId', chunk.toolCallId).build();
-      return { name: 'tool-output-denied', data: '', headers: h };
-    }
+  constructor(writer: ChannelWriter, options: EncoderCoreOptions = {}) {
+    this._core = createEncoderCore(writer, options);
+    this._messageId = options.messageId;
+  }
 
-    default: {
-      if (chunk.type.startsWith('data-')) {
-        // CAST: data-* chunks always have id, transient, and data fields per AI SDK types.
-        // TypeScript can't narrow the template literal union in a default case.
-        const dataChunk = chunk as Extract<AI.UIMessageChunk, { type: `data-${string}` }>;
-        const h = headerWriter().str('id', dataChunk.id).bool('transient', dataChunk.transient).build();
-        const ephemeral = dataChunk.transient === true;
-        return { name: chunk.type, data: dataChunk.data, headers: h, ephemeral };
+  async publishInput(input: VercelInput, options?: WriteOptions): Promise<void> {
+    switch (input.kind) {
+      case 'user-message': {
+        await this._publishUserMessage(input, options);
+        return;
+      }
+      case 'regenerate': {
+        await this._publishRegenerate(options);
+        return;
+      }
+      case 'tool-result': {
+        await this._publishToolResult(input, options);
+        return;
+      }
+      case 'tool-result-error': {
+        await this._publishToolResultError(input, options);
+        return;
+      }
+      case 'tool-approval-response': {
+        await this._publishToolApprovalResponse(input, options);
+        return;
       }
-      throw new Ably.ErrorInfo(
-        `unable to write event; unsupported chunk type '${chunk.type}'`,
-        ErrorCode.InvalidArgument,
-        400,
-      );
     }
   }
-};
 
-// ---------------------------------------------------------------------------
-// Default implementation
-// ---------------------------------------------------------------------------
+  async publishOutput(output: VercelOutput, options?: WriteOptions): Promise<void> {
+    await this._publishChunk(output, options);
+  }
 
-class DefaultUIMessageEncoder implements StreamEncoder<AI.UIMessageChunk, AI.UIMessage> {
-  private readonly _core: EncoderCore;
-  private readonly _messageId: string | undefined;
-  private _aborted = false;
+  async cancel(reason?: string): Promise<void> {
+    if (this._cancelled) return;
+    this._cancelled = true;
+    await this._core.cancelAllStreams();
+    await this._core.publishDiscrete({
+      name: EVENT_AI_OUTPUT,
+      data: reason ?? '',
+      codecHeaders: headerWriter().str('type', 'abort').build(),
+      transportHeaders: { [HEADER_STATUS]: 'cancelled' },
+    });
+  }
 
-  constructor(writer: ChannelWriter, options: EncoderCoreOptions = {}) {
-    this._core = createEncoderCore(writer, options);
-    this._messageId = options.messageId;
+  async close(): Promise<void> {
+    await this._core.close();
   }
 
-  async appendEvent(chunk: AI.UIMessageChunk, perWrite?: WriteOptions): Promise<void> {
-    switch (chunk.type) {
-      // -- Stream start: open a message stream with persistent headers -------
+  // -------------------------------------------------------------------------
+  // VercelOutput routing — UIMessageChunk
+  // -------------------------------------------------------------------------
 
+  private async _publishChunk(chunk: AI.UIMessageChunk, perWrite?: WriteOptions): Promise<void> {
+    switch (chunk.type) {
+      // -- Stream start -----------------------------------------------------
       case 'text-start': {
-        const h = headerWriter().str('id', chunk.id).json('providerMetadata', chunk.providerMetadata).build();
-        await this._core.startStream(chunk.id, { name: 'text', data: '', headers: h }, perWrite);
-        break;
+        const h = headerWriter()
+          .str('type', 'text')
+          .str('id', chunk.id)
+          .json('providerMetadata', chunk.providerMetadata)
+          .build();
+        await this._core.startStream(chunk.id, { name: EVENT_AI_OUTPUT, data: '', codecHeaders: h }, perWrite);
+        return;
       }
-
       case 'reasoning-start': {
-        const h = headerWriter().str('id', chunk.id).json('providerMetadata', chunk.providerMetadata).build();
-        await this._core.startStream(chunk.id, { name: 'reasoning', data: '', headers: h }, perWrite);
-        break;
+        const h = headerWriter()
+          .str('type', 'reasoning')
+          .str('id', chunk.id)
+          .json('providerMetadata', chunk.providerMetadata)
+          .build();
+        await this._core.startStream(chunk.id, { name: EVENT_AI_OUTPUT, data: '', codecHeaders: h }, perWrite);
+        return;
       }
-
       case 'tool-input-start': {
         const h = headerWriter()
+          .str('type', 'tool-input')
           .str('toolCallId', chunk.toolCallId)
           .str('toolName', chunk.toolName)
           .bool('dynamic', chunk.dynamic)
@@ -139,57 +139,59 @@ class DefaultUIMessageEncoder implements StreamEncoder<AI.UIMessageChunk, AI.UIM
           .bool('providerExecuted', chunk.providerExecuted)
           .json('providerMetadata', chunk.providerMetadata)
           .build();
-        await this._core.startStream(chunk.toolCallId, { name: 'tool-input', data: '', headers: h }, perWrite);
-        break;
+        await this._core.startStream(chunk.toolCallId, { name: EVENT_AI_OUTPUT, data: '', codecHeaders: h }, perWrite);
+        return;
       }
 
-      // -- Stream append: data only, core carries persistent headers --------
-
+      // -- Stream append ----------------------------------------------------
       case 'text-delta': {
         this._core.appendStream(chunk.id, chunk.delta);
-        break;
+        return;
       }
-
       case 'reasoning-delta': {
         this._core.appendStream(chunk.id, chunk.delta);
-        break;
+        return;
       }
-
       case 'tool-input-delta': {
         this._core.appendStream(chunk.toolCallId, chunk.inputTextDelta);
-        break;
+        return;
       }
 
-      // -- Stream close: pass all chunk headers, core merges with persistent
-
+      // -- Stream close -----------------------------------------------------
       case 'text-end': {
-        const h = headerWriter().str('id', chunk.id).json('providerMetadata', chunk.providerMetadata).build();
-        await this._core.closeStream(chunk.id, { name: 'text', data: '', headers: h });
-        break;
+        const h = headerWriter()
+          .str('type', 'text')
+          .str('id', chunk.id)
+          .json('providerMetadata', chunk.providerMetadata)
+          .build();
+        await this._core.closeStream(chunk.id, { name: EVENT_AI_OUTPUT, data: '', codecHeaders: h });
+        return;
       }
-
       case 'reasoning-end': {
-        const h = headerWriter().str('id', chunk.id).json('providerMetadata', chunk.providerMetadata).build();
-        await this._core.closeStream(chunk.id, { name: 'reasoning', data: '', headers: h });
-        break;
+        const h = headerWriter()
+          .str('type', 'reasoning')
+          .str('id', chunk.id)
+          .json('providerMetadata', chunk.providerMetadata)
+          .build();
+        await this._core.closeStream(chunk.id, { name: EVENT_AI_OUTPUT, data: '', codecHeaders: h });
+        return;
       }
-
       case 'tool-input-available': {
-        // If a stream tracker exists, this tool call was streamed — close it.
-        // Otherwise it's a non-streaming tool call — publish discrete.
         try {
           const h = headerWriter()
+            .str('type', 'tool-input')
             .str('toolCallId', chunk.toolCallId)
             .str('toolName', chunk.toolName)
             .json('providerMetadata', chunk.providerMetadata)
             .build();
-          await this._core.closeStream(chunk.toolCallId, { name: 'tool-input', data: '', headers: h });
+          await this._core.closeStream(chunk.toolCallId, { name: EVENT_AI_OUTPUT, data: '', codecHeaders: h });
         } catch (error: unknown) {
-          // Only fall through to discrete for "no active stream" — rethrow real failures
+          // closeStream raises ErrorCode.InvalidArgument when there is no active stream for this id; fall through to a discrete publish in that case and rethrow any other error.
           if (!(error instanceof Ably.ErrorInfo && errorInfoIs(error, ErrorCode.InvalidArgument))) {
             throw error;
           }
           const h = headerWriter()
+            .str('type', 'tool-input')
             .str('toolCallId', chunk.toolCallId)
             .str('toolName', chunk.toolName)
             .bool('dynamic', chunk.dynamic)
@@ -197,60 +199,69 @@ class DefaultUIMessageEncoder implements StreamEncoder<AI.UIMessageChunk, AI.UIM
             .bool('providerExecuted', chunk.providerExecuted)
             .json('providerMetadata', chunk.providerMetadata)
             .build();
-          await this._core.publishDiscrete({ name: 'tool-input', data: chunk.input, headers: h });
+          await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: chunk.input, codecHeaders: h });
         }
-        break;
+        return;
       }
 
-      // -- Discrete: lifecycle events ---------------------------------------
-
+      // -- Lifecycle (discrete) ---------------------------------------------
       case 'start': {
         const h = headerWriter()
+          .str('type', 'start')
           .str('messageId', chunk.messageId ?? this._messageId)
           .json('messageMetadata', chunk.messageMetadata)
           .build();
-        await this._core.publishDiscrete({ name: 'start', data: '', headers: h }, perWrite);
-        break;
+        await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: '', codecHeaders: h }, perWrite);
+        return;
       }
-
       case 'start-step': {
-        await this._core.publishDiscrete({ name: 'start-step', data: '' }, perWrite);
-        break;
+        const h = headerWriter().str('type', 'start-step').build();
+        await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: '', codecHeaders: h }, perWrite);
+        return;
       }
-
       case 'finish-step': {
-        await this._core.publishDiscrete({ name: 'finish-step', data: '' }, perWrite);
-        break;
+        const h = headerWriter().str('type', 'finish-step').build();
+        await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: '', codecHeaders: h }, perWrite);
+        return;
       }
-
       case 'finish': {
         const h = headerWriter()
+          .str('type', 'finish')
           .str('finishReason', chunk.finishReason)
           .json('messageMetadata', chunk.messageMetadata)
           .build();
-        await this._core.publishDiscrete({ name: 'finish', data: '', headers: h }, perWrite);
-        break;
+        await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: '', codecHeaders: h }, perWrite);
+        return;
       }
-
       case 'error': {
-        await this._core.publishDiscrete({ name: 'error', data: chunk.errorText }, perWrite);
-        break;
+        const h = headerWriter().str('type', 'error').build();
+        await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: chunk.errorText, codecHeaders: h }, perWrite);
+        return;
       }
-
       case 'abort': {
-        this._aborted = true;
-        await this._core.abortAllStreams(perWrite);
+        this._cancelled = true;
+        await this._core.cancelAllStreams(perWrite);
         await this._core.publishDiscrete(
-          { name: 'abort', data: chunk.reason ?? '', headers: { [HEADER_STATUS]: 'aborted' } },
+          {
+            name: EVENT_AI_OUTPUT,
+            data: chunk.reason ?? '',
+            codecHeaders: headerWriter().str('type', 'abort').build(),
+            transportHeaders: { [HEADER_STATUS]: 'cancelled' },
+          },
           perWrite,
         );
-        break;
+        return;
+      }
+      case 'message-metadata': {
+        const h = headerWriter().str('type', 'message-metadata').json('messageMetadata', chunk.messageMetadata).build();
+        await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: '', codecHeaders: h }, perWrite);
+        return;
       }
 
-      // -- Discrete: tool lifecycle events ----------------------------------
-
+      // -- Tool lifecycle (discrete) ----------------------------------------
       case 'tool-input-error': {
         const h = headerWriter()
+          .str('type', 'tool-input-error')
           .str('toolCallId', chunk.toolCallId)
           .str('toolName', chunk.toolName)
           .bool('dynamic', chunk.dynamic)
@@ -258,101 +269,215 @@ class DefaultUIMessageEncoder implements StreamEncoder<AI.UIMessageChunk, AI.UIM
           .bool('providerExecuted', chunk.providerExecuted)
           .json('providerMetadata', chunk.providerMetadata)
           .build();
-        await this._core.publishDiscrete({
-          name: 'tool-input-error',
-          data: { errorText: chunk.errorText, input: chunk.input },
-          headers: h,
-        });
-        break;
+        await this._core.publishDiscrete(
+          { name: EVENT_AI_OUTPUT, data: { errorText: chunk.errorText, input: chunk.input }, codecHeaders: h },
+          perWrite,
+        );
+        return;
       }
-
       case 'tool-output-available':
       case 'tool-output-error':
       case 'tool-approval-request':
       case 'tool-output-denied': {
-        await this._core.publishDiscrete(buildDiscretePayload(chunk), perWrite);
-        break;
+        await this._core.publishDiscrete(buildToolOutputPayload(chunk), perWrite);
+        return;
       }
 
-      // -- Discrete: content parts ------------------------------------------
-
+      // -- Content parts (discrete) -----------------------------------------
       case 'file': {
         const h = headerWriter()
+          .str('type', 'file')
           .str('mediaType', chunk.mediaType)
           .json('providerMetadata', chunk.providerMetadata)
           .build();
-        await this._core.publishDiscrete({ name: 'file', data: chunk.url, headers: h }, perWrite);
-        break;
+        await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: chunk.url, codecHeaders: h }, perWrite);
+        return;
       }
-
       case 'source-url': {
         const h = headerWriter()
+          .str('type', 'source-url')
           .str('sourceId', chunk.sourceId)
           .str('title', chunk.title)
           .json('providerMetadata', chunk.providerMetadata)
           .build();
-        await this._core.publishDiscrete({ name: 'source-url', data: chunk.url, headers: h }, perWrite);
-        break;
+        await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: chunk.url, codecHeaders: h }, perWrite);
+        return;
       }
-
       case 'source-document': {
         const h = headerWriter()
+          .str('type', 'source-document')
           .str('sourceId', chunk.sourceId)
           .str('mediaType', chunk.mediaType)
           .str('title', chunk.title)
           .str('filename', chunk.filename)
           .json('providerMetadata', chunk.providerMetadata)
           .build();
-        await this._core.publishDiscrete({ name: 'source-document', data: '', headers: h }, perWrite);
-        break;
-      }
-
-      case 'message-metadata': {
-        const h = headerWriter().json('messageMetadata', chunk.messageMetadata).build();
-        await this._core.publishDiscrete({ name: 'message-metadata', data: '', headers: h }, perWrite);
-        break;
+        await this._core.publishDiscrete({ name: EVENT_AI_OUTPUT, data: '', codecHeaders: h }, perWrite);
+        return;
       }
 
-      // -- Discrete: data-* custom chunks -----------------------------------
-
+      // -- data-* (discrete) ------------------------------------------------
       default: {
         if (chunk.type.startsWith('data-')) {
-          const h = headerWriter().str('id', chunk.id).bool('transient', chunk.transient).build();
-          const ephemeral = chunk.transient === true;
-          await this._core.publishDiscrete({ name: chunk.type, data: chunk.data, headers: h, ephemeral }, perWrite);
+          // CAST: data-* chunks always have id, transient, and data fields per AI SDK types.
+          // TypeScript can't narrow the template literal union in a default case.
+          const dataChunk = chunk;
+          const h = headerWriter()
+            .str('type', dataChunk.type)
+            .str('id', dataChunk.id)
+            .bool('transient', dataChunk.transient)
+            .build();
+          const ephemeral = dataChunk.transient === true;
+          await this._core.publishDiscrete(
+            { name: EVENT_AI_OUTPUT, data: dataChunk.data, codecHeaders: h, ephemeral },
+            perWrite,
+          );
+          return;
         }
-        break;
+        throw new Ably.ErrorInfo(
+          `unable to publish output; unsupported chunk type '${chunk.type}'`,
+          ErrorCode.InvalidArgument,
+          400,
+        );
       }
     }
   }
 
-  async writeEvent(chunk: AI.UIMessageChunk, perWrite?: WriteOptions): Promise<Ably.PublishResult> {
-    return this._core.publishDiscrete(buildDiscretePayload(chunk), perWrite);
+  // -------------------------------------------------------------------------
+  // VercelInput routing
+  // -------------------------------------------------------------------------
+
+  /**
+   * Publish a user-message input as a batch of per-part discrete Ably
+   * messages on the `ai-input` wire. Wire format matches the multi-part
+   * user-message convention; the receive-side decoder fans the parts back
+   * out into a single `UserMessage`.
+   * @param input - The user-message input carrying the UIMessage to encode.
+   * @param perWrite - Optional per-write overrides.
+   */
+  private async _publishUserMessage(input: UserMessage<AI.UIMessage>, perWrite?: WriteOptions): Promise<void> {
+    const payloads = encodeMessagePayloads(input.message);
+    // Stamp role (a transport header) on every payload so the decoder can
+    // reconstruct a `role: 'user'` UIMessage.
+    for (const payload of payloads) {
+      payload.transportHeaders = { ...payload.transportHeaders, [HEADER_ROLE]: 'user' };
+    }
+    await this._core.publishDiscreteBatch(payloads, perWrite);
   }
 
-  async writeMessages(messages: AI.UIMessage[], perWrite?: WriteOptions): Promise<Ably.PublishResult> {
-    const payloads = messages.flatMap((msg) => encodeMessagePayloads(msg));
-    return this._core.publishDiscreteBatch(payloads, perWrite);
+  /**
+   * Publish a regenerate input as a discrete `ai-input` Ably message
+   * carrying codec `type: 'regenerate'`. The wire carries no domain
+   * payload — `parent` / `target` are stamped on the transport headers by
+   * the client-session (it reads them off the input directly and builds
+   * `buildTransportHeaders`).
+   * @param perWrite - Per-write overrides carrying the transport headers built by client-session.
+   */
+  private async _publishRegenerate(perWrite?: WriteOptions): Promise<void> {
+    const h = headerWriter().str('type', 'regenerate').build();
+    await this._core.publishDiscrete({ name: EVENT_AI_INPUT, data: '', codecHeaders: h }, perWrite);
   }
 
-  async abort(reason?: string): Promise<void> {
-    if (this._aborted) return;
-    this._aborted = true;
-    await this._core.abortAllStreams();
-    await this._core.publishDiscrete({
-      name: 'abort',
-      data: reason ?? '',
-      headers: { [HEADER_STATUS]: 'aborted' },
-    });
+  /**
+   * Publish a client-side tool output on the `ai-input` wire. Targets the
+   * assistant addressed by `input.codecMessageId`; the wire's
+   * `codec-message-id` is stamped via `perWrite.messageId` by the
+   * client-session.
+   * @param input - The tool-output input.
+   * @param perWrite - Per-write overrides carrying the wire codecMessageId.
+   */
+  private async _publishToolResult(input: ToolResult<VercelToolResultPayload>, perWrite?: WriteOptions): Promise<void> {
+    const h = headerWriter().str('type', 'tool-result').str('toolCallId', input.payload.toolCallId).build();
+    await this._core.publishDiscrete(
+      { name: EVENT_AI_INPUT, data: { output: input.payload.output }, codecHeaders: h },
+      perWrite,
+    );
   }
 
-  async close(): Promise<void> {
-    await this._core.close();
+  /**
+   * Publish a client-side tool error on the `ai-input` wire. Targets the
+   * assistant addressed by `input.codecMessageId`.
+   * @param input - The tool-result-error input.
+   * @param perWrite - Per-write overrides.
+   */
+  private async _publishToolResultError(
+    input: ToolResultError<VercelToolResultErrorPayload>,
+    perWrite?: WriteOptions,
+  ): Promise<void> {
+    const h = headerWriter().str('type', 'tool-result-error').str('toolCallId', input.payload.toolCallId).build();
+    await this._core.publishDiscrete(
+      { name: EVENT_AI_INPUT, data: { message: input.payload.message }, codecHeaders: h },
+      perWrite,
+    );
+  }
+
+  /**
+   * Publish a client-side tool approval response on the `ai-input` wire.
+   * Targets the assistant addressed by `input.codecMessageId`.
+   * @param input - The approval-response input.
+   * @param perWrite - Per-write overrides.
+   */
+  private async _publishToolApprovalResponse(
+    input: ToolApprovalResponse<VercelToolApprovalResponsePayload>,
+    perWrite?: WriteOptions,
+  ): Promise<void> {
+    const h = headerWriter()
+      .str('type', 'tool-approval-response')
+      .str('toolCallId', input.payload.toolCallId)
+      .bool('approved', input.payload.approved)
+      .str('reason', input.payload.reason)
+      .build();
+    await this._core.publishDiscrete({ name: EVENT_AI_INPUT, data: '', codecHeaders: h }, perWrite);
   }
 }
 
 // ---------------------------------------------------------------------------
-// Message payload encoding (stateless helper)
+// Tool output discrete payload builder (agent-side `ai-output` wire)
+// ---------------------------------------------------------------------------
+
+const buildToolOutputPayload = (
+  chunk: Extract<
+    AI.UIMessageChunk,
+    { type: 'tool-output-available' | 'tool-output-error' | 'tool-approval-request' | 'tool-output-denied' }
+  >,
+): MessagePayload => {
+  switch (chunk.type) {
+    case 'tool-output-available': {
+      const h = headerWriter()
+        .str('type', 'tool-output-available')
+        .str('toolCallId', chunk.toolCallId)
+        .bool('dynamic', chunk.dynamic)
+        .bool('providerExecuted', chunk.providerExecuted)
+        .bool('preliminary', chunk.preliminary)
+        .build();
+      return { name: EVENT_AI_OUTPUT, data: { output: chunk.output }, codecHeaders: h };
+    }
+    case 'tool-output-error': {
+      const h = headerWriter()
+        .str('type', 'tool-output-error')
+        .str('toolCallId', chunk.toolCallId)
+        .bool('dynamic', chunk.dynamic)
+        .bool('providerExecuted', chunk.providerExecuted)
+        .build();
+      return { name: EVENT_AI_OUTPUT, data: { errorText: chunk.errorText }, codecHeaders: h };
+    }
+    case 'tool-approval-request': {
+      const h = headerWriter()
+        .str('type', 'tool-approval-request')
+        .str('toolCallId', chunk.toolCallId)
+        .str('approvalId', chunk.approvalId)
+        .build();
+      return { name: EVENT_AI_OUTPUT, data: '', codecHeaders: h };
+    }
+    case 'tool-output-denied': {
+      const h = headerWriter().str('type', 'tool-output-denied').str('toolCallId', chunk.toolCallId).build();
+      return { name: EVENT_AI_OUTPUT, data: '', codecHeaders: h };
+    }
+  }
+};
+
+// ---------------------------------------------------------------------------
+// User-message per-part payload encoding
 // ---------------------------------------------------------------------------
 
 const encodeMessagePayloads = (message: AI.UIMessage): MessagePayload[] => {
@@ -362,23 +487,31 @@ const encodeMessagePayloads = (message: AI.UIMessage): MessagePayload[] => {
   for (const part of message.parts) {
     switch (part.type) {
       case 'text': {
-        payloads.push({ name: 'text', data: part.text, headers: headerWriter().str('messageId', messageId).build() });
+        payloads.push({
+          name: EVENT_AI_INPUT,
+          data: part.text,
+          codecHeaders: headerWriter().str('type', 'text').str('messageId', messageId).build(),
+        });
         break;
       }
       case 'file': {
         payloads.push({
-          name: 'file',
+          name: EVENT_AI_INPUT,
           data: part.url,
-          headers: headerWriter().str('messageId', messageId).str('mediaType', part.mediaType).build(),
+          codecHeaders: headerWriter()
+            .str('type', 'file')
+            .str('messageId', messageId)
+            .str('mediaType', part.mediaType)
+            .build(),
         });
         break;
       }
       default: {
         if (isDataUIPart(part)) {
           payloads.push({
-            name: part.type,
+            name: EVENT_AI_INPUT,
             data: part.data,
-            headers: headerWriter().str('messageId', messageId).str('id', part.id).build(),
+            codecHeaders: headerWriter().str('type', part.type).str('messageId', messageId).str('id', part.id).build(),
           });
         }
         break;
@@ -387,7 +520,12 @@ const encodeMessagePayloads = (message: AI.UIMessage): MessagePayload[] => {
   }
 
   if (payloads.length === 0) {
-    payloads.push({ name: 'text', data: '', headers: headerWriter().str('messageId', messageId).build() });
+    // Always emit at least one part so the decoder can reconstruct the codec-message-id and role from headers, even when the user-message carried no encodable parts.
+    payloads.push({
+      name: EVENT_AI_INPUT,
+      data: '',
+      codecHeaders: headerWriter().str('type', 'text').str('messageId', messageId).build(),
+    });
   }
 
   return payloads;
@@ -398,13 +536,13 @@ const encodeMessagePayloads = (message: AI.UIMessage): MessagePayload[] => {
 // ---------------------------------------------------------------------------
 
 /**
- * Create a Vercel AI SDK encoder that maps UIMessageChunk events to Ably
- * channel operations via the encoder core.
+ * Create a Vercel AI SDK encoder that maps VercelInput / VercelOutput to
+ * Ably channel operations via the encoder core.
  * @param writer - The channel writer to publish messages through.
  * @param options - Encoder configuration (clientId, extras, hooks, logger).
- * @returns A {@link StreamEncoder} for UIMessageChunk/UIMessage.
+ * @returns An {@link Encoder} typed in both directions for the Vercel codec.
  */
 export const createEncoder = (
   writer: ChannelWriter,
   options: EncoderCoreOptions = {},
-): StreamEncoder<AI.UIMessageChunk, AI.UIMessage> => new DefaultUIMessageEncoder(writer, options);
+): Encoder<VercelInput, VercelOutput> => new DefaultUIMessageEncoder(writer, options);