assistant-cloud 0.1.16 → 0.1.18

package/dist/index.d.ts

@@ -1,3 +1,8 @@
 export type { CloudMessage } from "./AssistantCloudThreadMessages.js";
+export type { AssistantCloudTelemetryConfig } from "./AssistantCloudAPI.js";
+export type { AssistantCloudRunReport } from "./AssistantCloudRuns.js";
 export { AssistantCloud } from "./AssistantCloud.js";
+export { CloudMessagePersistence } from "./CloudMessagePersistence.js";
+export { createFormattedPersistence, type MessageFormatAdapter, } from "./FormattedCloudPersistence.js";
+export { wrapSamplingHandler, createSamplingCollector, type SamplingCallData, type McpSamplingHandler, } from "./instrumentMcpSampling.js";
 //# sourceMappingURL=index.d.ts.map

package/dist/index.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,YAAY,EAAE,YAAY,EAAE,0CAAuC;AACnE,OAAO,EAAE,cAAc,EAAE,4BAAyB"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,YAAY,EAAE,YAAY,EAAE,0CAAuC;AACnE,YAAY,EAAE,6BAA6B,EAAE,+BAA4B;AACzE,YAAY,EAAE,uBAAuB,EAAE,gCAA6B;AACpE,OAAO,EAAE,cAAc,EAAE,4BAAyB;AAClD,OAAO,EAAE,uBAAuB,EAAE,qCAAkC;AACpE,OAAO,EACL,0BAA0B,EAC1B,KAAK,oBAAoB,GAC1B,uCAAoC;AACrC,OAAO,EACL,mBAAmB,EACnB,uBAAuB,EACvB,KAAK,gBAAgB,EACrB,KAAK,kBAAkB,GACxB,mCAAgC"}

package/dist/index.js

@@ -1,2 +1,5 @@
 export { AssistantCloud } from "./AssistantCloud.js";
+export { CloudMessagePersistence } from "./CloudMessagePersistence.js";
+export { createFormattedPersistence, } from "./FormattedCloudPersistence.js";
+export { wrapSamplingHandler, createSamplingCollector, } from "./instrumentMcpSampling.js";
 //# sourceMappingURL=index.js.map

package/dist/index.js.map

