assistant-cloud 0.2.0 → 0.2.2

package/src/ai-sdk/index.test.ts

@@ -0,0 +1,258 @@
+import type { UIMessage } from "ai";
+import { describe, expect, it } from "vitest";
+import {
+  aiSDKV6FormatAdapter,
+  extractAISDKRunTelemetry,
+  type AISDKMessageLike,
+} from "./index";
+
+const assistant = (
+  parts: UIMessage["parts"],
+  metadata?: Record<string, unknown>,
+  id = "a1",
+): UIMessage =>
+  ({
+    id,
+    role: "assistant",
+    parts,
+    ...(metadata ? { metadata } : {}),
+  }) as UIMessage;
+
+describe("aiSDKV6FormatAdapter", () => {
+  it("stores a message without its id and restores it from the row", () => {
+    const message = assistant([{ type: "text", text: "hi" }], { a: 1 });
+    const content = aiSDKV6FormatAdapter.encode({ parentId: null, message });
+    expect(content).toEqual({
+      role: "assistant",
+      parts: [{ type: "text", text: "hi" }],
+      metadata: { a: 1 },
+    });
+    expect(
+      aiSDKV6FormatAdapter.decode({
+        id: "a1",
+        parent_id: "u1",
+        format: "ai-sdk/v6",
+        content,
+      }),
+    ).toEqual({ parentId: "u1", message });
+    expect(aiSDKV6FormatAdapter.getId(message)).toBe("a1");
+  });
+});
+
+describe("extractAISDKRunTelemetry", () => {
+  it("returns null without an assistant message", () => {
+    expect(
+      extractAISDKRunTelemetry([
+        { role: "user", parts: [{ type: "text", text: "q" }] },
+      ]),
+    ).toBeNull();
+  });
+
+  it("reads tool calls, steps and usage from one streamed message", () => {
+    const message = assistant(
+      [
+        { type: "step-start" },
+        {
+          type: "tool-weather",
+          toolCallId: "call_1",
+          state: "output-available",
+          input: { city: "Oslo" },
+          output: { temp: 3 },
+        } as UIMessage["parts"][number],
+        { type: "step-start" },
+        {
+          type: "dynamic-tool",
+          toolName: "search",
+          toolCallId: "call_2",
+          state: "output-available",
+          input: { q: "x" },
+          output: [{ type: "text", text: "found" }],
+        } as UIMessage["parts"][number],
+        { type: "step-start" },
+        { type: "text", text: "Cold." },
+      ],
+      {
+        usage: { inputTokens: 10, outputTokens: 5 },
+        steps: [
+          { usage: { promptTokens: 4, completionTokens: 1 } },
+          { usage: { inputTokens: 6, outputTokens: 4 } },
+          {},
+        ],
+        samplingCalls: { call_2: [{ duration_ms: 20 }] },
+        modelId: "gpt-5.6-luna",
+      },
+    );
+
+    expect(extractAISDKRunTelemetry([message])).toEqual({
+      assistantMessageId: "a1",
+      status: "completed",
+      toolCalls: [
+        {
+          tool_name: "weather",
+          tool_call_id: "call_1",
+          tool_args: '{"city":"Oslo"}',
+          tool_result: '{"temp":3}',
+          tool_source: "frontend",
+        },
+        {
+          tool_name: "search",
+          tool_call_id: "call_2",
+          tool_args: '{"q":"x"}',
+          tool_result: '[{"type":"text","text":"found"}]',
+          tool_source: "mcp",
+          sampling_calls: [{ duration_ms: 20 }],
+        },
+      ],
+      steps: [
+        {
+          usage: { promptTokens: 4, completionTokens: 1 },
+          toolCalls: [expect.objectContaining({ tool_call_id: "call_1" })],
+          finishReason: "tool-calls",
+        },
+        {
+          usage: { inputTokens: 6, outputTokens: 4 },
+          toolCalls: [expect.objectContaining({ tool_call_id: "call_2" })],
+          finishReason: "tool-calls",
+        },
+        {},
+      ],
+      totalSteps: 3,
+      outputText: "Cold.",
+      usage: { inputTokens: 10, outputTokens: 5 },
+      modelId: "gpt-5.6-luna",
+      metadata: message.metadata,
+    });
+  });
+
+  it("sums usage over stored rows without ids and reads legacy tool fields", () => {
+    const rows: AISDKMessageLike[] = [
+      {
+        role: "assistant",
+        parts: [
+          { type: "step-start" },
+          {
+            type: "tool-lookup",
+            toolCallId: "call_1",
+            state: "output-available",
+            args: { id: 1 },
+            result: "ok",
+          } as unknown as UIMessage["parts"][number],
+        ],
+        metadata: { steps: [{ usage: { inputTokens: 2, outputTokens: 1 } }] },
+      },
+      {
+        role: "assistant",
+        parts: [{ type: "step-start" }, { type: "text", text: "done" }],
+        metadata: {
+          steps: [{ usage: { inputTokens: 3, reasoningTokens: 7 } }],
+        },
+      },
+    ];
+    const result = extractAISDKRunTelemetry(rows);
+    expect(result?.assistantMessageId).toBeUndefined();
+    expect(result?.status).toBe("completed");
+    expect(result?.totalSteps).toBe(2);
+    expect(result?.usage).toEqual({
+      inputTokens: 5,
+      outputTokens: 1,
+      reasoningTokens: 7,
+    });
+    expect(result?.toolCalls).toEqual([
+      {
+        tool_name: "lookup",
+        tool_call_id: "call_1",
+        tool_args: '{"id":1}',
+        tool_result: '"ok"',
+        tool_source: "frontend",
+      },
+    ]);
+  });
+
+  it("attaches sampling calls from the row that made the tool call", () => {
+    const rows: AISDKMessageLike[] = [
+      {
+        role: "assistant",
+        parts: [
+          {
+            type: "dynamic-tool",
+            toolName: "search",
+            toolCallId: "call_1",
+            state: "output-available",
+            input: {},
+            output: "x",
+          } as UIMessage["parts"][number],
+        ],
+        metadata: { samplingCalls: { call_1: [{ duration_ms: 5 }] } },
+      },
+      {
+        role: "assistant",
+        parts: [{ type: "text", text: "done" }],
+        metadata: { samplingCalls: {} },
+      },
+    ];
+    expect(extractAISDKRunTelemetry(rows)?.toolCalls).toEqual([
+      expect.objectContaining({
+        tool_call_id: "call_1",
+        tool_source: "mcp",
+        sampling_calls: [{ duration_ms: 5 }],
+      }),
+    ]);
+  });
+
+  it("reads a tool only run as completed with its step closed by the calls", () => {
+    const call = (toolCallId: string, state: string, output?: unknown) =>
+      ({
+        type: "tool-search",
+        toolCallId,
+        state,
+        input: {},
+        ...(output !== undefined ? { output } : {}),
+      }) as UIMessage["parts"][number];
+
+    const answered = extractAISDKRunTelemetry([
+      assistant([{ type: "step-start" }, call("a", "output-available", [])]),
+    ]);
+    expect(answered?.status).toBe("completed");
+    expect(answered?.steps?.[0]?.finishReason).toBe("tool-calls");
+    expect(answered?.outputText).toBeUndefined();
+
+    const pending = extractAISDKRunTelemetry([
+      assistant([{ type: "step-start" }, call("b", "input-available")]),
+    ]);
+    expect(pending?.status).toBe("completed");
+    expect(pending?.steps?.[0]?.finishReason).toBe("tool-calls");
+  });
+
+  it("reads incomplete for a run that produced nothing", () => {
+    const empty = extractAISDKRunTelemetry([
+      assistant([{ type: "step-start" }]),
+    ]);
+    expect(empty?.status).toBe("incomplete");
+    expect(empty?.totalSteps).toBe(1);
+    expect(empty?.steps?.[0]?.finishReason).toBeUndefined();
+    expect(empty?.outputText).toBeUndefined();
+  });
+
+  it("reads no steps without step markers", () => {
+    const result = extractAISDKRunTelemetry([
+      assistant([{ type: "text", text: "hi" }]),
+    ]);
+    expect(result?.steps).toBeUndefined();
+    expect(result?.totalSteps).toBeUndefined();
+  });
+
+  it("skips a stored part without a type", () => {
+    const result = extractAISDKRunTelemetry([
+      {
+        role: "assistant",
+        parts: [
+          null,
+          { toolCallId: "a" },
+          { type: "text", text: "hi" },
+        ] as never,
+      },
+    ]);
+    expect(result?.outputText).toBe("hi");
+    expect(result?.toolCalls).toBeUndefined();
+  });
+});

