assistant-cloud 0.2.0 → 0.2.1

package/src/CloudEngagementReporter.ts

@@ -0,0 +1,231 @@
+import type { AssistantCloud } from "./AssistantCloud";
+import type {
+  AssistantCloudEvent,
+  AssistantCloudEventKind,
+} from "./AssistantCloudEvents";
+
+export type EngagementEventIds = Pick<
+  AssistantCloudEvent,
+  "thread_id" | "message_id" | "run_id"
+>;
+
+/**
+ * Turns the ids an integration knows (its own thread and message ids) into the
+ * ids the cloud stores. A send is the one event that may create the remote
+ * thread, so it asks for `awaitThread`; every other event reads the ids that
+ * already exist.
+ */
+export type EngagementIdResolver = (
+  threadId: string,
+  messageId: string | undefined,
+  options: { awaitThread: boolean },
+) => EngagementEventIds | Promise<EngagementEventIds>;
+
+type EngagementEventInit = Pick<AssistantCloudEvent, "value" | "props">;
+
+const MAX_REMEMBERED_THREADS = 256;
+
+function remember<T>(map: Map<string, T>, threadId: string, value: T): void {
+  map.delete(threadId);
+  map.set(threadId, value);
+  if (map.size > MAX_REMEMBERED_THREADS) {
+    map.delete(map.keys().next().value!);
+  }
+}
+
+function mark(set: Set<string>, threadId: string): void {
+  set.delete(threadId);
+  set.add(threadId);
+  if (set.size > MAX_REMEMBERED_THREADS) {
+    set.delete(set.values().next().value!);
+  }
+}
+
+const passThroughIds: EngagementIdResolver = (threadId, messageId) => ({
+  thread_id: threadId,
+  ...(messageId !== undefined ? { message_id: messageId } : undefined),
+});
+
+/**
+ * Derives engagement events from what a chat integration observes and keeps
+ * the per thread state the events need: a run's start for the stop duration,
+ * a run's end for the time to the next message, one error and one suggestion
+ * list per run or thread. A started run is kept until it ends or stops; the
+ * rest is kept for the 256 most recently touched threads, so a long session
+ * does not grow it without bound. Delivery goes through the cloud's event
+ * buffer, so a disabled telemetry setting drops everything here as well.
+ */
+export class CloudEngagementReporter {
+  private readonly runStartedAt = new Map<string, number>();
+  private readonly runEndedAt = new Map<string, number>();
+  private readonly shownErrors = new Set<string>();
+  private readonly shownSuggestions = new Set<string>();
+
+  private readonly getCloud: () => AssistantCloud;
+  private readonly resolveIds: EngagementIdResolver;
+
+  constructor(
+    cloud: AssistantCloud | (() => AssistantCloud),
+    resolveIds: EngagementIdResolver = passThroughIds,
+  ) {
+    this.getCloud = typeof cloud === "function" ? cloud : () => cloud;
+    this.resolveIds = resolveIds;
+  }
+
+  public runStarted(threadId: string): void {
+    this.runStartedAt.set(threadId, Date.now());
+    this.shownErrors.delete(threadId);
+  }
+
+  public runEnded(threadId: string): void {
+    this.runStartedAt.delete(threadId);
+    remember(this.runEndedAt, threadId, Date.now());
+  }
+
+  /** Reported once per started run, with the time the run had been going. */
+  public runStopped(threadId: string): void {
+    const startedAt = this.runStartedAt.get(threadId);
+    if (startedAt === undefined) return;
+    this.runStartedAt.delete(threadId);
+    this.track("run_stopped", threadId, undefined, {
+      value: Math.max(0, Date.now() - startedAt),
+    });
+  }
+
+  /** Carries the time since the previous run of the thread ended, when known. */
+  public messageSent(
+    threadId: string,
+    init: {
+      messageId?: string | undefined;
+      chars: number;
+      attachments: number;
+    },
+  ): void {
+    const previousRunEndedAt = this.runEndedAt.get(threadId);
+    this.track(
+      "message_sent",
+      threadId,
+      init.messageId,
+      {
+        props: { chars: init.chars, attachments: init.attachments },
+        ...(previousRunEndedAt !== undefined
+          ? { value: Math.max(0, Date.now() - previousRunEndedAt) }
+          : undefined),
+      },
+      { awaitThread: true },
+    );
+  }
+
+  public messageEdited(
+    threadId: string,
+    init: { messageId: string; chars: number },
+  ): void {
+    this.track("message_edited", threadId, init.messageId, {
+      props: { chars: init.chars },
+    });
+  }
+
+  public messageRegenerated(threadId: string, messageId?: string): void {
+    this.track("message_regenerated", threadId, messageId);
+  }
+
+  /** Reported once per run, so a retried render of the same error stays one event. */
+  public errorShown(
+    threadId: string,
+    init: { messageId?: string | undefined; reason: string },
+  ): void {
+    if (this.shownErrors.has(threadId)) return;
+    mark(this.shownErrors, threadId);
+    this.track("error_shown", threadId, init.messageId, {
+      props: { reason: init.reason },
+    });
+  }
+
+  /** Reported once per thread, with the number of suggestions on offer. */
+  public suggestionsShown(threadId: string, count: number): void {
+    if (this.shownSuggestions.has(threadId)) return;
+    mark(this.shownSuggestions, threadId);
+    this.track("suggestions_shown", threadId, undefined, { value: count });
+  }
+
+  public suggestionClicked(threadId: string): void {
+    this.track("suggestion_clicked", threadId);
+  }
+
+  public attachmentAdded(
+    threadId: string,
+    init: { messageId?: string | undefined; contentType?: string | undefined },
+  ): void {
+    this.track("attachment_added", threadId, init.messageId, {
+      ...(init.contentType ? { props: { type: init.contentType } } : undefined),
+    });
+  }
+
+  public attachmentFailed(
+    threadId: string,
+    init: { messageId?: string | undefined; contentType?: string | undefined },
+  ): void {
+    this.track("attachment_failed", threadId, init.messageId, {
+      ...(init.contentType ? { props: { type: init.contentType } } : undefined),
+    });
+  }
+
+  public voiceStarted(threadId: string): void {
+    this.track("voice_started", threadId);
+  }
+
+  public speechStarted(threadId: string, messageId?: string): void {
+    this.track("speech_started", threadId, messageId);
+  }
+
+  public branchSwitched(threadId: string, messageId?: string): void {
+    this.track("branch_switched", threadId, messageId);
+  }
+
+  public messageCopied(threadId: string, messageId?: string): void {
+    this.track("message_copied", threadId, messageId);
+  }
+
+  public toolApproved(
+    threadId: string,
+    messageId: string,
+    toolCallId: string,
+    toolName: string,
+  ): void {
+    this.track("tool_approved", threadId, messageId, {
+      props: { toolCallId, toolName },
+    });
+  }
+
+  public toolRejected(
+    threadId: string,
+    messageId: string,
+    toolCallId: string,
+    toolName: string,
+  ): void {
+    this.track("tool_rejected", threadId, messageId, {
+      props: { toolCallId, toolName },
+    });
+  }
+
+  /** Reported only for a thread the cloud already knows. */
+  public threadSwitched(threadId: string): void {
+    this.track("thread_switched", threadId);
+  }
+
+  private track(
+    kind: AssistantCloudEventKind,
+    threadId: string,
+    messageId?: string,
+    init: EngagementEventInit = {},
+    options: { awaitThread: boolean } = { awaitThread: false },
+  ): void {
+    void Promise.resolve()
+      .then(() => this.resolveIds(threadId, messageId, options))
+      .then((ids) => {
+        if (kind === "thread_switched" && !ids.thread_id) return;
+        this.getCloud().events.track({ kind, ...init, ...ids });
+      })
+      .catch(() => {});
+  }
+}

