@assistant-ui/core 0.3.15 → 0.3.17

package/src/react/runtimes/cloud/AssistantCloudThreadHistoryAdapter.test.tsx

@@ -21,12 +21,33 @@ const mocks = vi.hoisted(() => {
     };
     return {
       threadListItem,
-      threads: { item: vi.fn(() => threadListItem) },
+      threads: {
+        item: vi.fn(() => threadListItem),
+        getState: () => ({ threadItems: [{ id, remoteId }] }),
+      },
+    } as unknown as import("@assistant-ui/store").AssistantClient;
+  };
+
+  const makeSplitClient = (remoteId: string) => {
+    const itemShape = () => ({
+      source: "threads",
+      getState: () => ({ id: remoteId, remoteId }),
+      initialize: async () => ({ remoteId, externalId: undefined }),
+    });
+    const live = itemShape();
+    const listItem = itemShape();
+    return {
+      threadListItem: live,
+      threads: {
+        item: vi.fn(() => listItem),
+        getState: () => ({ threadItems: [{ id: remoteId, remoteId }] }),
+      },
     } as unknown as import("@assistant-ui/store").AssistantClient;
   };
 
   return {
     makeClient,
+    makeSplitClient,
     aui: makeClient("thread-1"),
   };
 });
@@ -188,6 +209,98 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
     );
   });
 
+  it("reports frontend and MCP sources for ai-sdk/v6 tool calls", () => {
+    mocks.aui = mocks.makeClient("thread-1");
+    const cloud = makeCloud();
+    const cloudRef = { current: cloud };
+    const { result } = renderHook(() =>
+      useAssistantCloudThreadHistoryAdapter(cloudRef),
+    );
+    const formatted = result.current.withFormat({
+      format: "ai-sdk/v6",
+      encode: ({ message }) => message,
+      decode: ({ parent_id, content }) => ({
+        parentId: parent_id,
+        message: content as { id: string },
+      }),
+      getId: (message: { id: string }) => message.id,
+    });
+
+    formatted.reportTelemetry([
+      {
+        parentId: null,
+        message: {
+          id: "message-1",
+          role: "assistant",
+          parts: [
+            { type: "step-start" },
+            {
+              type: "tool-search",
+              toolCallId: "static-1",
+              input: { query: "test" },
+              output: { result: "ok" },
+            },
+            {
+              type: "dynamic-tool",
+              toolName: "mcp-search",
+              toolCallId: "dynamic-1",
+              input: { query: "test" },
+              output: { result: "ok" },
+            },
+            { type: "step-start" },
+            {
+              type: "tool-search",
+              toolCallId: "static-2",
+              input: { query: "again" },
+              output: { result: "ok" },
+            },
+          ],
+        },
+      },
+    ]);
+
+    expect(cloud.runs.report).toHaveBeenCalledWith(
+      expect.objectContaining({
+        tool_calls: [
+          expect.objectContaining({
+            tool_call_id: "static-1",
+            tool_source: "frontend",
+          }),
+          expect.objectContaining({
+            tool_call_id: "dynamic-1",
+            tool_source: "mcp",
+          }),
+          expect.objectContaining({
+            tool_call_id: "static-2",
+            tool_source: "frontend",
+          }),
+        ],
+        steps: [
+          {
+            tool_calls: [
+              expect.objectContaining({
+                tool_call_id: "static-1",
+                tool_source: "frontend",
+              }),
+              expect.objectContaining({
+                tool_call_id: "dynamic-1",
+                tool_source: "mcp",
+              }),
+            ],
+          },
+          {
+            tool_calls: [
+              expect.objectContaining({
+                tool_call_id: "static-2",
+                tool_source: "frontend",
+              }),
+            ],
+          },
+        ],
+      }),
+    );
+  });
+
   it("initializes the pinned item instead of the new main thread", async () => {
     mocks.aui = mocks.makeClient(undefined, "new-thread", "thread-1");
     const cloud = makeCloud();
@@ -216,4 +329,47 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
       expect.anything(),
     );
   });