@@ -1 +1 @@
-{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AACA,OAAO,EAAE,cAAc,EAAE,4BAAyB"}
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAGA,OAAO,EAAE,cAAc,EAAE,4BAAyB;AAClD,OAAO,EAAE,uBAAuB,EAAE,qCAAkC;AACpE,OAAO,EACL,0BAA0B,GAE3B,uCAAoC;AACrC,OAAO,EACL,mBAAmB,EACnB,uBAAuB,GAGxB,mCAAgC"}

package/dist/instrumentMcpSampling.d.ts

@@ -0,0 +1,75 @@
+/**
+ * MCP sampling instrumentation utility.
+ *
+ * Wraps an MCP client's sampling handler to capture nested LLM calls
+ * (sampling/createMessage requests) made during tool execution.
+ * The captured data can be reported as child generation spans.
+ */
+export type SamplingCallData = {
+    model_id?: string;
+    input_tokens?: number;
+    output_tokens?: number;
+    duration_ms?: number;
+};
+export type McpSamplingHandler = (request: McpSamplingRequest) => Promise<McpSamplingResponse>;
+export type McpSamplingRequest = {
+    method: "sampling/createMessage";
+    params: {
+        messages: unknown[];
+        modelPreferences?: {
+            hints?: {
+                name?: string;
+            }[];
+        };
+        maxTokens?: number;
+        [key: string]: unknown;
+    };
+};
+export type McpSamplingResponse = {
+    model?: string;
+    content: unknown;
+    usage?: {
+        inputTokens?: number;
+        outputTokens?: number;
+        promptTokens?: number;
+        completionTokens?: number;
+    };
+    [key: string]: unknown;
+};
+/**
+ * Wraps an MCP sampling handler to intercept and measure sampling calls.
+ *
+ * @param handler - The original sampling handler from the MCP client
+ * @param onSamplingCall - Callback invoked with metrics for each sampling call
+ * @returns A wrapped handler that transparently captures sampling metrics
+ *
+ * @example
+ * ```ts
+ * const samplingCalls: SamplingCallData[] = [];
+ * const wrapped = wrapSamplingHandler(
+ *   originalHandler,
+ *   (data) => samplingCalls.push(data),
+ * );
+ * // Use `wrapped` as the MCP client's sampling handler
+ * // After tool execution, `samplingCalls` contains metrics for all nested LLM calls
+ * ```
+ */
+export declare function wrapSamplingHandler(handler: McpSamplingHandler, onSamplingCall: (data: SamplingCallData) => void): McpSamplingHandler;
+/**
+ * Creates a collector that accumulates sampling call data during tool execution.
+ * Use with `wrapSamplingHandler` to capture all sampling calls for a tool invocation.
+ *
+ * @example
+ * ```ts
+ * const collector = createSamplingCollector();
+ * const wrappedHandler = wrapSamplingHandler(handler, collector.collect);
+ * // ... execute MCP tool ...
+ * const calls = collector.getCalls(); // SamplingCallData[]
+ * ```
+ */
+export declare function createSamplingCollector(): {
+    collect: (data: SamplingCallData) => number;
+    getCalls: () => SamplingCallData[];
+    reset: () => void;
+};
+//# sourceMappingURL=instrumentMcpSampling.d.ts.map

package/dist/instrumentMcpSampling.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"instrumentMcpSampling.d.ts","sourceRoot":"","sources":["../src/instrumentMcpSampling.ts"],"names":[],"mappings":"AAAA;;;;;;GAMG;AAEH,MAAM,MAAM,gBAAgB,GAAG;IAC7B,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB,aAAa,CAAC,EAAE,MAAM,CAAC;IACvB,WAAW,CAAC,EAAE,MAAM,CAAC;CACtB,CAAC;AAEF,MAAM,MAAM,kBAAkB,GAAG,CAC/B,OAAO,EAAE,kBAAkB,KACxB,OAAO,CAAC,mBAAmB,CAAC,CAAC;AAElC,MAAM,MAAM,kBAAkB,GAAG;IAC/B,MAAM,EAAE,wBAAwB,CAAC;IACjC,MAAM,EAAE;QACN,QAAQ,EAAE,OAAO,EAAE,CAAC;QACpB,gBAAgB,CAAC,EAAE;YAAE,KAAK,CAAC,EAAE;gBAAE,IAAI,CAAC,EAAE,MAAM,CAAA;aAAE,EAAE,CAAA;SAAE,CAAC;QACnD,SAAS,CAAC,EAAE,MAAM,CAAC;QACnB,CAAC,GAAG,EAAE,MAAM,GAAG,OAAO,CAAC;KACxB,CAAC;CACH,CAAC;AAEF,MAAM,MAAM,mBAAmB,GAAG;IAChC,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,OAAO,EAAE,OAAO,CAAC;IACjB,KAAK,CAAC,EAAE;QACN,WAAW,CAAC,EAAE,MAAM,CAAC;QACrB,YAAY,CAAC,EAAE,MAAM,CAAC;QACtB,YAAY,CAAC,EAAE,MAAM,CAAC;QACtB,gBAAgB,CAAC,EAAE,MAAM,CAAC;KAC3B,CAAC;IACF,CAAC,GAAG,EAAE,MAAM,GAAG,OAAO,CAAC;CACxB,CAAC;AAEF;;;;;;;;;;;;;;;;;GAiBG;AACH,wBAAgB,mBAAmB,CACjC,OAAO,EAAE,kBAAkB,EAC3B,cAAc,EAAE,CAAC,IAAI,EAAE,gBAAgB,KAAK,IAAI,GAC/C,kBAAkB,CAuBpB;AAED;;;;;;;;;;;GAWG;AACH,wBAAgB,uBAAuB;oBAGnB,gBAAgB;;;EAMnC"}

package/dist/instrumentMcpSampling.js