package/src/CloudMessagePersistence.ts

@@ -5,15 +5,9 @@ import type { CloudMessage } from "./AssistantCloudThreadMessages";
 const CLOUD_MESSAGE_PAGE_SIZE = 200;
 
 /**
- * Shared persistence logic for cloud message storage.
- *
- * Handles ID mapping (local → remote) and parent_id chaining for both:
- * - AssistantCloudThreadHistoryAdapter (assistant-ui runtime)
- * - useCloudChat (standalone AI SDK hook)
- *
- * The promise-based ID resolution handles concurrent appends — if message B's
- * parent is message A, and A is still being created, we await A's promise
- * to get its remote ID before creating B.
+ * Appends, updates and loads cloud messages while mapping local ids to cloud
+ * ids and chaining parent_id. A parent that is still being created is awaited,
+ * so concurrent appends land under the right parent.
  */
 export class CloudMessagePersistence {
   private idMapping = new Map<string, string | Promise<string>>();
@@ -126,11 +120,14 @@ export class CloudMessagePersistence {
    * The ID mapping is populated so that `isPersisted()` returns true for
    * loaded messages, preventing re-persistence of already-stored messages.
    *
+   * A loaded ID that an append already maps keeps the remote ID from that append, and falls back to the loaded ID if the append fails.
+   *
    * @param threadId - Remote thread ID
    * @param format - Optional format filter
    * @returns Array of cloud messages
    */
   async load(threadId: string, format?: string) {
+    const idMapping = this.idMapping;
     const cloud = this.getCloud();
     const messages: CloudMessage[] = [];
     const seen = new Set<string>();
@@ -156,17 +153,31 @@ export class CloudMessagePersistence {
       after = last.id;
     }
 
-    // Populate ID mapping so isPersisted() recognizes loaded messages
-    for (const m of messages) {
-      this.idMapping.set(m.id, m.id);
+    if (this.idMapping === idMapping) {
+      for (const m of messages) {
+        const entry = idMapping.get(m.id);
+        if (entry === undefined) {
+          idMapping.set(m.id, m.id);
+        } else if (entry instanceof Promise) {
+          void entry.catch(() => {
+            const current = idMapping.get(m.id);
+            if (current === undefined || current === entry) {
+              idMapping.set(m.id, m.id);
+            }
+          });
+        }
+      }
     }
     return messages;
   }
 
   /**
    * Reset the ID mapping (call when switching threads).
+   *
+   * Pending `load()` and `append()` calls are not cancelled and still settle
+   * normally, but their results no longer populate the ID mapping.
    */
   reset() {
-    this.idMapping.clear();
+    this.idMapping = new Map();
   }
 }

package/src/CloudRunReporter.ts

@@ -0,0 +1,38 @@
+import type { AssistantCloud } from "./AssistantCloud";
+import { createRunReport, type RunReportInit } from "./runTelemetry";
+
+export type CloudRunReportInit = Omit<RunReportInit, "telemetry">;
+
+/**
+ * Sends run reports the way every client integration has to: nothing while
+ * telemetry is off, the cloud's environment, release and tags on every report,
+ * the `beforeReport` hook applied last, and a failed send that never surfaces.
+ * A report given a key is sent once per key, so an integration that observes
+ * the same finished run twice reports it once.
+ */
+export class CloudRunReporter {
+  private readonly reported = new Set<string>();
+  private readonly getCloud: () => AssistantCloud;
+
+  constructor(cloud: AssistantCloud | (() => AssistantCloud)) {
+    this.getCloud = typeof cloud === "function" ? cloud : () => cloud;
+  }
+
+  public async report(init: CloudRunReportInit, key?: string): Promise<void> {
+    try {
+      const cloud = this.getCloud();
+      if (!cloud.telemetry.enabled) return;
+      if (key !== undefined && this.reported.has(key)) return;
+
+      const initial = createRunReport({ ...init, telemetry: cloud.telemetry });
+      const { beforeReport } = cloud.telemetry;
+      const report = beforeReport ? beforeReport(initial) : initial;
+      if (!report) return;
+
+      if (key !== undefined) this.reported.add(key);
+      await cloud.runs.report(report);
+    } catch {
+      return;
+    }
+  }
+}

package/src/FormattedCloudPersistence.ts

@@ -18,14 +18,9 @@ export type MessageFormatAdapter<TMessage, TStorageFormat> = {
 };
 
 /**
- * Wraps a CloudMessagePersistence instance with format-aware encode/decode.
- *
- * This centralizes the pattern used by both:
- * - useCloudChat (standalone AI SDK hook)
- * - AssistantCloudThreadHistoryAdapter.withFormat() (assistant-ui runtime)
- *
- * The persistence parameter is typed structurally (not by class) so callers
- * don't need to import CloudMessagePersistence directly.
+ * Wraps a CloudMessagePersistence with a MessageFormatAdapter's encode and
+ * decode. The persistence parameter is typed structurally, so a caller does
+ * not need to import the class.
  */
 export const createFormattedPersistence = <TMessage, TStorageFormat>(
   persistence: {

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();
+  });
+});