@ably/ai-transport 0.1.0 → 0.3.0

package/src/vercel/codec/encoder.ts

@@ -1,410 +0,0 @@
-/**
- * Vercel AI SDK Encoder
- *
- * Maps UIMessageChunk events and complete UIMessage objects to Ably channel
- * operations (publish, appendMessage, updateMessage).
- *
- * 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.
- *
- * 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.
- */
-
-import * as Ably from 'ably';
-import type * as AI from 'ai';
-import { isDataUIPart } from 'ai';
-
-import { 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 { ErrorCode, errorInfoIs } from '../../errors.js';
-import { headerWriter } from '../../utils.js';
-
-// ---------------------------------------------------------------------------
-// Discrete event payload builder
-// ---------------------------------------------------------------------------
-
-/**
- * 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 };
-    }
-
-    case 'tool-output-denied': {
-      const h = headerWriter().str('toolCallId', chunk.toolCallId).build();
-      return { name: 'tool-output-denied', data: '', headers: h };
-    }
-
-    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 };
-      }
-      throw new Ably.ErrorInfo(
-        `unable to write event; unsupported chunk type '${chunk.type}'`,
-        ErrorCode.InvalidArgument,
-        400,
-      );
-    }
-  }
-};
-
-// ---------------------------------------------------------------------------
-// Default implementation
-// ---------------------------------------------------------------------------
-
-class DefaultUIMessageEncoder implements StreamEncoder<AI.UIMessageChunk, AI.UIMessage> {
-  private readonly _core: EncoderCore;
-  private readonly _messageId: string | undefined;
-  private _aborted = false;
-
-  constructor(writer: ChannelWriter, options: EncoderCoreOptions = {}) {
-    this._core = createEncoderCore(writer, options);
-    this._messageId = options.messageId;
-  }
-
-  async appendEvent(chunk: AI.UIMessageChunk, perWrite?: WriteOptions): Promise<void> {
-    switch (chunk.type) {
-      // -- Stream start: open a message stream with persistent headers -------
-
-      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;
-      }
-
-      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;
-      }
-
-      case 'tool-input-start': {
-        const h = headerWriter()
-          .str('toolCallId', chunk.toolCallId)
-          .str('toolName', chunk.toolName)
-          .bool('dynamic', chunk.dynamic)
-          .str('title', chunk.title)
-          .bool('providerExecuted', chunk.providerExecuted)
-          .json('providerMetadata', chunk.providerMetadata)
-          .build();
-        await this._core.startStream(chunk.toolCallId, { name: 'tool-input', data: '', headers: h }, perWrite);
-        break;
-      }
-
-      // -- Stream append: data only, core carries persistent headers --------
-
-      case 'text-delta': {
-        this._core.appendStream(chunk.id, chunk.delta);
-        break;
-      }
-
-      case 'reasoning-delta': {
-        this._core.appendStream(chunk.id, chunk.delta);
-        break;
-      }
-
-      case 'tool-input-delta': {
-        this._core.appendStream(chunk.toolCallId, chunk.inputTextDelta);
-        break;
-      }
-
-      // -- Stream close: pass all chunk headers, core merges with persistent
-
-      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;
-      }
-
-      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;
-      }
-
-      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('toolCallId', chunk.toolCallId)
-            .str('toolName', chunk.toolName)
-            .json('providerMetadata', chunk.providerMetadata)
-            .build();
-          await this._core.closeStream(chunk.toolCallId, { name: 'tool-input', data: '', headers: h });
-        } catch (error: unknown) {
-          // Only fall through to discrete for "no active stream" — rethrow real failures
-          if (!(error instanceof Ably.ErrorInfo && errorInfoIs(error, ErrorCode.InvalidArgument))) {
-            throw error;
-          }
-          const h = headerWriter()
-            .str('toolCallId', chunk.toolCallId)
-            .str('toolName', chunk.toolName)
-            .bool('dynamic', chunk.dynamic)
-            .str('title', chunk.title)
-            .bool('providerExecuted', chunk.providerExecuted)
-            .json('providerMetadata', chunk.providerMetadata)
-            .build();
-          await this._core.publishDiscrete({ name: 'tool-input', data: chunk.input, headers: h });
-        }
-        break;
-      }
-
-      // -- Discrete: lifecycle events ---------------------------------------
-
-      case 'start': {
-        const h = headerWriter()
-          .str('messageId', chunk.messageId ?? this._messageId)
-          .json('messageMetadata', chunk.messageMetadata)
-          .build();
-        await this._core.publishDiscrete({ name: 'start', data: '', headers: h }, perWrite);
-        break;
-      }
-
-      case 'start-step': {
-        await this._core.publishDiscrete({ name: 'start-step', data: '' }, perWrite);
-        break;
-      }
-
-      case 'finish-step': {
-        await this._core.publishDiscrete({ name: 'finish-step', data: '' }, perWrite);
-        break;
-      }
-
-      case 'finish': {
-        const h = headerWriter()
-          .str('finishReason', chunk.finishReason)
-          .json('messageMetadata', chunk.messageMetadata)
-          .build();
-        await this._core.publishDiscrete({ name: 'finish', data: '', headers: h }, perWrite);
-        break;
-      }
-
-      case 'error': {
-        await this._core.publishDiscrete({ name: 'error', data: chunk.errorText }, perWrite);
-        break;
-      }
-
-      case 'abort': {
-        this._aborted = true;
-        await this._core.abortAllStreams(perWrite);
-        await this._core.publishDiscrete(
-          { name: 'abort', data: chunk.reason ?? '', headers: { [HEADER_STATUS]: 'aborted' } },
-          perWrite,
-        );
-        break;
-      }
-
-      // -- Discrete: tool lifecycle events ----------------------------------
-
-      case 'tool-input-error': {
-        const h = headerWriter()
-          .str('toolCallId', chunk.toolCallId)
-          .str('toolName', chunk.toolName)
-          .bool('dynamic', chunk.dynamic)
-          .str('title', chunk.title)
-          .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;
-      }
-
-      case 'tool-output-available':
-      case 'tool-output-error':
-      case 'tool-approval-request':
-      case 'tool-output-denied': {
-        await this._core.publishDiscrete(buildDiscretePayload(chunk), perWrite);
-        break;
-      }
-
-      // -- Discrete: content parts ------------------------------------------
-
-      case 'file': {
-        const h = headerWriter()
-          .str('mediaType', chunk.mediaType)
-          .json('providerMetadata', chunk.providerMetadata)
-          .build();
-        await this._core.publishDiscrete({ name: 'file', data: chunk.url, headers: h }, perWrite);
-        break;
-      }
-
-      case 'source-url': {
-        const h = headerWriter()
-          .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;
-      }
-
-      case 'source-document': {
-        const h = headerWriter()
-          .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;
-      }
-
-      // -- Discrete: data-* custom chunks -----------------------------------
-
-      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);
-        }
-        break;
-      }
-    }
-  }
-
-  async writeEvent(chunk: AI.UIMessageChunk, perWrite?: WriteOptions): Promise<Ably.PublishResult> {
-    return this._core.publishDiscrete(buildDiscretePayload(chunk), perWrite);
-  }
-
-  async writeMessages(messages: AI.UIMessage[], perWrite?: WriteOptions): Promise<Ably.PublishResult> {
-    const payloads = messages.flatMap((msg) => encodeMessagePayloads(msg));
-    return this._core.publishDiscreteBatch(payloads, 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' },
-    });
-  }
-
-  async close(): Promise<void> {
-    await this._core.close();
-  }
-}
-
-// ---------------------------------------------------------------------------
-// Message payload encoding (stateless helper)
-// ---------------------------------------------------------------------------
-
-const encodeMessagePayloads = (message: AI.UIMessage): MessagePayload[] => {
-  const messageId = message.id;
-  const payloads: 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() });
-        break;
-      }
-      case 'file': {
-        payloads.push({
-          name: 'file',
-          data: part.url,
-          headers: headerWriter().str('messageId', messageId).str('mediaType', part.mediaType).build(),
-        });
-        break;
-      }
-      default: {
-        if (isDataUIPart(part)) {
-          payloads.push({
-            name: part.type,
-            data: part.data,
-            headers: headerWriter().str('messageId', messageId).str('id', part.id).build(),
-          });
-        }
-        break;
-      }
-    }
-  }
-
-  if (payloads.length === 0) {
-    payloads.push({ name: 'text', data: '', headers: headerWriter().str('messageId', messageId).build() });
-  }
-
-  return payloads;
-};
-
-// ---------------------------------------------------------------------------
-// Factory
-// ---------------------------------------------------------------------------
-
-/**
- * Create a Vercel AI SDK encoder that maps UIMessageChunk events 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.
- */
-export const createEncoder = (
-  writer: ChannelWriter,
-  options: EncoderCoreOptions = {},
-): StreamEncoder<AI.UIMessageChunk, AI.UIMessage> => new DefaultUIMessageEncoder(writer, options);

package/src/vercel/react/use-staged-add-tool-approval-response.ts

@@ -1,87 +0,0 @@
-/**
- * useStagedAddToolApprovalResponse — wrap useChat's `addToolApprovalResponse`
- * so the approval response is also applied to the transport tree
- * synchronously at click time.
- *
- * Patching the tree at click time eliminates the useChat↔tree divergence
- * the ChatTransport would otherwise have to reconcile via a history
- * overlay, and closes the observer-turn race that could wipe the
- * approval state between `addToolApprovalResponse` and
- * `sendAutomaticallyWhen`'s evaluation.
- *
- * Use this in place of useChat's raw `addToolApprovalResponse` wherever
- * you wire Approve / Deny buttons.
- */
-
-import type * as AI from 'ai';
-import type { ChatAddToolApproveResponseFunction } from 'ai';
-import { useCallback } from 'react';
-
-import type { ClientTransport } from '../../core/transport/types.js';
-
-/**
- * Returns a function with the same signature as useChat's
- * `addToolApprovalResponse`, but additionally applies the approval
- * response to the transport tree via `stageMessage` before delegating.
- *
- * If the tool call identified by `opts.id` isn't found in the tree,
- * the tree update is skipped and the raw function is still called —
- * matches useChat's tolerant behavior for stale approval ids.
- * @param transport - The client transport whose tree to patch.
- * @param addToolApprovalResponse - The raw function from `useChat()`.
- * @returns A drop-in replacement that patches the tree then delegates.
- */
-export const useStagedAddToolApprovalResponse = (
-  transport: ClientTransport<AI.UIMessageChunk, AI.UIMessage>,
-  addToolApprovalResponse: ChatAddToolApproveResponseFunction,
-): ChatAddToolApproveResponseFunction =>
-  useCallback<ChatAddToolApproveResponseFunction>(
-    (opts) => {
-      stageApprovalResponseOnTree(transport, opts);
-      return addToolApprovalResponse(opts);
-    },
-    [transport, addToolApprovalResponse],
-  );
-
-/**
- * Locate the assistant message whose `dynamic-tool` part carries the
- * given `approval.id`, build a patched copy with the part transitioned
- * to `approval-responded`, and stage the patched message on the tree.
- * @param transport - The transport whose tree to patch.
- * @param opts - The approval response being applied.
- * @param opts.id - The approval id matching a dynamic-tool part in the tree.
- * @param opts.approved - Whether the user approved or denied.
- * @param opts.reason - Optional reason accompanying the response.
- */
-const stageApprovalResponseOnTree = (
-  transport: ClientTransport<AI.UIMessageChunk, AI.UIMessage>,
-  opts: { id: string; approved: boolean; reason?: string },
-): void => {
-  const nodes = transport.view.flattenNodes();
-  for (const node of nodes) {
-    const partIndex = node.message.parts.findIndex((p) => p.type === 'dynamic-tool' && p.approval?.id === opts.id);
-    if (partIndex === -1) continue;
-
-    // CAST: findIndex predicate above narrows this to a dynamic-tool part
-    // with a non-undefined approval.
-    const part = node.message.parts[partIndex] as AI.DynamicToolUIPart;
-
-    // Build the approval-responded variant directly rather than spreading
-    // `part`, which TypeScript narrows to whichever source-state variant
-    // the union discriminator inferred and then rejects when we change
-    // `state` to a variant with different approval/output constraints.
-    const patchedPart: AI.DynamicToolUIPart = {
-      type: 'dynamic-tool',
-      toolName: part.toolName,
-      toolCallId: part.toolCallId,
-      state: 'approval-responded',
-      input: part.input,
-      approval: { id: opts.id, approved: opts.approved, reason: opts.reason },
-    };
-
-    const patchedParts = [...node.message.parts];
-    patchedParts[partIndex] = patchedPart;
-    transport.stageMessage(node.msgId, { ...node.message, parts: patchedParts });
-    return;
-  }
-};