@@ -0,0 +1,65 @@
+/**
+ * MCP sampling instrumentation utility.
+ *
+ * Wraps an MCP client's sampling handler to capture nested LLM calls
+ * (sampling/createMessage requests) made during tool execution.
+ * The captured data can be reported as child generation spans.
+ */
+/**
+ * Wraps an MCP sampling handler to intercept and measure sampling calls.
+ *
+ * @param handler - The original sampling handler from the MCP client
+ * @param onSamplingCall - Callback invoked with metrics for each sampling call
+ * @returns A wrapped handler that transparently captures sampling metrics
+ *
+ * @example
+ * ```ts
+ * const samplingCalls: SamplingCallData[] = [];
+ * const wrapped = wrapSamplingHandler(
+ *   originalHandler,
+ *   (data) => samplingCalls.push(data),
+ * );
+ * // Use `wrapped` as the MCP client's sampling handler
+ * // After tool execution, `samplingCalls` contains metrics for all nested LLM calls
+ * ```
+ */
+export function wrapSamplingHandler(handler, onSamplingCall) {
+    return async (request) => {
+        const startTime = Date.now();
+        const response = await handler(request);
+        const durationMs = Date.now() - startTime;
+        const modelId = response.model ?? request.params.modelPreferences?.hints?.[0]?.name;
+        const inputTokens = response.usage?.inputTokens ?? response.usage?.promptTokens;
+        const outputTokens = response.usage?.outputTokens ?? response.usage?.completionTokens;
+        onSamplingCall({
+            ...(modelId ? { model_id: modelId } : undefined),
+            ...(inputTokens != null ? { input_tokens: inputTokens } : undefined),
+            ...(outputTokens != null ? { output_tokens: outputTokens } : undefined),
+            duration_ms: durationMs,
+        });
+        return response;
+    };
+}
+/**
+ * Creates a collector that accumulates sampling call data during tool execution.
+ * Use with `wrapSamplingHandler` to capture all sampling calls for a tool invocation.
+ *
+ * @example
+ * ```ts
+ * const collector = createSamplingCollector();
+ * const wrappedHandler = wrapSamplingHandler(handler, collector.collect);
+ * // ... execute MCP tool ...
+ * const calls = collector.getCalls(); // SamplingCallData[]
+ * ```
+ */
+export function createSamplingCollector() {
+    const calls = [];
+    return {
+        collect: (data) => calls.push(data),
+        getCalls: () => [...calls],
+        reset: () => {
+            calls.length = 0;
+        },
+    };
+}
+//# sourceMappingURL=instrumentMcpSampling.js.map