package/src/ai-sdk/index.ts

@@ -0,0 +1,222 @@
+import type { UIMessage } from "ai";
+import type { MessageFormatAdapter } from "../FormattedCloudPersistence";
+import type { SamplingCallData } from "../instrumentMcpSampling";
+import {
+  type AssistantCloudRunReportToolCall,
+  createRunTelemetryToolCall,
+  extractRunTelemetryModelId,
+  normalizeRunTelemetryUsage,
+  type RunMessageTelemetry,
+  type RunReportStepInit,
+  type RunTelemetryUsage,
+  type RunTelemetryUsageInit,
+  truncateRunTelemetryText,
+} from "../runTelemetry";
+
+export type AISDKStorageFormat = Omit<UIMessage, "id">;
+
+/** The stored form of an AI SDK message: the message without its id. */
+export const aiSDKV6FormatAdapter: MessageFormatAdapter<
+  UIMessage,
+  AISDKStorageFormat
+> = {
+  format: "ai-sdk/v6",
+  encode: ({ message: { id: _id, ...message } }) => message,
+  decode: (stored) => ({
+    parentId: stored.parent_id,
+    message: { id: stored.id, ...stored.content } as UIMessage,
+  }),
+  getId: (message) => message.id,
+};
+
+/**
+ * A message as an AI SDK integration holds it, or as the cloud stored it under
+ * the ai-sdk/v6 format, which drops the id. Parts are typed by their `type`
+ * alone, so this shape and the telemetry read from it need nothing from `ai`.
+ */
+export type AISDKMessageLike = {
+  id?: string | undefined;
+  role: string;
+  parts: readonly { type: string; [key: string]: unknown }[];
+  metadata?: unknown;
+};
+
+type Part = Record<string, unknown> & { type: string };
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object";
+}
+
+function isPart(value: unknown): value is Part {
+  return isRecord(value) && typeof value.type === "string";
+}
+
+/**
+ * The AI SDK's own tool part rules, kept here so the entry loads without the
+ * `ai` runtime: a static tool part is `tool-<name>`, a dynamic one is
+ * `dynamic-tool` with its name in `toolName`.
+ */
+function toolCallOf(part: Part): AssistantCloudRunReportToolCall | undefined {
+  if (typeof part.toolCallId !== "string") return undefined;
+  const isStatic = part.type.startsWith("tool-");
+  if (!isStatic && part.type !== "dynamic-tool") return undefined;
+  const toolName = isStatic
+    ? part.type.slice("tool-".length)
+    : typeof part.toolName === "string"
+      ? part.toolName
+      : undefined;
+  if (!toolName) return undefined;
+  return createRunTelemetryToolCall({
+    toolName,
+    toolCallId: part.toolCallId,
+    args: part.input ?? part.args,
+    result: part.output ?? part.result,
+    toolSource: isStatic ? "frontend" : "mcp",
+  });
+}
+
+function attachSamplingCalls(
+  toolCalls: readonly AssistantCloudRunReportToolCall[],
+  metadata: Record<string, unknown> | undefined,
+): void {
+  const samplingCalls = isRecord(metadata?.samplingCalls)
+    ? (metadata.samplingCalls as Record<string, SamplingCallData[]>)
+    : undefined;
+  if (!samplingCalls) return;
+  for (const toolCall of toolCalls) {
+    const calls = samplingCalls[toolCall.tool_call_id];
+    if (Array.isArray(calls) && calls.length > 0) {
+      toolCall.sampling_calls = calls;
+    }
+  }
+}
+
+function stepUsages(
+  metadata: Record<string, unknown> | undefined,
+): (RunTelemetryUsageInit | undefined)[] {
+  const steps = metadata?.steps;
+  if (!Array.isArray(steps)) return [];
+  return steps.map((step) =>
+    isRecord(step) && isRecord(step.usage)
+      ? (step.usage as RunTelemetryUsageInit)
+      : undefined,
+  );
+}
+
+function sumUsage(usages: readonly RunTelemetryUsage[]): RunTelemetryUsage {
+  const total: RunTelemetryUsage = {};
+  for (const usage of usages) {
+    if (usage.inputTokens != null) {
+      total.inputTokens = (total.inputTokens ?? 0) + usage.inputTokens;
+    }
+    if (usage.outputTokens != null) {
+      total.outputTokens = (total.outputTokens ?? 0) + usage.outputTokens;
+    }
+    if (usage.reasoningTokens != null) {
+      total.reasoningTokens =
+        (total.reasoningTokens ?? 0) + usage.reasoningTokens;
+    }
+    if (usage.cachedInputTokens != null) {
+      total.cachedInputTokens =
+        (total.cachedInputTokens ?? 0) + usage.cachedInputTokens;
+    }
+  }
+  return total;
+}
+
+/**
+ * The usage a message reports: `metadata.usage` when the integration copied
+ * the run total there, else the sum over `metadata.steps[].usage`.
+ */
+function messageUsage(
+  metadata: Record<string, unknown> | undefined,
+): RunTelemetryUsage | undefined {
+  const total = isRecord(metadata?.usage)
+    ? normalizeRunTelemetryUsage(metadata.usage as RunTelemetryUsageInit)
+    : undefined;
+  if (total) return total;
+  const perStep = stepUsages(metadata).flatMap((usage) => {
+    const normalized = usage ? normalizeRunTelemetryUsage(usage) : undefined;
+    return normalized ? [normalized] : [];
+  });
+  return perStep.length > 0 ? sumUsage(perStep) : undefined;
+}
+
+/**
+ * Reads the run report fields out of the assistant messages of one run. A run
+ * the AI SDK streamed as one message is one element; a run the cloud stored as
+ * several assistant rows is aggregated, step by step, in order. Returns null
+ * when no assistant message is present. Status reads completed when the run
+ * produced text or tool calls, as a live finish with reason `tool-calls` does;
+ * an integration that observed the finish event overrides it.
+ */
+export function extractAISDKRunTelemetry(
+  messages: readonly AISDKMessageLike[],
+): RunMessageTelemetry | null {
+  const textParts: string[] = [];
+  const toolCalls: AssistantCloudRunReportToolCall[] = [];
+  const steps: RunReportStepInit[] = [];
+  const usages: RunTelemetryUsage[] = [];
+  let assistant: AISDKMessageLike | undefined;
+
+  for (const message of messages) {
+    if (message.role !== "assistant") continue;
+    assistant = message;
+    const metadata = isRecord(message.metadata) ? message.metadata : undefined;
+    const usagePerStep = stepUsages(metadata);
+    const messageToolCalls: AssistantCloudRunReportToolCall[] = [];
+    let step: RunReportStepInit | undefined;
+    let stepIndex = -1;
+
+    for (const part of message.parts) {
+      if (!isPart(part)) continue;
+      if (part.type === "step-start") {
+        stepIndex += 1;
+        const usage = usagePerStep[stepIndex];
+        step = usage ? { usage } : {};
+        steps.push(step);
+        continue;
+      }
+      if (part.type === "text" && typeof part.text === "string" && part.text) {
+        textParts.push(part.text);
+        continue;
+      }
+      const toolCall = toolCallOf(part);
+      if (!toolCall) continue;
+      toolCalls.push(toolCall);
+      messageToolCalls.push(toolCall);
+      if (step) {
+        step.toolCalls = [...(step.toolCalls ?? []), toolCall];
+        step.finishReason = "tool-calls";
+      }
+    }
+
+    attachSamplingCalls(messageToolCalls, metadata);
+    const usage = messageUsage(metadata);
+    if (usage) usages.push(usage);
+  }
+
+  if (!assistant) return null;
+
+  const metadata = isRecord(assistant.metadata)
+    ? assistant.metadata
+    : undefined;
+
+  const usage = usages.length > 0 ? sumUsage(usages) : undefined;
+  const modelId = extractRunTelemetryModelId(metadata);
+  const completed = textParts.length > 0 || toolCalls.length > 0;
+  return {
+    ...(assistant.id !== undefined
+      ? { assistantMessageId: assistant.id }
+      : undefined),
+    status: completed ? "completed" : "incomplete",
+    ...(toolCalls.length > 0 ? { toolCalls } : undefined),
+    ...(steps.length > 0 ? { steps, totalSteps: steps.length } : undefined),
+    ...(textParts.length > 0
+      ? { outputText: truncateRunTelemetryText(textParts.join("")) }
+      : undefined),
+    ...(usage ? { usage } : undefined),
+    ...(modelId ? { modelId } : undefined),
+    ...(metadata ? { metadata } : undefined),
+  };
+}

