@convex-dev/agent 0.5.0-alpha.1 → 0.6.0-alpha.0

package/src/deltas.ts

@@ -0,0 +1,570 @@
+import {
+  readUIMessageStream,
+  type DynamicToolUIPart,
+  type ProviderMetadata,
+  type ReasoningUIPart,
+  type TextStreamPart,
+  type TextUIPart,
+  type ToolSet,
+  type ToolUIPart,
+  type UIMessageChunk,
+} from "ai";
+import { assert, pick } from "convex-helpers";
+import { type UIMessage } from "./UIMessages.js";
+import { joinText, sorted } from "./shared.js";
+import {
+  type MessageStatus,
+  type StreamDelta,
+  type StreamMessage,
+} from "./validators.js";
+import { getErrorMessage } from "@ai-sdk/provider-utils";
+
+export function blankUIMessage<METADATA = unknown>(
+  streamMessage: StreamMessage & { metadata?: METADATA },
+  threadId: string,
+): UIMessage<METADATA> {
+  return {
+    id: `stream:${streamMessage.streamId}`,
+    key: `${threadId}-${streamMessage.order}-${streamMessage.stepOrder}`,
+    order: streamMessage.order,
+    stepOrder: streamMessage.stepOrder,
+    status: statusFromStreamStatus(streamMessage.status),
+    agentName: streamMessage.agentName,
+    text: "",
+    _creationTime: Date.now(),
+    role: "assistant",
+    parts: [],
+    ...(streamMessage.metadata ? { metadata: streamMessage.metadata } : {}),
+  };
+}
+
+export function statusFromStreamStatus(
+  status: StreamMessage["status"],
+): MessageStatus | "streaming" {
+  switch (status) {
+    case "streaming":
+      return "streaming";
+    case "finished":
+      return "success";
+    case "aborted":
+      return "failed";
+    default:
+      return "pending";
+  }
+}
+
+export async function updateFromUIMessageChunks(
+  uiMessage: UIMessage,
+  parts: UIMessageChunk[],
+) {
+  const partsStream = new ReadableStream<UIMessageChunk>({
+    start(controller) {
+      for (const part of parts) {
+        controller.enqueue(part);
+      }
+      controller.close();
+    },
+  });
+  let failed = false;
+  let suppressError = false;
+  const messageStream = readUIMessageStream({
+    message: uiMessage,
+    stream: partsStream,
+    onError: (e) => {
+      const errorMessage = e instanceof Error ? e.message : String(e);
+      // Tool invocation errors can be safely ignored when streaming continuation
+      // after tool approval - the stored messages have the complete tool context
+      if (errorMessage.toLowerCase().includes("no tool invocation found")) {
+        // Silently suppress - this is expected after tool approval when the
+        // continuation stream has tool-result without the original tool-call
+        suppressError = true;
+        return;
+      }
+      failed = true;
+      console.error("Error in stream", e);
+    },
+    terminateOnError: true,
+  });
+  let message = uiMessage;
+  try {
+    for await (const messagePart of messageStream) {
+      assert(
+        messagePart.id === message.id,
+        `Expecting to only make one UIMessage in a stream`,
+      );
+      message = messagePart;
+    }
+  } catch (e) {
+    // If we've already handled this error in onError and marked it as suppressed,
+    // don't rethrow - the stored messages provide the fallback
+    if (!suppressError) {
+      throw e;
+    }
+  }
+  if (failed) {
+    message.status = "failed";
+  }
+  message.text = joinText(message.parts);
+  return message;
+}
+
+export async function deriveUIMessagesFromDeltas(
+  threadId: string,
+  streamMessages: StreamMessage[],
+  allDeltas: StreamDelta[],
+): Promise<UIMessage[]> {
+  const messages: UIMessage[] = [];
+  for (const streamMessage of streamMessages) {
+    if (streamMessage.format === "UIMessageChunk") {
+      const { parts } = getParts<UIMessageChunk>(
+        allDeltas.filter((d) => d.streamId === streamMessage.streamId),
+        0,
+      );
+      const uiMessage = await updateFromUIMessageChunks(
+        blankUIMessage(streamMessage, threadId),
+        parts,
+      );
+      // TODO: this fails on partial tool calls
+      messages.push(uiMessage);
+    } else {
+      const [uiMessages] = deriveUIMessagesFromTextStreamParts(
+        threadId,
+        [streamMessage],
+        [],
+        allDeltas,
+      );
+      messages.push(...uiMessages);
+    }
+  }
+  return sorted(messages);
+}
+
+/**
+ *
+ */
+
+export function deriveUIMessagesFromTextStreamParts(
+  threadId: string,
+  streamMessages: StreamMessage[],
+  existingStreams: Array<{
+    streamId: string;
+    cursor: number;
+    message: UIMessage;
+  }>,
+  allDeltas: StreamDelta[],
+): [
+  UIMessage[],
+  Array<{ streamId: string; cursor: number; message: UIMessage }>,
+  boolean,
+] {
+  const newStreams: Array<{
+    streamId: string;
+    cursor: number;
+    message: UIMessage;
+  }> = [];
+  // Seed the existing chunks
+  let changed = false;
+  for (const streamMessage of streamMessages) {
+    const deltas = allDeltas.filter(
+      (d) => d.streamId === streamMessage.streamId,
+    );
+    const existing = existingStreams.find(
+      (s) => s.streamId === streamMessage.streamId,
+    );
+    const [newStream, messageChanged] = updateFromTextStreamParts(
+      threadId,
+      streamMessage,
+      existing,
+      deltas,
+    );
+    newStreams.push(newStream);
+    if (messageChanged) changed = true;
+  }
+  for (const { streamId } of existingStreams) {
+    if (!newStreams.find((s) => s.streamId === streamId)) {
+      // There's a stream that's no longer active.
+      changed = true;
+    }
+  }
+  const messages = sorted(newStreams.map((s) => s.message));
+  return [messages, newStreams, changed];
+}
+
+export function getParts<T extends StreamDelta["parts"][number]>(
+  deltas: StreamDelta[],
+  fromCursor?: number,
+): { parts: T[]; cursor: number } {
+  const parts: T[] = [];
+  let cursor = fromCursor ?? 0;
+  for (const delta of deltas.sort((a, b) => a.start - b.start)) {
+    if (delta.parts.length === 0) {
+      console.debug(`Got delta with no parts: ${JSON.stringify(delta)}`);
+      continue;
+    }
+    if (cursor !== delta.start) {
+      if (cursor >= delta.end) {
+        continue;
+      } else if (cursor < delta.start) {
+        console.warn(
+          `Got delta for stream ${delta.streamId} that has a gap ${cursor} -> ${delta.start}`,
+        );
+        break;
+      } else {
+        throw new Error(
+          `Got unexpected delta for stream ${delta.streamId}: delta: ${delta.start} -> ${delta.end} existing cursor: ${cursor}`,
+        );
+      }
+    }
+    parts.push(...delta.parts);
+    cursor = delta.end;
+  }
+  return { parts, cursor };
+}
+
+/**
+ * This is historically from when we would use the onChunk callback instead of
+ * consuming the full UIMessageStream.
+ */
+
+// exported for testing
+export function updateFromTextStreamParts(
+  threadId: string,
+  streamMessage: StreamMessage,
+  existing:
+    | { streamId: string; cursor: number; message: UIMessage }
+    | undefined,
+  deltas: StreamDelta[],
+): [{ streamId: string; cursor: number; message: UIMessage }, boolean] {
+  const { cursor, parts } = getParts<TextStreamPart<ToolSet>>(
+    deltas,
+    existing?.cursor,
+  );
+  const changed =
+    parts.length > 0 ||
+    (existing &&
+      statusFromStreamStatus(streamMessage.status) !== existing.message.status);
+  const existingMessage =
+    existing?.message ?? blankUIMessage(streamMessage, threadId);
+  if (!changed) {
+    return [
+      existing ?? {
+        streamId: streamMessage.streamId,
+        cursor,
+        message: existingMessage,
+      },
+      false,
+    ];
+  }
+
+  const message: UIMessage = structuredClone(existingMessage);
+  message.status = statusFromStreamStatus(streamMessage.status);
+
+  const textPartsById = new Map<string, TextUIPart>();
+  const toolPartsById = new Map<string, ToolUIPart | DynamicToolUIPart>(
+    message.parts
+      .filter(
+        (p): p is ToolUIPart | DynamicToolUIPart =>
+          p.type.startsWith("tool-") || p.type === "dynamic-tool",
+      )
+      .map((p) => [p.toolCallId, p]),
+  );
+  const reasoningPartsById = new Map<string, ReasoningUIPart>();
+
+  for (const part of parts) {
+    switch (part.type) {
+      case "text-start":
+      case "text-delta": {
+        if (!textPartsById.has(part.id)) {
+          const lastPart = message.parts.at(-1);
+          if (lastPart?.type === "text") {
+            textPartsById.set(part.id, lastPart);
+          } else {
+            const newPart = {
+              type: "text",
+              text: "",
+              providerMetadata: part.providerMetadata,
+            } satisfies TextUIPart;
+            textPartsById.set(part.id, newPart);
+            message.parts.push(newPart);
+          }
+        }
+        if (part.type === "text-delta") {
+          const textPart = textPartsById.get(part.id)!;
+          textPart.text += part.text;
+          textPart.providerMetadata = mergeProviderMetadata(
+            textPart.providerMetadata,
+            part.providerMetadata,
+          );
+        }
+        break;
+      }
+      case "tool-input-start": {
+        let newPart: ToolUIPart | DynamicToolUIPart;
+        if (part.dynamic) {
+          newPart = {
+            type: "dynamic-tool",
+            toolCallId: part.id,
+            toolName: part.toolName,
+            state: "input-streaming",
+            input: "",
+          } satisfies DynamicToolUIPart;
+        } else {
+          newPart = {
+            type: `tool-${part.toolName}`,
+            toolCallId: part.id,
+            state: "input-streaming",
+            input: "",
+            providerExecuted: part.providerExecuted,
+          } satisfies ToolUIPart;
+        }
+        toolPartsById.set(part.id, newPart);
+        message.parts.push(newPart);
+        break;
+      }
+      case "tool-input-delta":
+        {
+          const toUpdate = toolPartsById.get(part.id);
+          assert(
+            toUpdate,
+            `Expected to find tool call part ${part.id} to update`,
+          );
+          toUpdate.input = (toUpdate.input ?? "") + part.delta;
+        }
+        break;
+      case "tool-input-end":
+        {
+          const toUpdate = toolPartsById.get(part.id);
+          assert(
+            toUpdate,
+            `Expected to find tool call part ${part.id} to update`,
+          );
+          toUpdate.state = "input-available";
+          if (part.providerMetadata) {
+            const updatable = toUpdate as Extract<
+              ToolUIPart | DynamicToolUIPart,
+              { state: "input-available" }
+            >;
+            updatable.callProviderMetadata = mergeProviderMetadata(
+              updatable.callProviderMetadata,
+              part.providerMetadata,
+            );
+          }
+        }
+        break;
+      case "tool-call": {
+        let newPart: ToolUIPart | DynamicToolUIPart;
+        if (part.dynamic) {
+          newPart = {
+            type: "dynamic-tool",
+            toolCallId: part.toolCallId,
+            toolName: part.toolName,
+            input: part.input,
+            state: "input-available",
+          };
+        } else {
+          newPart = {
+            type: `tool-${part.toolName}`,
+            toolCallId: part.toolCallId,
+            input: part.input,
+            state: "input-available",
+          };
+          if (part.providerExecuted) {
+            newPart.providerExecuted = part.providerExecuted;
+          }
+        }
+        if (part.providerMetadata) {
+          newPart.callProviderMetadata = part.providerMetadata;
+        }
+        if (toolPartsById.has(part.toolCallId)) {
+          const toUpdate = toolPartsById.get(part.toolCallId)!;
+          Object.assign(toUpdate, newPart);
+        } else {
+          toolPartsById.set(part.toolCallId, newPart);
+          message.parts.push(newPart);
+        }
+        break;
+      }
+      case "tool-result": {
+        const toolCall = toolPartsById.get(part.toolCallId);
+        assert(
+          toolCall,
+          `Expected to find tool call part ${part.toolCallId} to update with result`,
+        );
+        let newPart: ToolUIPart | DynamicToolUIPart;
+        if (toolCall.type === "dynamic-tool") {
+          newPart = {
+            ...toolCall,
+            state: "output-available",
+            input: part.input ?? toolCall.input,
+            output: part.output ?? toolCall.output,
+            ...pick(part, ["preliminary"]),
+          } as DynamicToolUIPart;
+        } else {
+          newPart = {
+            ...toolCall,
+            state: "output-available",
+            input: part.input ?? toolCall.input,
+            output: part.output ?? toolCall.output,
+            preliminary: part.preliminary,
+          } as ToolUIPart;
+        }
+        Object.assign(toolCall, newPart);
+        break;
+      }
+      case "reasoning-start":
+      case "reasoning-delta": {
+        if (!reasoningPartsById.has(part.id)) {
+          const lastPart = message.parts.at(-1);
+          if (lastPart?.type === "reasoning") {
+            reasoningPartsById.set(part.id, lastPart);
+          } else {
+            const newPart = {
+              type: "reasoning",
+              state: "streaming",
+              text: "",
+              providerMetadata: part.providerMetadata,
+            } satisfies ReasoningUIPart;
+            reasoningPartsById.set(part.id, newPart);
+            message.parts.push(newPart);
+          }
+        }
+        const reasoningPart = reasoningPartsById.get(part.id)!;
+        if (part.type === "reasoning-delta") {
+          reasoningPart.text += part.text;
+          reasoningPart.providerMetadata = mergeProviderMetadata(
+            reasoningPart.providerMetadata,
+            part.providerMetadata,
+          );
+        }
+        break;
+      }
+      case "reasoning-end": {
+        const reasoningPart =
+          reasoningPartsById.get(part.id) ??
+          message.parts.find(
+            (p): p is ReasoningUIPart =>
+              p.type === "reasoning" && p.state === "streaming",
+          )!;
+        if (reasoningPart) {
+          reasoningPart.state = "done";
+        } else {
+          console.warn(
+            `Expected to find reasoning part ${part.id} to finish, but found none`,
+          );
+        }
+        break;
+      }
+      case "source":
+        if (part.sourceType === "url") {
+          message.parts.push({
+            type: "source-url",
+            url: part.url,
+            sourceId: part.id,
+            providerMetadata: part.providerMetadata,
+            title: part.title,
+          });
+        } else if (part.sourceType === "document") {
+          message.parts.push({
+            type: "source-document",
+            mediaType: part.mediaType,
+            sourceId: part.id,
+            title: part.title,
+            filename: part.filename,
+            providerMetadata: part.providerMetadata,
+          });
+        } else {
+          console.warn("Got source part with unknown source type", part);
+        }
+        break;
+      case "abort":
+        message.status = "failed";
+        break;
+      case "error":
+        message.status = "failed";
+        console.warn("Generation failed with error", part.error);
+        break;
+      case "tool-error": {
+        const toolPart = toolPartsById.get(part.toolCallId);
+        if (toolPart) {
+          toolPart.errorText = getErrorMessage(part.error);
+        }
+        break;
+      }
+      case "tool-approval-request": {
+        const typedPart = part as unknown as {
+          type: "tool-approval-request";
+          toolCallId: string;
+          approvalId: string;
+        };
+        const toolPart = toolPartsById.get(typedPart.toolCallId);
+        if (toolPart) {
+          toolPart.state = "approval-requested";
+          (toolPart as ToolUIPart & { approval?: object }).approval = {
+            id: typedPart.approvalId,
+          };
+        } else {
+          console.warn(
+            `Expected tool call part ${typedPart.toolCallId} for approval request`,
+          );
+        }
+        break;
+      }
+      case "file":
+      case "text-end":
+      case "finish-step":
+      case "finish":
+      case "raw":
+      case "start-step":
+      case "start":
+        // ignore
+        break;
+      default: {
+        // Exhaustiveness check disabled intentionally for forwards compatibility.
+        // New TextStreamPart types from future AI SDK versions will trigger a
+        // runtime warning rather than a compile error, allowing graceful degradation.
+        // const _: never = part;
+        console.warn(`Received unexpected part: ${JSON.stringify(part)}`);
+        break;
+      }
+    }
+  }
+  // Consider reasoning done once something else happens
+  for (let i = 0; i < message.parts.length - 1; i++) {
+    const part = message.parts[i];
+    if (part.type === "reasoning") {
+      part.state = "done";
+    }
+  }
+  message.text = joinText(message.parts);
+  return [
+    {
+      streamId: streamMessage.streamId,
+      cursor,
+      message,
+    },
+    true,
+  ];
+}
+
+function mergeProviderMetadata(
+  existing: ProviderMetadata | undefined,
+  part: ProviderMetadata | undefined,
+): ProviderMetadata | undefined {
+  if (!existing && !part) {
+    return undefined;
+  }
+  if (!existing) {
+    return part;
+  }
+  if (!part) {
+    return existing;
+  }
+  const merged: ProviderMetadata = existing;
+  for (const [provider, metadata] of Object.entries(part)) {
+    merged[provider] = {
+      ...merged[provider],
+      ...metadata,
+    };
+  }
+  return merged;
+}