package/dist/instrumentMcpSampling.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"instrumentMcpSampling.js","sourceRoot":"","sources":["../src/instrumentMcpSampling.ts"],"names":[],"mappings":"AAAA;;;;;;GAMG;AAmCH;;;;;;;;;;;;;;;;;GAiBG;AACH,MAAM,UAAU,mBAAmB,CACjC,OAA2B,EAC3B,cAAgD;IAEhD,OAAO,KAAK,EAAE,OAAO,EAAE,EAAE;QACvB,MAAM,SAAS,GAAG,IAAI,CAAC,GAAG,EAAE,CAAC;QAC7B,MAAM,QAAQ,GAAG,MAAM,OAAO,CAAC,OAAO,CAAC,CAAC;QACxC,MAAM,UAAU,GAAG,IAAI,CAAC,GAAG,EAAE,GAAG,SAAS,CAAC;QAE1C,MAAM,OAAO,GACX,QAAQ,CAAC,KAAK,IAAI,OAAO,CAAC,MAAM,CAAC,gBAAgB,EAAE,KAAK,EAAE,CAAC,CAAC,CAAC,EAAE,IAAI,CAAC;QAEtE,MAAM,WAAW,GACf,QAAQ,CAAC,KAAK,EAAE,WAAW,IAAI,QAAQ,CAAC,KAAK,EAAE,YAAY,CAAC;QAC9D,MAAM,YAAY,GAChB,QAAQ,CAAC,KAAK,EAAE,YAAY,IAAI,QAAQ,CAAC,KAAK,EAAE,gBAAgB,CAAC;QAEnE,cAAc,CAAC;YACb,GAAG,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,QAAQ,EAAE,OAAO,EAAE,CAAC,CAAC,CAAC,SAAS,CAAC;YAChD,GAAG,CAAC,WAAW,IAAI,IAAI,CAAC,CAAC,CAAC,EAAE,YAAY,EAAE,WAAW,EAAE,CAAC,CAAC,CAAC,SAAS,CAAC;YACpE,GAAG,CAAC,YAAY,IAAI,IAAI,CAAC,CAAC,CAAC,EAAE,aAAa,EAAE,YAAY,EAAE,CAAC,CAAC,CAAC,SAAS,CAAC;YACvE,WAAW,EAAE,UAAU;SACxB,CAAC,CAAC;QAEH,OAAO,QAAQ,CAAC;IAClB,CAAC,CAAC;AACJ,CAAC;AAED;;;;;;;;;;;GAWG;AACH,MAAM,UAAU,uBAAuB;IACrC,MAAM,KAAK,GAAuB,EAAE,CAAC;IACrC,OAAO;QACL,OAAO,EAAE,CAAC,IAAsB,EAAE,EAAE,CAAC,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC;QACrD,QAAQ,EAAE,GAAG,EAAE,CAAC,CAAC,GAAG,KAAK,CAAC;QAC1B,KAAK,EAAE,GAAG,EAAE;YACV,KAAK,CAAC,MAAM,GAAG,CAAC,CAAC;QACnB,CAAC;KACF,CAAC;AACJ,CAAC"}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "assistant-cloud",
-  "version": "0.1.16",
+  "version": "0.1.18",
   "description": "Cloud integration for assistant-ui",
   "keywords": [
     "assistant",
@@ -28,10 +28,10 @@
   ],
   "sideEffects": false,
   "dependencies": {
-    "assistant-stream": "^0.3.1"
+    "assistant-stream": "^0.3.3"
   },
   "devDependencies": {
-    "@types/node": "^25.2.0",
+    "@types/node": "^25.2.1",
     "vitest": "^4.0.18",
     "@assistant-ui/x-buildutils": "0.0.1"
   },

package/src/{AssistantCloud.tsx → AssistantCloud.ts}

@@ -1,4 +1,8 @@
-import { AssistantCloudAPI, AssistantCloudConfig } from "./AssistantCloudAPI";
+import {
+  AssistantCloudAPI,
+  AssistantCloudConfig,
+  AssistantCloudTelemetryConfig,
+} from "./AssistantCloudAPI";
 import { AssistantCloudAuthTokens } from "./AssistantCloudAuthTokens";
 import { AssistantCloudRuns } from "./AssistantCloudRuns";
 import { AssistantCloudThreads } from "./AssistantCloudThreads";
@@ -9,6 +13,7 @@ export class AssistantCloud {
   public readonly auth;
   public readonly runs;
   public readonly files;
+  public readonly telemetry: AssistantCloudTelemetryConfig;
 
   constructor(config: AssistantCloudConfig) {
     const api = new AssistantCloudAPI(config);
@@ -18,5 +23,13 @@ export class AssistantCloud {
     };
     this.runs = new AssistantCloudRuns(api);
     this.files = new AssistantCloudFiles(api);
+
+    const t = config.telemetry;
+    this.telemetry =
+      t === false
+        ? { enabled: false }
+        : t === true || t === undefined
+          ? { enabled: true }
+          : { enabled: t.enabled !== false, ...t };
   }
 }

package/src/{AssistantCloudAPI.tsx → AssistantCloudAPI.ts}

@@ -4,8 +4,21 @@ import {
   AssistantCloudAPIKeyAuthStrategy,
   AssistantCloudAnonymousAuthStrategy,
 } from "./AssistantCloudAuthStrategy";
+import type { AssistantCloudRunReport } from "./AssistantCloudRuns";
 