package/src/generateThreadTitle.test.ts

@@ -68,4 +68,36 @@ describe("generateThreadTitle", () => {
 
     expect(update).not.toHaveBeenCalled();
   });
+
+  it("returns null without updating for a whitespace-only title", async () => {
+    const { cloud, update } = createCloud(
+      titleStream({ type: "text-delta", textDelta: " \n\t" }),
+    );
+
+    await expect(
+      generateThreadTitle(cloud, {
+        threadId: "thread-1",
+        messages: [],
+      }),
+    ).resolves.toBeNull();
+
+    expect(update).not.toHaveBeenCalled();
+  });
+
+  it("trims surrounding whitespace from a generated title", async () => {
+    const { cloud, update } = createCloud(
+      titleStream({ type: "text-delta", textDelta: "  Weather chat \n" }),
+    );
+
+    await expect(
+      generateThreadTitle(cloud, {
+        threadId: "thread-1",
+        messages: [],
+      }),
+    ).resolves.toBe("Weather chat");
+
+    expect(update).toHaveBeenCalledExactlyOnceWith("thread-1", {
+      title: "Weather chat",
+    });
+  });
 });

package/src/generateThreadTitle.ts

@@ -30,6 +30,7 @@ export async function generateThreadTitle(
     reader.releaseLock();
   }
 
