@assistant-ui/react 0.12.12 → 0.12.15

package/src/legacy-runtime/cloud/AssistantCloudThreadHistoryAdapter.ts

@@ -64,10 +64,12 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
           stepTimestamps?: StepTimestamp[];
         },
       ) {
-        const encodedContents = items.map((item) => formatAdapter.encode(item));
-        adapter._reportBatchTelemetry(
+        const encodedRunMessages = items.map((item) =>
+          formatAdapter.encode(item),
+        );
+        adapter._reportRunTelemetry(
           formatAdapter.format,
-          encodedContents,
+          encodedRunMessages,
           options,
         );
       },
@@ -109,9 +111,9 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
     };
   }
 
-  private _reportBatchTelemetry<T>(
+  private _reportRunTelemetry<T>(
     format: string,
-    contents: T[],
+    runMessages: T[],
     options?: {
       durationMs?: number;
       stepTimestamps?: StepTimestamp[];
@@ -122,7 +124,7 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
     const remoteId = this.aui.threadListItem().getState().remoteId;
     if (!remoteId) return;
 
-    const extracted = extractBatchTelemetry(format, contents);
+    const extracted = extractRunTelemetry(format, runMessages);
     if (!extracted) return;
 
     this._sendReport(
@@ -147,6 +149,8 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
     stepTimestamps?: StepTimestamp[],
   ) {
     const mergedSteps = mergeStepTimestamps(data.steps, stepTimestamps);
+    // Keep in sync with assistant-cloud createRunSchema
+    // (apps/aui-cloud-api/src/endpoints/runs/create.ts).
     const initial: Parameters<typeof this.cloudRef.current.runs.report>[0] = {
       thread_id: remoteId,
       status: data.status,
@@ -161,6 +165,12 @@ class AssistantCloudThreadHistoryAdapter implements ThreadHistoryAdapter {
       ...(data.outputTokens != null
         ? { output_tokens: data.outputTokens }
         : undefined),
+      ...(data.reasoningTokens != null
+        ? { reasoning_tokens: data.reasoningTokens }
+        : undefined),
+      ...(data.cachedInputTokens != null
+        ? { cached_input_tokens: data.cachedInputTokens }
+        : undefined),
       ...(durationMs != null ? { duration_ms: durationMs } : undefined),
       ...(data.outputText != null
         ? { output_text: data.outputText }
@@ -253,6 +263,8 @@ function buildToolCall(
 type TelemetryStepData = {
   input_tokens?: number;
   output_tokens?: number;
+  reasoning_tokens?: number;
+  cached_input_tokens?: number;
   tool_calls?: TelemetryToolCall[];
   start_ms?: number;
   end_ms?: number;
@@ -280,6 +292,8 @@ type TelemetryData = {
   totalSteps?: number;
   inputTokens?: number;
   outputTokens?: number;
+  reasoningTokens?: number;
+  cachedInputTokens?: number;
   outputText?: string;
   metadata?: Record<string, unknown>;
   steps?: TelemetryStepData[];
@@ -297,15 +311,15 @@ function extractTelemetry<T>(format: string, content: T): TelemetryData | null {
   }
 }
 
-function extractBatchTelemetry<T>(
+function extractRunTelemetry<T>(
   format: string,
-  contents: T[],
+  runMessages: T[],
 ): TelemetryData | null {
   if (format === "ai-sdk/v6") {
-    return extractAiSdkV6Batch(contents);
+    return aggregateAiSdkV6RunSteps(runMessages);
   }
-  for (let i = contents.length - 1; i >= 0; i--) {
-    const result = extractTelemetry(format, contents[i]!);
+  for (let i = runMessages.length - 1; i >= 0; i--) {
+    const result = extractTelemetry(format, runMessages[i]!);
     if (result) return result;
   }
   return null;
@@ -332,7 +346,12 @@ function extractAuiV0<T>(content: T): TelemetryData | null {
     metadata?: {
       modelId?: string;
       steps?: readonly {
-        usage?: { inputTokens?: number; outputTokens?: number };
+        usage?: {
+          inputTokens?: number;
+          outputTokens?: number;
+          reasoningTokens?: number;
+          cachedInputTokens?: number;
+        };
       }[];
       custom?: Record<string, unknown> & { modelId?: string };
     };
@@ -355,13 +374,39 @@ function extractAuiV0<T>(content: T): TelemetryData | null {
   const steps = msg.metadata?.steps;
   let inputTokens: number | undefined;
   let outputTokens: number | undefined;
+  let reasoningTokens: number | undefined;
+  let cachedInputTokens: number | undefined;
   if (steps && steps.length > 0) {
-    inputTokens = 0;
-    outputTokens = 0;
+    let totalInput = 0;
+    let totalOutput = 0;
+    let totalReasoning = 0;
+    let totalCachedInput = 0;
+    let hasInput = false;
+    let hasOutput = false;
+    let hasReasoning = false;
+    let hasCachedInput = false;
     for (const step of steps) {
-      inputTokens += step.usage?.inputTokens ?? 0;
-      outputTokens += step.usage?.outputTokens ?? 0;
+      if (step.usage?.inputTokens != null) {
+        totalInput += step.usage.inputTokens;
+        hasInput = true;
+      }
+      if (step.usage?.outputTokens != null) {
+        totalOutput += step.usage.outputTokens;
+        hasOutput = true;
+      }
+      if (step.usage?.reasoningTokens != null) {
+        totalReasoning += step.usage.reasoningTokens;
+        hasReasoning = true;
+      }
+      if (step.usage?.cachedInputTokens != null) {
+        totalCachedInput += step.usage.cachedInputTokens;
+        hasCachedInput = true;
+      }
     }
+    inputTokens = hasInput ? totalInput : undefined;
+    outputTokens = hasOutput ? totalOutput : undefined;
+    reasoningTokens = hasReasoning ? totalReasoning : undefined;
+    cachedInputTokens = hasCachedInput ? totalCachedInput : undefined;
   }
 
   const statusType = msg.status?.type;
@@ -384,6 +429,12 @@ function extractAuiV0<T>(content: T): TelemetryData | null {
           ...(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),
         }))
       : undefined;
 
@@ -393,6 +444,8 @@ function extractAuiV0<T>(content: T): TelemetryData | null {
     ...(steps?.length ? { totalSteps: steps.length } : undefined),
     ...(inputTokens != null ? { inputTokens } : undefined),
     ...(outputTokens != null ? { outputTokens } : undefined),
+    ...(reasoningTokens != null ? { reasoningTokens } : undefined),
+    ...(cachedInputTokens != null ? { cachedInputTokens } : undefined),
     ...(outputText != null ? { outputText } : undefined),
     ...(metadata ? { metadata } : undefined),
     ...(telemetrySteps ? { steps: telemetrySteps } : undefined),
@@ -491,7 +544,12 @@ function buildAiSdkV6Result(
   totalSteps: number,
   metadata?: Record<string, unknown>,
   stepsData?: { tool_calls: TelemetryToolCall[] }[],
-  usage?: { inputTokens?: number; outputTokens?: number },
+  usage?: {
+    inputTokens?: number;
+    outputTokens?: number;
+    reasoningTokens?: number;
+    cachedInputTokens?: number;
+  },
 ): TelemetryData {
   const hasText = textParts.length > 0;
   const outputText = hasText ? truncateStr(textParts.join("")) : undefined;
@@ -516,6 +574,12 @@ function buildAiSdkV6Result(
     ...(usage?.outputTokens != null
       ? { outputTokens: usage.outputTokens }
       : undefined),
+    ...(usage?.reasoningTokens != null
+      ? { reasoningTokens: usage.reasoningTokens }
+      : undefined),
+    ...(usage?.cachedInputTokens != null
+      ? { cachedInputTokens: usage.cachedInputTokens }
+      : undefined),
     ...(outputText != null ? { outputText } : undefined),
     ...(metadata ? { metadata } : undefined),
     ...(steps ? { steps } : undefined),
@@ -528,23 +592,49 @@ type UsageFields = {
   outputTokens?: number;
   promptTokens?: number;
   completionTokens?: number;
+  reasoningTokens?: number;
+  cachedInputTokens?: number;
 };
 
-function normalizeUsage(
-  u: UsageFields,
-): { inputTokens: number; outputTokens: number } | undefined {
+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) return undefined;
+  if (
+    input == null &&
+    output == null &&
+    u.reasoningTokens == null &&
+    u.cachedInputTokens == null
+  ) {
+    return undefined;
+  }
+
   return {
-    inputTokens: input ?? 0,
-    outputTokens: output ?? 0,
+    ...(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; outputTokens?: number } | undefined {
+function extractAiSdkV6Usage(metadata?: Record<string, unknown>):
+  | {
+      inputTokens?: number;
+      outputTokens?: number;
+      reasoningTokens?: number;
+      cachedInputTokens?: number;
+    }
+  | undefined {
   // Try top-level metadata.usage
   const usage = metadata?.usage as UsageFields | undefined;
   if (usage) {
@@ -559,17 +649,44 @@ function extractAiSdkV6Usage(
   if (steps && steps.length > 0) {
     let inputTokens = 0;
     let outputTokens = 0;
+    let reasoningTokens = 0;
+    let cachedInputTokens = 0;
+    let hasInput = false;
+    let hasOutput = false;
+    let hasReasoning = false;
+    let hasCachedInput = false;
     let hasAny = false;
     for (const s of steps) {
       if (!s.usage) continue;
       const n = normalizeUsage(s.usage);
       if (n) {
-        inputTokens += n.inputTokens;
-        outputTokens += n.outputTokens;
+        if (n.inputTokens != null) {
+          inputTokens += n.inputTokens;
+          hasInput = true;
+        }
+        if (n.outputTokens != null) {
+          outputTokens += n.outputTokens;
+          hasOutput = true;
+        }
+        if (n.reasoningTokens != null) {
+          reasoningTokens += n.reasoningTokens;
+          hasReasoning = true;
+        }
+        if (n.cachedInputTokens != null) {
+          cachedInputTokens += n.cachedInputTokens;
+          hasCachedInput = true;
+        }
         hasAny = true;
       }
     }
-    if (hasAny) return { inputTokens, outputTokens };
+    if (hasAny) {
+      return {
+        ...(hasInput ? { inputTokens } : undefined),
+        ...(hasOutput ? { outputTokens } : undefined),
+        ...(hasReasoning ? { reasoningTokens } : undefined),
+        ...(hasCachedInput ? { cachedInputTokens } : undefined),
+      };
+    }
   }
 
   return undefined;
@@ -592,17 +709,22 @@ function extractAiSdkV6<T>(content: T): TelemetryData | null {
   );
 }
 
-function extractAiSdkV6Batch<T>(contents: T[]): TelemetryData | null {
+function aggregateAiSdkV6RunSteps<T>(stepMessages: T[]): TelemetryData | null {
   const allTextParts: string[] = [];
   const allToolCalls: TelemetryToolCall[] = [];
   const allStepsData: { tool_calls: TelemetryToolCall[] }[] = [];
   let hasAssistant = false;
   let metadata: Record<string, unknown> | undefined;
-  let aggregatedUsage:
-    | { inputTokens: number; outputTokens: number }
-    | undefined;
-
-  for (const content of contents) {
+  let inputTokens = 0;
+  let outputTokens = 0;
+  let reasoningTokens = 0;
+  let cachedInputTokens = 0;
+  let hasInput = false;
+  let hasOutput = false;
+  let hasReasoning = false;
+  let hasCachedInput = false;
+
+  for (const content of stepMessages) {
     const msg = content as AiSdkV6Message;
     if (msg.role !== "assistant") continue;
     hasAssistant = true;
@@ -617,10 +739,22 @@ function extractAiSdkV6Batch<T>(contents: T[]): TelemetryData | null {
 
     const usage = extractAiSdkV6Usage(msg.metadata);
     if (usage) {
-      if (!aggregatedUsage)
-        aggregatedUsage = { inputTokens: 0, outputTokens: 0 };
-      aggregatedUsage.inputTokens += usage.inputTokens ?? 0;
-      aggregatedUsage.outputTokens += usage.outputTokens ?? 0;
+      if (usage.inputTokens != null) {
+        inputTokens += usage.inputTokens;
+        hasInput = true;
+      }
+      if (usage.outputTokens != null) {
+        outputTokens += usage.outputTokens;
+        hasOutput = true;
+      }
+      if (usage.reasoningTokens != null) {
+        reasoningTokens += usage.reasoningTokens;
+        hasReasoning = true;
+      }
+      if (usage.cachedInputTokens != null) {
+        cachedInputTokens += usage.cachedInputTokens;
+        hasCachedInput = true;
+      }
     }
   }
 
@@ -631,7 +765,12 @@ function extractAiSdkV6Batch<T>(contents: T[]): TelemetryData | null {
     allStepsData.length,
     metadata,
     allStepsData,
-    aggregatedUsage,
+    {
+      ...(hasInput ? { inputTokens } : undefined),
+      ...(hasOutput ? { outputTokens } : undefined),
+      ...(hasReasoning ? { reasoningTokens } : undefined),
+      ...(hasCachedInput ? { cachedInputTokens } : undefined),
+    },
   );
 }
 

package/src/legacy-runtime/runtime-cores/assistant-transport/useToolInvocations.test.ts

@@ -4,19 +4,22 @@ import type { ThreadAssistantMessage } from "@assistant-ui/core";
 import type { Tool } from "assistant-stream";
 import { act, renderHook, waitFor } from "@testing-library/react";
 import { describe, expect, it, vi } from "vitest";
-import type { AssistantTransportState, ToolExecutionStatus } from "./types";
-import { useToolInvocations } from "./useToolInvocations";
+import type { AssistantTransportState } from "./types";
+import { ToolExecutionStatus, useToolInvocations } from "./useToolInvocations";
+import { ReadonlyJSONObject, ReadonlyJSONValue } from "assistant-stream/utils";
 
 const createState = (
   messages: ThreadAssistantMessage[],
+  isRunning: boolean = true,
 ): AssistantTransportState => ({
   messages,
-  isRunning: true,
+  isRunning,
 });
 
 const createAssistantMessage = (
   argsText: string,
   args: Record<string, unknown>,
+  options?: { result?: ReadonlyJSONValue; isError?: boolean },
 ): ThreadAssistantMessage => ({
   id: "m-1",
   role: "assistant",
@@ -34,8 +37,10 @@ const createAssistantMessage = (
       type: "tool-call",
       toolCallId: "tool-1",
       toolName: "weatherSearch",
-      args,
+      args: args as ReadonlyJSONObject,
       argsText,
+      ...(options?.result !== undefined && { result: options.result }),
+      ...(options?.isError !== undefined && { isError: options.isError }),
     },
   ],
 });
@@ -45,6 +50,7 @@ describe("useToolInvocations", () => {
     const execute = vi.fn(async () => ({ forecast: "ok" }));
     const getTools = () => ({
       weatherSearch: {
+        parameters: { type: "object", properties: {} },
         execute,
       } satisfies Tool,
     });
@@ -137,6 +143,7 @@ describe("useToolInvocations", () => {
     );
     const getTools = () => ({
       weatherSearch: {
+        parameters: { type: "object", properties: {} },
         execute,
       } satisfies Tool,
     });
@@ -220,4 +227,223 @@ describe("useToolInvocations", () => {
     });
     expect(Object.keys(statuses)).not.toContain("tool-1:rewrite:0");
   });
+
+  it("does not close args stream early for non-executable tool snapshots", () => {
+    const getTools = () => ({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+      } satisfies Tool,
+    });
+    const onResult = vi.fn();
+    const setToolStatuses = vi.fn();
+    const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
+
+    try {
+      const { rerender } = renderHook(
+        ({ state }: { state: AssistantTransportState }) =>
+          useToolInvocations({
+            state,
+            getTools,
+            onResult,
+            setToolStatuses,
+          }),
+        {
+          initialProps: {
+            state: createState([]),
+          },
+        },
+      );
+
+      act(() => {
+        rerender({
+          state: createState([createAssistantMessage("{}", {})]),
+        });
+      });
+
+      act(() => {
+        rerender({
+          state: createState([
+            createAssistantMessage('{"title":"Weekly"', {
+              title: "Weekly",
+            }),
+          ]),
+        });
+      });
+
+      act(() => {
+        rerender({
+          state: createState([
+            createAssistantMessage('{"title":"Weekly","columns":["name"]}', {
+              title: "Weekly",
+              columns: ["name"],
+            }),
+          ]),
+        });
+      });
+
+      expect(warnSpy).not.toHaveBeenCalledWith(
+        "argsText updated after controller was closed:",
+        expect.anything(),
+      );
+      expect(warnSpy).not.toHaveBeenCalledWith(
+        "argsText updated after controller was closed, restarting tool args stream:",
+        expect.anything(),
+      );
+      expect(onResult).not.toHaveBeenCalled();
+    } finally {
+      warnSpy.mockRestore();
+    }
+  });
+
+  it("closes non-executable complete args stream after run settles", () => {
+    const getTools = () => ({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+      } satisfies Tool,
+    });
+    const onResult = vi.fn();
+    const setToolStatuses = vi.fn();
+    const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
+
+    try {
+      const { rerender } = renderHook(
+        ({ state }: { state: AssistantTransportState }) =>
+          useToolInvocations({
+            state,
+            getTools,
+            onResult,
+            setToolStatuses,
+          }),
+        {
+          initialProps: {
+            state: createState([]),
+          },
+        },
+      );
+
+      act(() => {
+        rerender({
+          state: createState(
+            [
+              createAssistantMessage('{"title":"Weekly"}', {
+                title: "Weekly",
+              }),
+            ],
+            true,
+          ),
+        });
+      });
+
+      act(() => {
+        rerender({
+          state: createState(
+            [
+              createAssistantMessage('{"title":"Weekly"}', {
+                title: "Weekly",
+              }),
+            ],
+            false,
+          ),
+        });
+      });
+
+      act(() => {
+        rerender({
+          state: createState(
+            [
+              createAssistantMessage('{"title":"Weekly","columns":["name"]}', {
+                title: "Weekly",
+                columns: ["name"],
+              }),
+            ],
+            false,
+          ),
+        });
+      });
+
+      expect(warnSpy).toHaveBeenCalledWith(
+        "argsText updated after controller was closed, restarting tool args stream:",
+        expect.objectContaining({
+          previous: '{"title":"Weekly"}',
+          next: '{"title":"Weekly","columns":["name"]}',
+        }),
+      );
+      expect(onResult).not.toHaveBeenCalled();
+    } finally {
+      warnSpy.mockRestore();
+    }
+  });
+
+  it("handles backend result when equivalent complete argsText reorders keys", async () => {
+    let resolveExecute: ((value: unknown) => void) | undefined;
+    const execute = vi.fn(
+      () =>
+        new Promise<unknown>((resolve) => {
+          resolveExecute = resolve;
+        }),
+    );
+    const getTools = () => ({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+        execute,
+      } satisfies Tool,
+    });
+    const onResult = vi.fn();
+    const setToolStatuses = vi.fn();
+
+    const { rerender } = renderHook(
+      ({ state }: { state: AssistantTransportState }) =>
+        useToolInvocations({
+          state,
+          getTools,
+          onResult,
+          setToolStatuses,
+        }),
+      {
+        initialProps: {
+          state: createState([]),
+        },
+      },
+    );
+
+    act(() => {
+      rerender({
+        state: createState([
+          createAssistantMessage('{"a":1,"b":2}', {
+            a: 1,
+            b: 2,
+          }),
+        ]),
+      });
+    });
+
+    await waitFor(() => {
+      expect(execute).toHaveBeenCalledTimes(1);
+    });
+
+    act(() => {
+      rerender({
+        state: createState([
+          createAssistantMessage(
+            '{"b":2,"a":1}',
+            {
+              a: 1,
+              b: 2,
+            },
+            {
+              result: { source: "backend" },
+            },
+          ),
+        ]),
+      });
+    });
+
+    await act(async () => {
+      resolveExecute?.({ source: "client" });
+    });
+
+    await waitFor(() => {
+      expect(onResult).not.toHaveBeenCalled();
+    });
+  });
 });