@ably/ai-transport 0.1.0 → 0.3.0

package/src/vercel/tool-approvals.ts

@@ -1,380 +0,0 @@
-/**
- * Server-side helpers for processing a tool-approval turn.
- *
- * When a Vercel AI SDK tool is marked `needsApproval`, `streamText` pauses
- * after emitting a `dynamic-tool` part in state `approval-requested`. To
- * resume, the server must:
- *
- *   1. Patch the UIMessage history so the pending tool part reflects the
- *      user's decision (`approval-responded` or `output-denied`).
- *   2. Strip the client-appended "Approved: …" user message, because
- *      `streamText`'s multi-step loop only auto-executes pending tool
- *      calls when the conversation ends on a tool/assistant message.
- *   3. Disable `needsApproval` on just-approved tools so the multi-step
- *      loop doesn't immediately pause again on the same tool.
- *   4. Redirect the resulting `tool-output-available` / `tool-output-error`
- *      chunks back to the ORIGINAL assistant message (the one that held
- *      the `approval-requested` part) via `x-ably-amend`, instead of
- *      letting them land on the new assistant message this turn produces.
- *
- * `prepareApprovalTurn` covers steps 1–3; `streamResponseWithApprovalRedirect`
- * covers step 4.
- */
-
-import type * as AI from 'ai';
-import { convertToModelMessages } from 'ai';
-
-import { HEADER_AMEND } from '../constants.js';
-import type { MessageNode, StreamResponseOptions, StreamResult, Turn } from '../core/transport/types.js';
-import { stripUndefined } from '../utils.js';
-import { toolBase } from './codec/tool-transitions.js';
-
-// ---------------------------------------------------------------------------
-// Tool-part transition helpers (private — only used by applyToolApprovalsToHistory)
-// ---------------------------------------------------------------------------
-
-// Build the `approval-responded` variant of a DynamicToolUIPart. Pure.
-const applyApprovalResponseToPart = (
-  part: AI.DynamicToolUIPart,
-  approvalId: string,
-  approved: boolean,
-  reason: string | undefined,
-): AI.DynamicToolUIPart =>
-  stripUndefined({
-    ...toolBase(part),
-    state: 'approval-responded' as const,
-    input: part.input,
-    approval: stripUndefined({ id: approvalId, approved, reason }),
-  });
-
-// Build the `output-denied` variant of a DynamicToolUIPart. Pure.
-const applyApprovalDeniedToPart = (part: AI.DynamicToolUIPart, approvalId: string): AI.DynamicToolUIPart => ({
-  ...toolBase(part),
-  state: 'output-denied',
-  input: part.input,
-  approval: { id: approvalId, approved: false as const },
-});
-
-// ---------------------------------------------------------------------------
-// Wire type
-// ---------------------------------------------------------------------------
-
-/**
- * A user's decision on a pending tool approval. The client ships an array of
- * these to the server in the POST body; the server feeds them to
- * `prepareApprovalTurn` (to patch history) and
- * `streamResponseWithApprovalRedirect` (to route tool outputs back to the
- * original assistant message).
- *
- * Intentionally does not carry `toolName` or `input` — those are redundant
- * with what's already on the UIMessage history part.
- */
-export interface ToolApprovalDecision {
-  /**
-   * The `toolCallId` of the pending `dynamic-tool` part being approved/denied.
-   * Must match a part already in the history; decisions that don't match any
-   * part are ignored by {@link applyToolApprovalsToHistory}.
-   */
-  toolCallId: string;
-  /** Whether the user approved or denied the tool call. */
-  approved: boolean;
-  /**
-   * The `x-ably-msg-id` of the assistant message whose `dynamic-tool` part
-   * is being responded to. When approved and the tool executes successfully,
-   * the output is published cross-turn targeting this message.
-   */
-  targetMsgId: string;
-  /** Optional reason accompanying the response. */
-  reason?: string;
-}
-
-// ---------------------------------------------------------------------------
-// History patching
-// ---------------------------------------------------------------------------
-
-/**
- * Patch `dynamic-tool` parts in the history to reflect a batch of approval
- * decisions. Pure — returns a new array; input is not mutated.
- *
- * Approved decisions transition the matching part to `approval-responded`,
- * which `convertToModelMessages` will expand into a `tool-approval-response`
- * model message for `streamText`'s multi-step loop. Denied decisions
- * transition to `output-denied`.
- *
- * Messages and parts whose `toolCallId` is not referenced by any decision
- * are passed through by reference.
- * @param messages - The UIMessage history (user + assistant messages).
- * @param decisions - Approval decisions keyed by `toolCallId`.
- * @returns A new array with matching tool parts transitioned.
- */
-export const applyToolApprovalsToHistory = (
-  messages: AI.UIMessage[],
-  decisions: ToolApprovalDecision[],
-): AI.UIMessage[] => {
-  if (decisions.length === 0) return messages;
-  const byToolCallId = new Map(decisions.map((d) => [d.toolCallId, d]));
-
-  return messages.map((msg) => {
-    let patchedParts: AI.UIMessage['parts'] | undefined;
-
-    for (const [index, part] of msg.parts.entries()) {
-      if (part.type !== 'dynamic-tool') continue;
-      const decision = byToolCallId.get(part.toolCallId);
-      if (!decision) continue;
-
-      // Preserve an existing approval id if the part already has one
-      // (it was set when the approval-request chunk arrived); otherwise mint
-      // a new id so the emitted tool-approval-response has a stable handle.
-      const approvalId = part.approval?.id ?? crypto.randomUUID();
-      const replacement = decision.approved
-        ? applyApprovalResponseToPart(part, approvalId, true, decision.reason)
-        : applyApprovalDeniedToPart(part, approvalId);
-
-      patchedParts ??= [...msg.parts];
-      patchedParts[index] = replacement;
-    }
-
-    return patchedParts ? { ...msg, parts: patchedParts } : msg;
-  });
-};
-
-// ---------------------------------------------------------------------------
-// Tool manipulation
-// ---------------------------------------------------------------------------
-
-/**
- * Derive the set of tool names that have just been approved by walking the
- * (pre-patch) history for `dynamic-tool` parts whose `toolCallId` matches an
- * approved decision.
- * @param messages - The full UIMessage history.
- * @param decisions - Approval decisions for this request.
- * @returns The set of tool names that were just approved.
- */
-const approvedToolNames = (messages: AI.UIMessage[], decisions: ToolApprovalDecision[]): Set<string> => {
-  const approvedIds = new Set(decisions.filter((d) => d.approved).map((d) => d.toolCallId));
-  if (approvedIds.size === 0) return new Set();
-
-  const names = new Set<string>();
-  for (const msg of messages) {
-    for (const part of msg.parts) {
-      if (part.type === 'dynamic-tool' && approvedIds.has(part.toolCallId)) {
-        names.add(part.toolName);
-      }
-    }
-  }
-  return names;
-};
-
-/**
- * Return a tool dict with `needsApproval: false` forced on any tool whose
- * name is in `approvedNames`. Prevents an infinite approval loop when
- * `streamText`'s multi-step loop calls an approved tool again after
- * executing it.
- *
- * The generic uses `object` (not `AI.Tool`) for its value constraint so
- * duplicate peer-dep resolutions — common when the SDK and the consuming app
- * each pull their own copy of `ai` — still type-check. Every real Vercel Tool
- * is structurally an object, so the constraint holds in practice.
- * @param tools - The tool dictionary.
- * @param approvedNames - Names of tools whose `needsApproval` should be disabled.
- * @returns A new tool dict with the flag cleared on matching entries; input returned unchanged when the set is empty.
- */
-const disableApprovalFor = <T extends Record<string, object>>(tools: T, approvedNames: ReadonlySet<string>): T => {
-  if (approvedNames.size === 0) return tools;
-  const entries = Object.entries(tools).map(([name, def]) =>
-    approvedNames.has(name) ? ([name, { ...def, needsApproval: false }] as const) : ([name, def] as const),
-  );
-  // CAST: Object.fromEntries loses the exact T shape in its return type, but
-  // we preserve every key and only set an existing optional field, so the T
-  // contract holds at runtime.
-  return Object.fromEntries(entries) as T;
-};
-
-// ---------------------------------------------------------------------------
-// Orchestration
-// ---------------------------------------------------------------------------
-
-/** Options for {@link prepareApprovalTurn}. */
-export interface PrepareApprovalTurnOptions<T extends Record<string, object>> {
-  /** The full UIMessage history (user + assistant messages for this conversation). */
-  messages: AI.UIMessage[];
-  /** The user's approval decisions for this request, if any. */
-  decisions: ToolApprovalDecision[] | undefined;
-  /**
-   * The tool dictionary that will be passed to `streamText`. Typed with a
-   * structural `object` value constraint so it accepts `Record<string, Tool>`
-   * regardless of which copy of the `ai` peer dep typed it.
-   */
-  tools: T;
-}
-
-/** Result of {@link prepareApprovalTurn}. */
-export interface PrepareApprovalTurnResult<T extends Record<string, object>> {
-  /** Model-format messages ready to pass to `streamText({ messages })`. */
-  modelMessages: AI.ModelMessage[];
-  /** Tools with `needsApproval` disabled for any tool that was just approved. */
-  tools: T;
-}
-
-/**
- * One-shot transform to ready a history + tool dict for a `streamText` call
- * on an approval turn. Returns the patched model-message array and the
- * effective tools dict.
- *
- * When `decisions` is absent or empty, this is a thin wrapper around
- * `convertToModelMessages(messages)` that returns the original tools — so
- * callers can use it uniformly regardless of whether the request carries
- * approvals.
- * @param options - See {@link PrepareApprovalTurnOptions}.
- * @returns See {@link PrepareApprovalTurnResult}.
- */
-export const prepareApprovalTurn = async <T extends Record<string, object>>(
-  options: PrepareApprovalTurnOptions<T>,
-): Promise<PrepareApprovalTurnResult<T>> => {
-  const { messages, decisions, tools } = options;
-
-  if (!decisions || decisions.length === 0) {
-    return { modelMessages: await convertToModelMessages(messages), tools };
-  }
-
-  const patched = applyToolApprovalsToHistory(messages, decisions);
-  const converted = await convertToModelMessages(patched);
-
-  // Strip the client-appended "Approved: …" / "Denied: …" user message so
-  // `streamText`'s multi-step loop auto-executes the pending tool call.
-  const modelMessages = converted.at(-1)?.role === 'user' ? converted.slice(0, -1) : converted;
-
-  const effectiveTools = disableApprovalFor(tools, approvedToolNames(messages, decisions));
-
-  return { modelMessages, tools: effectiveTools };
-};
-
-// ---------------------------------------------------------------------------
-// Stream response with cross-turn redirect
-// ---------------------------------------------------------------------------
-
-/** Options for {@link streamResponseWithApprovalRedirect}. */
-export interface StreamResponseWithApprovalRedirectOptions extends StreamResponseOptions<AI.UIMessageChunk> {
-  /**
-   * The approval decisions this turn is resolving. Only approved decisions
-   * redirect tool outputs — denied decisions have already been reflected
-   * in the history and produce no tool output to capture.
-   */
-  decisions: ToolApprovalDecision[] | undefined;
-}
-
-/**
- * Pipe a UIMessage chunk stream through the turn's encoder, but redirect
- * `tool-output-available` / `tool-output-error` chunks for approved tools to
- * the original assistant message via `x-ably-amend`.
- *
- * Without this redirect, the tool output would land on the new assistant
- * message produced this turn — leaving the original message stuck in
- * `approval-responded` state. The redirect uses a per-event
- * {@link StreamResponseOptions.resolveWriteOptions} hook: when a matching
- * chunk reaches the encoder, it is published with the target's `msgId`
- * and an `x-ably-amend` header so the client merges the output onto the
- * original message instead of the current-turn one.
- *
- * To preserve "no amendments on cancel" semantics — a partial turn must
- * not leave torn-off tool outputs on the original message — redirect-
- * target chunks are held in a small TransformStream buffer and only
- * released to the encoder when the source stream closes normally. If the
- * turn's `abortSignal` fires before the flush, the buffer is discarded.
- * Non-redirect chunks are enqueued inline and are unaffected by the buffer.
- * @param turn - The active server turn.
- * @param stream - The UIMessage chunk stream to pipe through the encoder.
- * @param options - Stream options plus the approval decisions to redirect.
- * @returns The underlying `streamResponse` result.
- */
-// The redirect-eligible subset of UIMessageChunk — narrow enough for the type
-// guard below to tell TypeScript that `event.toolCallId` is defined.
-type RedirectTargetChunk = Extract<AI.UIMessageChunk, { type: 'tool-output-available' | 'tool-output-error' }>;
-
-export const streamResponseWithApprovalRedirect = (
-  turn: Turn<AI.UIMessageChunk, AI.UIMessage>,
-  stream: ReadableStream<AI.UIMessageChunk>,
-  options: StreamResponseWithApprovalRedirectOptions,
-  // eslint-disable-next-line @typescript-eslint/promise-function-async -- body only returns other promises; an async wrapper would add a pointless microtask hop
-): Promise<StreamResult> => {
-  const { decisions, ...streamOptions } = options;
-
-  const targets = new Map<string, string>();
-  for (const decision of decisions ?? []) {
-    if (decision.approved) targets.set(decision.toolCallId, decision.targetMsgId);
-  }
-
-  if (targets.size === 0) return turn.streamResponse(stream, streamOptions);
-
-  const isRedirectTarget = (event: AI.UIMessageChunk): event is RedirectTargetChunk =>
-    (event.type === 'tool-output-available' || event.type === 'tool-output-error') && targets.has(event.toolCallId);
-
-  const buffer: AI.UIMessageChunk[] = [];
-  const guarded = stream.pipeThrough(
-    new TransformStream<AI.UIMessageChunk, AI.UIMessageChunk>({
-      transform: (chunk, controller) => {
-        if (isRedirectTarget(chunk)) {
-          buffer.push(chunk);
-          return;
-        }
-        controller.enqueue(chunk);
-      },
-      flush: (controller) => {
-        if (turn.abortSignal.aborted) return;
-        for (const chunk of buffer) controller.enqueue(chunk);
-      },
-    }),
-  );
-
-  return turn.streamResponse(guarded, {
-    ...streamOptions,
-    resolveWriteOptions: (event) => {
-      if (!isRedirectTarget(event)) return;
-      const target = targets.get(event.toolCallId);
-      if (target === undefined) return;
-      return { messageId: target, extras: { headers: { [HEADER_AMEND]: target } } };
-    },
-  });
-};
-
-// ---------------------------------------------------------------------------
-// History-scan helper (useChat-style routes)
-// ---------------------------------------------------------------------------
-
-/**
- * Walk the conversation history and synthesize a {@link ToolApprovalDecision}
- * for each `dynamic-tool` part in `approval-responded` (approved) or
- * `output-denied` (denied) state.
- *
- * Use in server routes where the client flips the tool part state directly
- * (via useChat's `addToolApprovalResponse` and our
- * `useStagedAddToolApprovalResponse`) and ships it through the history
- * overlay instead of a separate `toolApprovals` body field.
- * @param history - The conversation history nodes from the POST body.
- * @returns Approval decisions derived from the history, in walk order.
- */
-export const extractApprovalDecisionsFromHistory = (
-  history: readonly MessageNode<AI.UIMessage>[],
-): ToolApprovalDecision[] => {
-  const decisions: ToolApprovalDecision[] = [];
-  for (const node of history) {
-    for (const part of node.message.parts) {
-      if (part.type !== 'dynamic-tool') continue;
-      if (part.state === 'approval-responded') {
-        decisions.push({
-          toolCallId: part.toolCallId,
-          approved: true,
-          targetMsgId: node.msgId,
-          ...(part.approval.reason === undefined ? {} : { reason: part.approval.reason }),
-        });
-      } else if (part.state === 'output-denied') {
-        decisions.push({
-          toolCallId: part.toolCallId,
-          approved: false,
-          targetMsgId: node.msgId,
-        });
-      }
-    }
-  }
-  return decisions;
-};