+
+  it("updates a history-loaded message when the list item identity differs from the live item", async () => {
+    mocks.aui = mocks.makeSplitClient("thread-split");
+    const cloud = makeCloud();
+    (cloud.threads.messages.list as ReturnType<typeof vi.fn>).mockResolvedValue(
+      {
+        messages: [
+          {
+            id: "m1",
+            parent_id: null,
+            format: "test",
+            content: { id: "m1" },
+          },
+        ],
+      },
+    );
+    const cloudRef = { current: cloud };
+    const { result } = renderHook(() =>
+      useAssistantCloudThreadHistoryAdapter(cloudRef),
+    );
+    const formatted = result.current.withFormat<
+      { id: string },
+      Record<string, unknown>
+    >({
+      format: "test",
+      encode: ({ message }) => message,
+      decode: ({ parent_id, content }) => ({
+        parentId: parent_id,
+        message: content as { id: string },
+      }),
+      getId: (message) => message.id,
+    });
+
+    formatted.pin!();
+    await formatted.load();
+    await formatted.update!({ parentId: null, message: { id: "m1" } }, "m1");
+
+    expect(cloud.threads.messages.update).toHaveBeenCalledWith(
+      "thread-split",
+      "m1",
+      expect.anything(),
+    );
+  });
 });

package/src/react/runtimes/cloud/AssistantCloudThreadHistoryAdapter.ts

@@ -9,8 +9,13 @@ import type {
 import type { ExportedMessageRepositoryItem } from "../../../runtime/utils/message-repository";
 import {
   type AssistantCloud,
+  type AssistantCloudRunReportToolCall,
   CloudMessagePersistence,
   createFormattedPersistence,
+  createRunTelemetryToolCall,
+  normalizeRunTelemetryUsage,
+  type RunTelemetryUsageInit,
+  truncateRunTelemetryText,
 } from "assistant-cloud";
 import { auiV0Decode, auiV0Encode } from "./auiV0";
 import { type AssistantClient, getClientId, useAui } from "@assistant-ui/store";
@@ -64,7 +69,12 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
     if (!live.source) return undefined;
     const id = live.getState().id;
     if (id === undefined) return undefined;
-    return this.aui.threads.item({ id });
+    // A body can resolve before the list's committed items include its
+    // thread; the live item is already the per-thread anchor in that window.
+    const listed = this.aui.threads
+      .getState()
+      .threadItems.some((item) => item.id === id || item.remoteId === id);
+    return listed ? this.aui.threads.item({ id }) : live;
   }
 
   withFormat<TMessage, TStorageFormat extends Record<string, unknown>>(
@@ -78,8 +88,6 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
       return threadListItem;
     };
     const resolvePinned = () => threadListItem ?? pinCurrent();
-    const getFormatted = () =>
-      createFormattedPersistence(adapter._persistence, formatAdapter);
     const getTargetFormatted = (item: CloudThreadListItem) =>
       createFormattedPersistence(adapter.getPersistence(item), formatAdapter);
     return {
@@ -130,11 +138,14 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
         );
       },
       async load(): Promise<MessageFormatRepository<TMessage>> {
-        pinCurrent();
+        // Loads re-pin and resolve through the pinned item, so the id mapping
+        // they populate lives on the same persistence instance later writes
+        // resolve, whichever of the list item or the live graft won the pin.
+        const pinned = pinCurrent();
         const live = adapter.aui.threadListItem;
         const remoteId = live.source ? live.getState().remoteId : undefined;
         if (!remoteId) return { messages: [] };
-        return getFormatted().load(remoteId);
+        return getTargetFormatted(pinned ?? live).load(remoteId);
       },
     };
   }
@@ -268,85 +279,12 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
   }
 }
 