+  title = title.trim();
   if (title) {
     await cloud.threads.update(options.threadId, { title });
   }

package/src/index.ts

@@ -3,7 +3,10 @@ export type {
   AssistantCloudThreadMessageFeedbackBody,
   AssistantCloudThreadMessageFeedbackResponse,
 } from "./AssistantCloudThreadMessages";
-export type { AssistantCloudTelemetryConfig } from "./AssistantCloudAPI";
+export type {
+  AssistantCloudTelemetryConfig,
+  SdkIdentity,
+} from "./AssistantCloudAPI";
 export {
   AssistantCloudEvents,
   type AssistantCloudEvent,
@@ -28,6 +31,7 @@ export {
   normalizeRunTelemetryUsage,
   truncateRunTelemetryText,
   type AssistantCloudRunReportToolCall,
+  type RunMessageTelemetry,
   type RunReportInit,
   type RunReportOutcome,
   type RunReportStepInit,
@@ -36,6 +40,12 @@ export {
   type RunTelemetryUsageInit,
 } from "./runTelemetry";
 export { AssistantCloud } from "./AssistantCloud";
+export { CloudRunReporter, type CloudRunReportInit } from "./CloudRunReporter";
+export {
+  CloudEngagementReporter,
+  type EngagementEventIds,
+  type EngagementIdResolver,
+} from "./CloudEngagementReporter";
 export { readAnonymousRefreshToken } from "./AssistantCloudAuthStrategy";
 export { CloudMessagePersistence } from "./CloudMessagePersistence";
 export {

package/src/instrumentMcpSampling.test.ts

@@ -0,0 +1,102 @@
+import { describe, expect, it, vi } from "vitest";
+import {
+  wrapSamplingHandler,
+  createSamplingCollector,
+  type McpSamplingRequest,
+} from "./instrumentMcpSampling";
+
+const request: McpSamplingRequest = {
+  method: "sampling/createMessage",
+  params: { messages: [], maxTokens: 10 },
+};
+
+describe("wrapSamplingHandler", () => {
+  it("returns the response and records sampling metrics", async () => {
+    const collector = createSamplingCollector();
+    const response = {
+      model: "test-model",
+      content: { type: "text", text: "hello" },
+      usage: { inputTokens: 2, outputTokens: 3 },
+    };
+    const handler = vi.fn(async () => response);
+    await expect(
+      wrapSamplingHandler(handler, collector.collect)(request),
+    ).resolves.toBe(response);
+    expect(handler).toHaveBeenCalledWith(request);
+    expect(collector.getCalls()).toEqual([
+      expect.objectContaining({
+        model_id: "test-model",
+        input_tokens: 2,
+        output_tokens: 3,
+        duration_ms: expect.any(Number),
+      }),
+    ]);
+  });
+
+  it.each([false, true])(
+    "preserves successful responses when the observer fails asynchronously=%s",
+    async (asynchronous) => {
+      const error = new Error("metrics unavailable");
+      const observer = asynchronous
+        ? async () => {
+            throw error;
+          }
+        : () => {
+            throw error;
+          };
+      const response = { content: "successful response" };
+      const log = vi.spyOn(console, "error").mockImplementation(() => {});
+      try {
+        await expect(
+          wrapSamplingHandler(async () => response, observer)(request),
+        ).resolves.toBe(response);
+        expect(log).toHaveBeenCalledWith(
+          "[assistant-cloud] onSamplingCall callback threw an error",
+          error,
+        );
+      } finally {
+        log.mockRestore();
+      }
+    },
+  );
+
+  it("does not await an asynchronous observer", async () => {
+    const response = { content: "successful response" };
+    await expect(
+      wrapSamplingHandler(
+        async () => response,
+        () => new Promise<void>(() => {}),
+      )(request),
+    ).resolves.toBe(response);
+  });
+
+  it("propagates model errors without calling the observer", async () => {
+    const error = new Error("model failed");
+    const observer = vi.fn();
+    await expect(
+      wrapSamplingHandler(async () => {
+        throw error;
+      }, observer)(request),
+    ).rejects.toBe(error);
+    expect(observer).not.toHaveBeenCalled();
+  });
+
+  it("preserves successful responses even if error reporting throws", async () => {
+    const response = { content: "successful response" };
+    const log = vi.spyOn(console, "error").mockImplementation(() => {
+      throw new Error("logging failed");
+    });
+    try {
+      await expect(
+        wrapSamplingHandler(
+          async () => response,
+          () => {
+            throw new Error("metrics failed");
+          },
+        )(request),
+      ).resolves.toBe(response);
+    } finally {
+      log.mockRestore();
+    }
+  });
+});

package/src/instrumentMcpSampling.ts

@@ -43,6 +43,15 @@ export type McpSamplingResponse = {
   [key: string]: unknown;
 };
 
+const reportSamplingError = (error: unknown): void => {
+  try {
+    console.error(
+      "[assistant-cloud] onSamplingCall callback threw an error",
+      error,
+    );
+  } catch {}
+};
+
 /**
  * Wraps an MCP sampling handler to intercept and measure sampling calls.
  *
@@ -80,7 +89,7 @@ export function wrapSamplingHandler(
     const reasoningTokens = response.usage?.reasoningTokens;
     const cachedInputTokens = response.usage?.cachedInputTokens;
 
-    onSamplingCall({
+    const data: SamplingCallData = {
       ...(modelId ? { model_id: modelId } : undefined),
       ...(inputTokens != null ? { input_tokens: inputTokens } : undefined),
       ...(outputTokens != null ? { output_tokens: outputTokens } : undefined),
@@ -91,7 +100,12 @@ export function wrapSamplingHandler(
         ? { cached_input_tokens: cachedInputTokens }
         : undefined),
       duration_ms: durationMs,
-    });
+    };
+    try {
+      void Promise.resolve(onSamplingCall(data)).catch(reportSamplingError);
+    } catch (error) {
+      reportSamplingError(error);
+    }
 
     return response;
   };