-export type AssistantCloudConfig =
+export type AssistantCloudTelemetryConfig = {
+  enabled?: boolean;
+  /**
+   * Called before each telemetry report is sent.
+   * Return a modified report to enrich it (e.g. add `model_id`),
+   * or return `null` to skip the report.
+   */
+  beforeReport?: (
+    report: AssistantCloudRunReport,
+  ) => AssistantCloudRunReport | null;
+};
+
+export type AssistantCloudConfig = (
   | {
       baseUrl: string;
       authToken: () => Promise<string | null>;
@@ -18,7 +31,21 @@ export type AssistantCloudConfig =
   | {
       baseUrl: string;
       anonymous: true;
-    };
+    }
+) & {
+  /**
+   * Client-side run telemetry reporting. Default: `true`.
+   *
+   * When enabled, the SDK automatically reports run metadata (status, step
+   * count, tool calls, and token usage) to Assistant Cloud after each
+   * assistant message is saved. No message content is sent.
+   *
+   * - `true` / `undefined` — enabled with defaults
+   * - `false` — disabled
+   * - `{ beforeReport }` — enabled with a hook to enrich or filter reports
+   */
+  telemetry?: boolean | AssistantCloudTelemetryConfig;
+};
 
 class CloudAPIError extends Error {
   constructor(message: string) {
@@ -104,7 +131,8 @@ export class AssistantCloudAPI {
       try {
         const body = JSON.parse(text);
         throw new CloudAPIError(body.message);
-      } catch {
+      } catch (error) {
+        if (error instanceof CloudAPIError) throw error;
         throw new Error(
           `Request failed with status ${response.status}, ${text}`,
         );

package/src/AssistantCloudRuns.ts

@@ -0,0 +1,99 @@
+import { AssistantCloudAPI } from "./AssistantCloudAPI";
+import { AssistantStream, PlainTextDecoder } from "assistant-stream";
+
+type AssistantCloudRunsStreamBody = {
+  thread_id: string;
+  assistant_id: "system/thread_title";
+  messages: readonly unknown[]; // TODO type
+};
+
+export type AssistantCloudRunReport = {
+  thread_id: string;
+  status: "completed" | "incomplete" | "error";
+  total_steps?: number;
+  tool_calls?: {
+    tool_name: string;
+    tool_call_id: string;
+    tool_args?: string;
+    tool_result?: string;
+    tool_source?: "mcp" | "frontend" | "backend";
+    start_ms?: number;
+    end_ms?: number;
+    sampling_calls?: {
+      model_id?: string;
+      input_tokens?: number;
+      output_tokens?: number;
+      duration_ms?: number;
+    }[];
+  }[];
+  steps?: {
+    input_tokens?: number;
+    output_tokens?: number;
+    tool_calls?: {
+      tool_name: string;
+      tool_call_id: string;
+      tool_args?: string;
+      tool_result?: string;
+      tool_source?: "mcp" | "frontend" | "backend";
+      start_ms?: number;
+      end_ms?: number;
+      sampling_calls?: {
+        model_id?: string;
+        input_tokens?: number;
+        output_tokens?: number;
+        duration_ms?: number;
+      }[];
+    }[];
+    start_ms?: number;
+    end_ms?: number;
+  }[];
+  input_tokens?: number;
+  output_tokens?: number;
+  model_id?: string;
+  provider_type?: string;
+  duration_ms?: number;
+  output_text?: string;
+  metadata?: Record<string, unknown>;
+};
+
+export class AssistantCloudRuns {
+  constructor(private cloud: AssistantCloudAPI) {}
+
+  public __internal_getAssistantOptions(assistantId: string) {
+    return {
+      api: `${this.cloud._baseUrl}/v1/runs/stream`,
+      headers: async () => {
+        const headers = await this.cloud._auth.getAuthHeaders();
+        if (!headers) throw new Error("Authorization failed");
+        return {
+          ...headers,
+          Accept: "text/plain",
+        };
+      },
+      body: {
+        assistant_id: assistantId,
+        response_format: "vercel-ai-data-stream/v1",
+        thread_id: "unstable_todo",
+      },
+    };
+  }
+
+  public async stream(
+    body: AssistantCloudRunsStreamBody,
+  ): Promise<AssistantStream> {
+    const response = await this.cloud.makeRawRequest("/runs/stream", {
+      method: "POST",
+      headers: {
+        Accept: "text/plain",
+      },
+      body,
+    });
+    return AssistantStream.fromResponse(response, new PlainTextDecoder());
+  }
+
+  public async report(
+    body: AssistantCloudRunReport,
+  ): Promise<{ run_id: string }> {
+    return this.cloud.makeRequest("/runs", { method: "POST", body });
+  }
+}

package/src/{AssistantCloudThreadMessages.tsx → AssistantCloudThreadMessages.ts}

@@ -29,6 +29,10 @@ type AssistantCloudMessageCreateResponse = {
   message_id: string;
 };
 
+type AssistantCloudThreadMessageUpdateBody = {
+  content: ReadonlyJSONObject;
+};
+
 export class AssistantCloudThreadMessages {
   constructor(private cloud: AssistantCloudAPI) {}
 
@@ -51,4 +55,15 @@ export class AssistantCloudThreadMessages {
       { method: "POST", body },
     );
   }
+
+  public async update(
+    threadId: string,
+    messageId: string,
+    body: AssistantCloudThreadMessageUpdateBody,
+  ): Promise<void> {
+    return this.cloud.makeRequest(
+      `/threads/${encodeURIComponent(threadId)}/messages/${encodeURIComponent(messageId)}`,
+      { method: "PUT", body },
+    );
+  }
 }

package/src/CloudMessagePersistence.ts

@@ -0,0 +1,123 @@
+import type { ReadonlyJSONObject } from "assistant-stream/utils";
+import type { AssistantCloud } from "./AssistantCloud";
+
+/**
+ * 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.
+ */
+export class CloudMessagePersistence {
+  private idMapping: Record<string, string | Promise<string>> = {};
+
+  constructor(private cloud: AssistantCloud) {}
+
+  /**
+   * Persist a message to the cloud.
+   *
+   * @param threadId - Remote thread ID
+   * @param messageId - Local message ID (used for tracking)
+   * @param parentId - Local parent message ID (or null for first message)
+   * @param format - Message format (e.g., "aui/v0", "ai-sdk/v6")
+   * @param content - Message content (format-specific)
+   */
+  async append(
+    threadId: string,
+    messageId: string,
+    parentId: string | null,
+    format: string,
+    content: ReadonlyJSONObject,
+  ): Promise<void> {
+    // Resolve parent's remote ID if it exists (may be a promise if concurrent)
+    const resolvedParentId = parentId
+      ? ((await this.idMapping[parentId]) ?? parentId)
+      : null;
+
+    const task = this.cloud.threads.messages
+      .create(threadId, {
+        parent_id: resolvedParentId,
+        format,
+        content,
+      })
+      .then(({ message_id }) => {
+        this.idMapping[messageId] = message_id;
+        return message_id;
+      })
+      .catch((err) => {
+        // Only delete if we're still the active task (avoids clobbering a retry)
+        if (this.idMapping[messageId] === task) {
+          delete this.idMapping[messageId];
+        }
+        throw err;
+      });
+
+    // Store the promise immediately so concurrent appends can await it
+    this.idMapping[messageId] = task;
+    return task.then(() => {});
+  }
+
+  /**
+   * Update an already-persisted message in the cloud.
+   */
+  async update(
+    threadId: string,
+    messageId: string,
+    _format: string,
+    content: ReadonlyJSONObject,
+  ): Promise<void> {
+    const remoteId = await this.getRemoteId(messageId);
+    if (!remoteId) return; // not persisted yet, skip
+    await this.cloud.threads.messages.update(threadId, remoteId, { content });
+  }
+
+  /**
+   * Check if a message has been persisted (or is currently being persisted).
+   */
+  isPersisted(messageId: string): boolean {
+    return messageId in this.idMapping;
+  }
+
+  /**
+   * Get the remote ID for a local message ID (resolved).
+   * Returns undefined if not persisted.
+   */
+  async getRemoteId(messageId: string): Promise<string | undefined> {
+    const entry = this.idMapping[messageId];
+    if (!entry) return undefined;
+    return entry;
+  }
+
+  /**
+   * Load messages from the cloud and populate the ID mapping.
+   *
+   * The ID mapping is populated so that `isPersisted()` returns true for
+   * loaded messages, preventing re-persistence of already-stored messages.
+   *
+   * @param threadId - Remote thread ID
+   * @param format - Optional format filter
+   * @returns Array of cloud messages
+   */
+  async load(threadId: string, format?: string) {
+    const { messages } = await this.cloud.threads.messages.list(
+      threadId,
+      format ? { format } : undefined,
+    );
+    // Populate ID mapping so isPersisted() recognizes loaded messages
+    for (const m of messages) {
+      this.idMapping[m.id] = m.id;
+    }
+    return messages;
+  }
+
+  /**
+   * Reset the ID mapping (call when switching threads).
+   */
+  reset() {
+    this.idMapping = {};
+  }
+}

package/src/FormattedCloudPersistence.ts

@@ -0,0 +1,96 @@
+import type { ReadonlyJSONObject } from "assistant-stream/utils";
+
+/**
+ * Format adapter shape — structurally identical to the MessageFormatAdapter
+ * in @assistant-ui/react, but defined here to avoid cross-package type moves.
+ * TypeScript's structural typing ensures these are interchangeable.
+ */
+export type MessageFormatAdapter<TMessage, TStorageFormat> = {
+  format: string;
+  encode(item: { parentId: string | null; message: TMessage }): TStorageFormat;
+  decode(stored: {
+    id: string;
+    parent_id: string | null;
+    format: string;
+    content: TStorageFormat;
+  }): { parentId: string | null; message: TMessage };
+  getId(message: TMessage): string;
+};
+
+/**
+ * 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.
+ */
+export const createFormattedPersistence = <TMessage, TStorageFormat>(
+  persistence: {
+    append: (
+      threadId: string,
+      messageId: string,
+      parentId: string | null,
+      format: string,
+      content: ReadonlyJSONObject,
+    ) => Promise<void>;
+    load: (threadId: string, format?: string) => Promise<any[]>;
+    isPersisted: (messageId: string) => boolean;
+    update?: (
+      threadId: string,
+      messageId: string,
+      format: string,
+      content: ReadonlyJSONObject,
+    ) => Promise<void>;
+  },
+  adapter: MessageFormatAdapter<TMessage, TStorageFormat>,
+) => ({
+  append: async (
+    threadId: string,
+    item: { parentId: string | null; message: TMessage },
+  ): Promise<void> => {
+    const messageId = adapter.getId(item.message);
+    const encoded = adapter.encode(item);
+    return persistence.append(
+      threadId,
+      messageId,
+      item.parentId,
+      adapter.format,
+      encoded as ReadonlyJSONObject,
+    );
+  },
+  update: persistence.update
+    ? async (
+        threadId: string,
+        item: { parentId: string | null; message: TMessage },
+        messageId: string,
+      ): Promise<void> => {
+        const encoded = adapter.encode(item);
+        return persistence.update!(
+          threadId,
+          messageId,
+          adapter.format,
+          encoded as ReadonlyJSONObject,
+        );
+      }
+    : undefined,
+  load: async (threadId: string) => {
+    const messages = await persistence.load(threadId, adapter.format);
+    return {
+      messages: messages
+        .filter((m) => m.format === adapter.format)
+        .map((m) =>
+          adapter.decode({
+            id: m.id,
+            parent_id: m.parent_id,
+            format: m.format,
+            content: m.content as TStorageFormat,
+          }),
+        )
+        .reverse(),
+    };
+  },
+  isPersisted: (messageId: string) => persistence.isPersisted(messageId),
+});

package/src/index.ts

@@ -1,2 +1,15 @@
 export type { CloudMessage } from "./AssistantCloudThreadMessages";
+export type { AssistantCloudTelemetryConfig } from "./AssistantCloudAPI";
+export type { AssistantCloudRunReport } from "./AssistantCloudRuns";
 export { AssistantCloud } from "./AssistantCloud";
+export { CloudMessagePersistence } from "./CloudMessagePersistence";
+export {
+  createFormattedPersistence,
+  type MessageFormatAdapter,
+} from "./FormattedCloudPersistence";
+export {
+  wrapSamplingHandler,
+  createSamplingCollector,
+  type SamplingCallData,
+  type McpSamplingHandler,
+} from "./instrumentMcpSampling";