-const MAX_SPAN_CONTENT = 50_000;
-
-function truncateStr(value: string): string {
-  if (value.length <= MAX_SPAN_CONTENT) return value;
-  return value.slice(0, MAX_SPAN_CONTENT);
-}
-
-function safeStringify(value: unknown): string | undefined {
-  if (value == null) return undefined;
-  try {
-    return truncateStr(JSON.stringify(value));
-  } catch {
-    return undefined;
-  }
-}
-
-type TelemetryToolCall = {
-  tool_name: string;
-  tool_call_id: string;
-  tool_args?: string;
-  tool_result?: string;
-  tool_source?: "mcp" | "frontend" | "backend";
-};
-
-const BASE64_PATTERN = /^[A-Za-z0-9+/]{100,}={0,2}$/;
-
-function summarizeMcpResult(value: unknown): string | undefined {
-  if (value == null) return undefined;
-  try {
-    const parsed = typeof value === "string" ? JSON.parse(value) : value;
-    if (Array.isArray(parsed)) {
-      const summarized = parsed.map((item) => {
-        if (item && typeof item === "object" && item.type) {
-          if (
-            (item.type === "image" || item.type === "audio") &&
-            typeof item.data === "string" &&
-            BASE64_PATTERN.test(item.data.slice(0, 200))
-          ) {
-            const sizeKB = ((item.data.length * 3) / 4 / 1024).toFixed(1);
-            return { ...item, data: `[${item.type}: ${sizeKB}KB]` };
-          }
-        }
-        return item;
-      });
-      return truncateStr(JSON.stringify(summarized));
-    }
-  } catch {
-    // not JSON array, fall through
-  }
-  return safeStringify(value);
-}
-
-function buildToolCall(
-  toolName: string,
-  toolCallId: string,
-  args: unknown,
-  result: unknown,
-  argsText?: string,
-  toolSource?: "mcp" | "frontend" | "backend",
-): TelemetryToolCall {
-  const call: TelemetryToolCall = {
-    tool_name: toolName,
-    tool_call_id: toolCallId,
-  };
-  const toolArgs = argsText ?? safeStringify(args);
-  if (toolArgs !== undefined) call.tool_args = toolArgs;
-  const toolResult =
-    toolSource === "mcp" ? summarizeMcpResult(result) : safeStringify(result);
-  if (toolResult !== undefined) call.tool_result = toolResult;
-  if (toolSource) call.tool_source = toolSource;
-  return call;
-}
-
 type TelemetryStepData = {
   input_tokens?: number;
   output_tokens?: number;
   reasoning_tokens?: number;
   cached_input_tokens?: number;
-  tool_calls?: TelemetryToolCall[];
+  tool_calls?: AssistantCloudRunReportToolCall[];
   start_ms?: number;
   end_ms?: number;
 };
