assistant-cloud 0.2.0 → 0.2.1

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;
   };

package/src/runTelemetry.ts

@@ -225,6 +225,23 @@ export type RunReportStepInit = {
   finishReason?: string | undefined;
 };
 
+/**
+ * The run report fields read from the messages of one run, in whichever
+ * format they were stored: the status the messages imply, the tool calls, the
+ * steps, the text, the usage and the model.
+ */
+export type RunMessageTelemetry = {
+  assistantMessageId?: string;
+  status: "completed" | "incomplete";
+  toolCalls?: AssistantCloudRunReportToolCall[];
+  steps?: RunReportStepInit[];
+  totalSteps?: number;
+  outputText?: string;
+  usage?: RunTelemetryUsage;
+  modelId?: string;
+  metadata?: Record<string, unknown>;
+};
+
 export type RunReportInit = {
   threadId: string;
   status: AssistantCloudRunReport["status"];

package/src/tests/AssistantCloud.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import { AssistantCloud } from "../AssistantCloud";
 import type { AssistantCloudTelemetryConfig } from "../AssistantCloudAPI";
 
@@ -13,6 +13,10 @@ const createCloud = (
   });
 
 describe("AssistantCloud telemetry config", () => {
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
   it("defaults to enabled", () => {
     expect(createCloud().telemetry.enabled).toBe(true);
     expect(createCloud(true).telemetry.enabled).toBe(true);
@@ -58,4 +62,32 @@ describe("AssistantCloud telemetry config", () => {
       tags: ["region:sg", "tier:paid"],
     });
   });
+
+  it("forwards registered SDK identities to requests and stream options", async () => {
+    const fetchMock = vi.fn().mockResolvedValue({
+      ok: true,
+      headers: new Headers(),
+      text: vi.fn().mockResolvedValue(JSON.stringify({ threads: [] })),
+    });
+    vi.stubGlobal("fetch", fetchMock);
+
+    const cloud = createCloud();
+    cloud.registerSdk({ name: "@assistant-ui/core", version: "0.3.18" });
+
+    await cloud.threads.list();
+
+    const [, init] = fetchMock.mock.calls[0]!;
+    expect(init.headers).toMatchObject({
+      "Aui-Sdk": expect.stringMatching(
+        /^assistant-cloud\/.* @assistant-ui\/core\/0\.3\.18$/,
+      ),
+    });
+    await expect(
+      cloud.runs.__internal_getAssistantOptions("assistant-id").headers(),
+    ).resolves.toMatchObject({
+      "Aui-Sdk": expect.stringMatching(
+        /^assistant-cloud\/.* @assistant-ui\/core\/0\.3\.18$/,
+      ),
+    });
+  });
 });

package/src/tests/AssistantCloudAPI.test.ts

@@ -45,6 +45,7 @@ describe("AssistantCloudAPI", () => {
       Authorization: "Bearer test-key",
       "Aui-User-Id": "u-1",
       "Aui-Workspace-Id": "w-1",
+      "Aui-Sdk": expect.stringMatching(/^assistant-cloud\//),
       "Content-Type": "application/json",
       "X-Test": "1",
     });
@@ -53,6 +54,56 @@ describe("AssistantCloudAPI", () => {
     expect(init.body).toBe(JSON.stringify({ hello: "world" }));
   });
 
+  it("ignores identities that cannot travel in a header", () => {
+    const api = new AssistantCloudAPI({
+      apiKey: "test-key",
+      userId: "u-1",
+      workspaceId: "w-1",
+    });
+    api.registerSdk({ name: "bad name", version: "1.0.0" });
+    api.registerSdk({ name: "@scope/pkg", version: "1.0.0 ok" });
+    api.registerSdk({ name: "@scope/pkg\ttab", version: "1.0.0" });
+    api.registerSdk({ name: "@scope/ok", version: " 1.0.0 " });
+
+    expect(api.sdkHeader().split(" ")).toEqual([
+      expect.stringMatching(/^assistant-cloud\//),
+      "@scope/ok/1.0.0",
+    ]);
+  });
+
+  it("sends each registered SDK identity once in registration order", async () => {
+    const fetchMock = vi.fn().mockResolvedValue({
+      ok: true,
+      headers: new Headers(),
+      json: vi.fn().mockResolvedValue({}),
+    });
+    vi.stubGlobal("fetch", fetchMock);
+
+    const api = new AssistantCloudAPI({
+      apiKey: "test-key",
+      userId: "u-1",
+      workspaceId: "w-1",
+    });
+
+    api.registerSdk({ name: " @assistant-ui/core ", version: " 0.3.18 " });
+    api.registerSdk({ name: "@assistant-ui/core", version: "0.3.18" });
+    api.registerSdk({ name: "@assistant-ui/ai-sdk", version: "0.0.5" });
+    api.registerSdk({ name: " ", version: "0.0.5" });
+    api.registerSdk({ name: "@assistant-ui/react-langgraph", version: " " });
+
+    await api.makeRawRequest("/threads", {
+      headers: { "Aui-Sdk": "overridden" },
+    });
+
+    const [, init] = fetchMock.mock.calls[0]!;
+    expect(init.headers).toMatchObject({ "Aui-Sdk": api.sdkHeader() });
+    expect(api.sdkHeader().split(" ")).toEqual([
+      expect.stringMatching(/^assistant-cloud\//),
+      "@assistant-ui/core/0.3.18",
+      "@assistant-ui/ai-sdk/0.0.5",
+    ]);
+  });
+
   it("uses custom baseUrl when provided with apiKey config", async () => {
     const fetchMock = vi.fn().mockResolvedValue({
       ok: true,

package/src/tests/AssistantCloudAuthStrategy.test.ts

@@ -950,4 +950,26 @@ describe("AssistantCloudJWTAuthStrategy", () => {
     });
     expect(authToken).toHaveBeenCalledTimes(2);
   });
+
+  it("does not cache a malformed rotated token", async () => {
+    const authToken = vi
+      .fn<() => Promise<string | null>>()
+      .mockResolvedValue(accessToken);
+    const strategy = new AssistantCloudJWTAuthStrategy(authToken);
+
+    await expect(strategy.getAuthHeaders()).resolves.toEqual({
+      Authorization: `Bearer ${accessToken}`,
+    });
+
+    expect(() =>
+      strategy.readAuthHeaders(
+        new Headers({ Authorization: "Bearer malformed" }),
+      ),
+    ).toThrow("Unable to determine the token expiry");
+
+    await expect(strategy.getAuthHeaders()).resolves.toEqual({
+      Authorization: `Bearer ${accessToken}`,
+    });
+    expect(authToken).toHaveBeenCalledTimes(1);
+  });
 });

package/src/tests/AssistantCloudProjects.test.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
+import { describe, it, expect, beforeEach, vi } from "vitest";
 import { AssistantCloudProjects } from "../AssistantCloudProjects";
 import type { AssistantCloudAPI } from "../AssistantCloudAPI";
 
@@ -19,10 +19,6 @@ describe("AssistantCloudProjects", () => {
     projects = new AssistantCloudProjects(mockApi);
   });
 
-  afterEach(() => {
-    vi.clearAllMocks();
-  });
-
   it("lists project threads with the query", async () => {
     vi.mocked(mockApi.makeRequest).mockResolvedValue({ threads: [] });