package/src/vercel/tool-events.ts

@@ -1,53 +0,0 @@
-/**
- * Server-side helper for folding client-shipped events into an in-memory
- * history array before handing it to `convertToModelMessages` / `streamText`.
- *
- * When a client-executed tool resolves, the client stages the resulting
- * `tool-output-available` / `tool-output-error` chunk via
- * `transport.stageEvents(msgId, [...])`. The next send flushes it into the
- * POST body's `events` field. The server republishes the event on the
- * channel via `turn.addEvents`, and must also merge it into the in-memory
- * history so the LLM sees the tool result this turn.
- */
-
-import type * as AI from 'ai';
-
-import type { EventsNode, MessageNode } from '../core/transport/types.js';
-import { createAccumulator } from './codec/accumulator.js';
-
-/**
- * Fold a batch of client-shipped events into an in-memory history array.
- *
- * Mirrors the optimistic tree update in
- * `DefaultClientTransport._internalSend` (src/core/transport/client-transport.ts)
- * so the server can rebuild the same message state before handing it to
- * `convertToModelMessages` / `streamText`.
- * @param events - The events shipped by the client.
- * @param nodes - The history messages from the POST body.
- * @returns A new array with tool-result events applied to the matching
- *   messages. Non-targeted messages are passed through unchanged.
- */
-export const applyToolEventsToHistory = (
-  events: EventsNode<AI.UIMessageChunk>[],
-  nodes: MessageNode<AI.UIMessage>[],
-): MessageNode<AI.UIMessage>[] => {
-  if (events.length === 0) return nodes;
-  const eventsByMsgId = new Map(events.map((e) => [e.msgId, e]));
-
-  return nodes.map((node) => {
-    const evNode = eventsByMsgId.get(node.msgId);
-    if (!evNode) return node;
-
-    const accumulator = createAccumulator();
-    accumulator.initMessage(node.msgId, node.message);
-    accumulator.processOutputs(
-      evNode.events.map((event) => ({
-        kind: 'event' as const,
-        event,
-        messageId: node.msgId,
-      })),
-    );
-    const updated = accumulator.messages.at(-1);
-    return updated ? { ...node, message: updated } : node;
-  });
-};