@@ -369,7 +307,7 @@ function mergeStepTimestamps(
 
 type TelemetryData = {
   status: "completed" | "incomplete" | "error";
-  toolCalls?: TelemetryToolCall[];
+  toolCalls?: AssistantCloudRunReportToolCall[];
   totalSteps?: number;
   inputTokens?: number;
   outputTokens?: number;
@@ -426,14 +364,7 @@ export function extractAuiV0<T>(content: T): TelemetryData | null {
     }[];
     metadata?: {
       modelId?: string;
-      steps?: readonly {
-        usage?: {
-          inputTokens?: number;
-          outputTokens?: number;
-          reasoningTokens?: number;
-          cachedInputTokens?: number;
-        };
-      }[];
+      steps?: readonly { usage?: RunTelemetryUsageInit }[];
       custom?: Record<string, unknown> & { modelId?: string };
     };
   };
@@ -446,13 +377,19 @@ export function extractAuiV0<T>(content: T): TelemetryData | null {
   const toolCalls = msg.content
     ?.filter((p) => p.type === "tool-call" && p.toolName && p.toolCallId)
     .map((p) =>
-      buildToolCall(p.toolName!, p.toolCallId!, p.args, p.result, p.argsText),
+      createRunTelemetryToolCall({
+        toolName: p.toolName!,
+        toolCallId: p.toolCallId!,
+        args: p.args,
+        result: p.result,
+        argsText: p.argsText,
+      }),
     );
 
   const textParts = msg.content?.filter((p) => p.type === "text" && p.text);
   const outputText =
     textParts && textParts.length > 0
-      ? truncateStr(textParts.map((p) => p.text).join(""))
+      ? truncateRunTelemetryText(textParts.map((p) => p.text).join(""))
       : undefined;
 
   const steps = msg.metadata?.steps;
@@ -470,20 +407,23 @@ export function extractAuiV0<T>(content: T): TelemetryData | null {
     let hasReasoning = false;
     let hasCachedInput = false;
     for (const step of steps) {
-      if (step.usage?.inputTokens != null) {
-        totalInput += step.usage.inputTokens;
+      if (!step.usage) continue;
+      const usage = normalizeRunTelemetryUsage(step.usage);
+      if (!usage) continue;
+      if (usage.inputTokens != null) {
+        totalInput += usage.inputTokens;
         hasInput = true;
       }
-      if (step.usage?.outputTokens != null) {
-        totalOutput += step.usage.outputTokens;
+      if (usage.outputTokens != null) {
+        totalOutput += usage.outputTokens;
         hasOutput = true;
       }
-      if (step.usage?.reasoningTokens != null) {
-        totalReasoning += step.usage.reasoningTokens;
+      if (usage.reasoningTokens != null) {
+        totalReasoning += usage.reasoningTokens;
         hasReasoning = true;
       }
-      if (step.usage?.cachedInputTokens != null) {
-        totalCachedInput += step.usage.cachedInputTokens;
+      if (usage.cachedInputTokens != null) {
+        totalCachedInput += usage.cachedInputTokens;
         hasCachedInput = true;
       }
     }
@@ -506,20 +446,25 @@ export function extractAuiV0<T>(content: T): TelemetryData | null {
 
   const telemetrySteps: TelemetryStepData[] | undefined =
     steps && steps.length > 1
-      ? steps.map((s) => ({
-          ...(s.usage?.inputTokens != null
-            ? { input_tokens: s.usage.inputTokens }
-            : undefined),
-          ...(s.usage?.outputTokens != null
-            ? { output_tokens: s.usage.outputTokens }
-            : undefined),
-          ...(s.usage?.reasoningTokens != null
-            ? { reasoning_tokens: s.usage.reasoningTokens }
-            : undefined),
-          ...(s.usage?.cachedInputTokens != null
-            ? { cached_input_tokens: s.usage.cachedInputTokens }
-            : undefined),
-        }))
+      ? steps.map((s) => {
+          const usage = s.usage
+            ? normalizeRunTelemetryUsage(s.usage)
+            : undefined;
+          return {
+            ...(usage?.inputTokens != null
+              ? { input_tokens: usage.inputTokens }
+              : undefined),
+            ...(usage?.outputTokens != null
+              ? { output_tokens: usage.outputTokens }
+              : undefined),
+            ...(usage?.reasoningTokens != null
+              ? { reasoning_tokens: usage.reasoningTokens }
+              : undefined),
+            ...(usage?.cachedInputTokens != null
+              ? { cached_input_tokens: usage.cachedInputTokens }
+              : undefined),
+          };
+        })
       : undefined;
 
   return {
@@ -564,29 +509,28 @@ function isDynamicToolPart(p: AiSdkV6Part): boolean {
   return p.type === "dynamic-tool" || p.type.startsWith("dynamic-tool-");
 }
 
-function partToToolCall(p: AiSdkV6Part): TelemetryToolCall {
-  const toolSource: "mcp" | undefined = isDynamicToolPart(p)
+function partToToolCall(p: AiSdkV6Part): AssistantCloudRunReportToolCall {
+  const toolSource: "mcp" | "frontend" = isDynamicToolPart(p)
     ? "mcp"
-    : undefined;
-  return buildToolCall(
-    p.toolName ?? p.type.slice(5),
-    p.toolCallId!,
-    p.args ?? p.input,
-    p.result ?? p.output,
-    undefined,
+    : "frontend";
+  return createRunTelemetryToolCall({
+    toolName: p.toolName ?? p.type.slice(5),
+    toolCallId: p.toolCallId!,
+    args: p.args ?? p.input,
+    result: p.result ?? p.output,
     toolSource,
-  );
+  });
 }
 
 function collectAiSdkV6Parts(parts: readonly AiSdkV6Part[]): {
   textParts: string[];
-  toolCalls: TelemetryToolCall[];
-  stepsData: { tool_calls: TelemetryToolCall[] }[];
+  toolCalls: AssistantCloudRunReportToolCall[];
+  stepsData: { tool_calls: AssistantCloudRunReportToolCall[] }[];
 } {
   const textParts: string[] = [];
-  const toolCalls: TelemetryToolCall[] = [];
-  const stepsData: { tool_calls: TelemetryToolCall[] }[] = [];
-  let currentStepToolCalls: TelemetryToolCall[] | null = null;
+  const toolCalls: AssistantCloudRunReportToolCall[] = [];
+  const stepsData: { tool_calls: AssistantCloudRunReportToolCall[] }[] = [];
+  let currentStepToolCalls: AssistantCloudRunReportToolCall[] | null = null;
 
   for (const p of parts) {
     if (p.type === "step-start") {
@@ -624,10 +568,10 @@ function extractModelId(
 
 function buildAiSdkV6Result(
   textParts: string[],
-  toolCalls: TelemetryToolCall[],
+  toolCalls: AssistantCloudRunReportToolCall[],
   totalSteps: number,
   metadata?: Record<string, unknown>,
-  stepsData?: { tool_calls: TelemetryToolCall[] }[],
+  stepsData?: { tool_calls: AssistantCloudRunReportToolCall[] }[],
   usage?: {
     inputTokens?: number;
     outputTokens?: number;
@@ -636,7 +580,9 @@ function buildAiSdkV6Result(
   },
 ): TelemetryData {
   const hasText = textParts.length > 0;
-  const outputText = hasText ? truncateStr(textParts.join("")) : undefined;
+  const outputText = hasText
+    ? truncateRunTelemetryText(textParts.join(""))
+    : undefined;
   const modelId = extractModelId(metadata);
 
   const steps: TelemetryStepData[] | undefined =
@@ -671,46 +617,6 @@ function buildAiSdkV6Result(
   };
 }
 
-type UsageFields = {
-  inputTokens?: number;
-  outputTokens?: number;
-  promptTokens?: number;
-  completionTokens?: number;
-  reasoningTokens?: number;
-  cachedInputTokens?: number;
-};
-
-function normalizeUsage(u: UsageFields):
-  | {
-      inputTokens?: number;
-      outputTokens?: number;
-      reasoningTokens?: number;
-      cachedInputTokens?: number;
-    }
-  | undefined {
-  const input = u.inputTokens ?? u.promptTokens;
-  const output = u.outputTokens ?? u.completionTokens;
-  if (
-    input == null &&
-    output == null &&
-    u.reasoningTokens == null &&
-    u.cachedInputTokens == null
-  ) {
-    return undefined;
-  }
-
-  return {
-    ...(input != null ? { inputTokens: input } : undefined),
-    ...(output != null ? { outputTokens: output } : undefined),
-    ...(u.reasoningTokens != null
-      ? { reasoningTokens: u.reasoningTokens }
-      : undefined),
-    ...(u.cachedInputTokens != null
-      ? { cachedInputTokens: u.cachedInputTokens }
-      : undefined),
-  };
-}
-
 function extractAiSdkV6Usage(metadata?: Record<string, unknown>):
   | {
       inputTokens?: number;
@@ -720,15 +626,15 @@ function extractAiSdkV6Usage(metadata?: Record<string, unknown>):
     }
   | undefined {
   // Try top-level metadata.usage
-  const usage = metadata?.usage as UsageFields | undefined;
+  const usage = metadata?.usage as RunTelemetryUsageInit | undefined;
   if (usage) {
-    const normalized = normalizeUsage(usage);
+    const normalized = normalizeRunTelemetryUsage(usage);
     if (normalized) return normalized;
   }
 
   // Try aggregating from metadata.steps[].usage
   const steps = metadata?.steps as
-    | readonly { usage?: UsageFields }[]
+    | readonly { usage?: RunTelemetryUsageInit }[]
     | undefined;
   if (steps && steps.length > 0) {
     let inputTokens = 0;
@@ -742,7 +648,7 @@ function extractAiSdkV6Usage(metadata?: Record<string, unknown>):
     let hasAny = false;
     for (const s of steps) {
       if (!s.usage) continue;
-      const n = normalizeUsage(s.usage);
+      const n = normalizeRunTelemetryUsage(s.usage);
       if (n) {
         if (n.inputTokens != null) {
           inputTokens += n.inputTokens;
@@ -795,8 +701,8 @@ function extractAiSdkV6<T>(content: T): TelemetryData | null {
 
 function aggregateAiSdkV6RunSteps<T>(stepMessages: T[]): TelemetryData | null {
   const allTextParts: string[] = [];
-  const allToolCalls: TelemetryToolCall[] = [];
-  const allStepsData: { tool_calls: TelemetryToolCall[] }[] = [];
+  const allToolCalls: AssistantCloudRunReportToolCall[] = [];
+  const allStepsData: { tool_calls: AssistantCloudRunReportToolCall[] }[] = [];
   let hasAssistant = false;
   let metadata: Record<string, unknown> | undefined;
   let inputTokens = 0;

package/src/react/runtimes/cloud/auiV0.ts

@@ -2,6 +2,7 @@ import type {
   MessageStatus,
   SourceProviderMetadata,
   ThreadMessage,
+  ToolApprovalDisplay,
   ToolApprovalOption,
 } from "../../../types/message";
 import type { CompleteAttachment } from "../../../types/attachment";
@@ -16,11 +17,15 @@ import type { ExportedMessageRepositoryItem } from "../../../runtime/utils/messa
 
 type AuiV0ToolApproval = {
   readonly id: string;
+  readonly prompt?: string;
+  readonly display?: ToolApprovalDisplay;
+  readonly allowFreeform?: boolean;
   readonly approved?: boolean;
   readonly reason?: string;
   readonly isAutomatic?: boolean;
   readonly options?: readonly ToolApprovalOption[];
   readonly optionId?: string;
+  readonly text?: string;
   readonly resolution?: "cancelled" | "expired";
 };