@assistant-ui/core 0.2.20 → 0.2.22

package/src/react/adapters/LocalStorageThreadListAdapter.tsx

@@ -8,6 +8,7 @@ import type {
   RemoteThreadMetadata,
   ThreadHistoryAdapter,
   ThreadMessage,
+  RunConfig,
 } from "../../index";
 import type {
   ExportedMessageRepository,
@@ -22,6 +23,37 @@ export type AsyncStorageLike = {
   removeItem(key: string): Promise<void>;
 };
 
+class KeyedMutationQueue {
+  private readonly tails = new Map<string, Promise<void>>();
+
+  run<T>(key: string, mutation: () => Promise<T>): Promise<T> {
+    const previous = this.tails.get(key);
+    const result = previous ? previous.then(mutation) : mutation();
+    const tail = result.then(
+      () => undefined,
+      () => undefined,
+    );
+
+    this.tails.set(key, tail);
+    void tail.then(() => {
+      if (this.tails.get(key) === tail) this.tails.delete(key);
+    });
+
+    return result;
+  }
+}
+
+const mutationQueues = new WeakMap<AsyncStorageLike, KeyedMutationQueue>();
+
+const getMutationQueue = (storage: AsyncStorageLike): KeyedMutationQueue => {
+  let queue = mutationQueues.get(storage);
+  if (!queue) {
+    queue = new KeyedMutationQueue();
+    mutationQueues.set(storage, queue);
+  }
+  return queue;
+};
+
 type LocalStorageAdapterOptions = {
   storage: AsyncStorageLike;
   prefix?: string | undefined;
@@ -36,11 +68,218 @@ type StoredThreadMetadata = {
   custom?: Record<string, unknown> | undefined;
 };
 
+type StoredSystemMessage = Extract<ThreadMessage, { role: "system" }>;
+type StoredUserMessage = Extract<ThreadMessage, { role: "user" }>;
+type StoredAssistantMessage = Extract<ThreadMessage, { role: "assistant" }>;
+
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === "object" && value !== null && !Array.isArray(value);
+
+const parseJSON = (raw: string | null): unknown => {
+  if (!raw) return undefined;
+  try {
+    return JSON.parse(raw) as unknown;
+  } catch {
+    return undefined;
+  }
+};
+
+const parseStoredThread = (value: unknown): StoredThreadMetadata | null => {
+  if (!isRecord(value) || typeof value.remoteId !== "string") return null;
+
+  const status = value.status ?? "regular";
+  if (status !== "regular" && status !== "archived") return null;
+
+  return {
+    remoteId: value.remoteId,
+    status,
+    ...(typeof value.externalId === "string"
+      ? { externalId: value.externalId }
+      : undefined),
+    ...(typeof value.title === "string" ? { title: value.title } : undefined),
+    ...(isRecord(value.custom) ? { custom: value.custom } : undefined),
+  };
+};
+
+const parseDate = (value: unknown): Date | null => {
+  if (value instanceof Date && !Number.isNaN(value.getTime())) return value;
+  if (typeof value !== "string" && typeof value !== "number") return null;
+
+  const date = new Date(value);
+  return Number.isNaN(date.getTime()) ? null : date;
+};
+
+const isMessageRole = (value: unknown): value is ThreadMessage["role"] =>
+  value === "system" || value === "user" || value === "assistant";
+
+const parseStoredThreadMessage = (value: unknown): ThreadMessage | null => {
+  if (!isRecord(value) || typeof value.id !== "string") return null;
+  if (!isMessageRole(value.role)) return null;
+  if (!Array.isArray(value.content)) return null;
+
+  const createdAt = parseDate(value.createdAt);
+  if (!createdAt) return null;
+
+  const metadata = value.metadata;
+  if (!isRecord(metadata) || !isRecord(metadata.custom)) return null;
+
+  if (value.role === "assistant") {
+    const status = value.status;
+    if (!isRecord(status) || typeof status.type !== "string") return null;
+
+    const submittedFeedback = isRecord(metadata.submittedFeedback)
+      ? metadata.submittedFeedback
+      : undefined;
+    const submittedFeedbackType = submittedFeedback?.type;
+
+    return {
+      id: value.id,
+      role: "assistant",
+      content: value.content as StoredAssistantMessage["content"],
+      status: status as StoredAssistantMessage["status"],
+      createdAt,
+      metadata: {
+        unstable_state: (metadata.unstable_state ??
+          null) as StoredAssistantMessage["metadata"]["unstable_state"],
+        unstable_annotations: Array.isArray(metadata.unstable_annotations)
+          ? (metadata.unstable_annotations as StoredAssistantMessage["metadata"]["unstable_annotations"])
+          : [],
+        unstable_data: Array.isArray(metadata.unstable_data)
+          ? (metadata.unstable_data as StoredAssistantMessage["metadata"]["unstable_data"])
+          : [],
+        steps: Array.isArray(metadata.steps)
+          ? (metadata.steps as StoredAssistantMessage["metadata"]["steps"])
+          : [],
+        ...(submittedFeedbackType === "positive" ||
+        submittedFeedbackType === "negative"
+          ? {
+              submittedFeedback: {
+                type: submittedFeedbackType,
+              },
+            }
+          : undefined),
+        ...(metadata.timing !== undefined
+          ? {
+              timing: metadata.timing as NonNullable<
+                StoredAssistantMessage["metadata"]["timing"]
+              >,
+            }
+          : undefined),
+        ...(metadata.isOptimistic === true
+          ? { isOptimistic: true }
+          : undefined),
+        custom: metadata.custom,
+      },
+    };
+  }
+
+  if (value.role === "user") {
+    return {
+      id: value.id,
+      role: "user",
+      content: value.content as StoredUserMessage["content"],
+      attachments: Array.isArray(value.attachments)
+        ? (value.attachments as StoredUserMessage["attachments"])
+        : [],
+      createdAt,
+      metadata: {
+        custom: metadata.custom,
+      },
+    };
+  }
+
+  if (value.content.length !== 1) return null;
+
+  return {
+    id: value.id,
+    role: "system",
+    content: [value.content[0] as StoredSystemMessage["content"][0]],
+    createdAt,
+    metadata: {
+      custom: metadata.custom,
+    },
+  };
+};
+
+export const parseStoredThreadMetadata = (
+  raw: string | null,
+): StoredThreadMetadata[] => {
+  const parsed = parseJSON(raw);
+  if (!Array.isArray(parsed)) return [];
+
+  return parsed.flatMap((item) => {
+    const thread = parseStoredThread(item);
+    return thread ? [thread] : [];
+  });
+};
+
+const parseStoredMessageRepositoryItem = (
+  value: unknown,
+): ExportedMessageRepositoryItem | null => {
+  if (!isRecord(value)) return null;
+
+  const message = parseStoredThreadMessage(value.message);
+  if (!message) return null;
+
+  const parentId = value.parentId;
+  if (
+    parentId !== undefined &&
+    parentId !== null &&
+    typeof parentId !== "string"
+  ) {
+    return null;
+  }
+
+  return {
+    message,
+    parentId: parentId ?? null,
+    ...(isRecord(value.runConfig)
+      ? { runConfig: value.runConfig as RunConfig }
+      : undefined),
+  };
+};
+
+export const parseStoredMessageRepository = (
+  raw: string | null,
+): ExportedMessageRepository => {
+  const parsed = parseJSON(raw);
+  if (!isRecord(parsed) || !Array.isArray(parsed.messages)) {
+    return { messages: [] };
+  }
+
+  const candidateMessages = parsed.messages.flatMap((item) => {
+    const parsedItem = parseStoredMessageRepositoryItem(item);
+    return parsedItem ? [parsedItem] : [];
+  });
+
+  const acceptedIds = new Set<string>();
+  const messages = candidateMessages.flatMap((item) => {
+    if (acceptedIds.has(item.message.id)) return [];
+    if (item.parentId !== null && !acceptedIds.has(item.parentId)) return [];
+
+    acceptedIds.add(item.message.id);
+    return [item];
+  });
+
+  const headId =
+    parsed.headId === null ||
+    (typeof parsed.headId === "string" &&
+      messages.some((item) => item.message.id === parsed.headId))
+      ? parsed.headId
+      : undefined;
+
+  return {
+    ...(headId !== undefined ? { headId } : undefined),
+    messages,
+  };
+};
+
 class AsyncStorageHistoryAdapter implements ThreadHistoryAdapter {
   constructor(
     private storage: AsyncStorageLike,
     private aui: ReturnType<typeof useAui>,
     private prefix: string,
+    private mutationQueue: KeyedMutationQueue,
   ) {}
 
   private _messagesKey(remoteId: string) {
@@ -52,41 +291,41 @@ class AsyncStorageHistoryAdapter implements ThreadHistoryAdapter {
     if (!remoteId) return { messages: [] };
 
     const raw = await this.storage.getItem(this._messagesKey(remoteId));
-    if (!raw) return { messages: [] };
-    return JSON.parse(raw) as ExportedMessageRepository;
+    return parseStoredMessageRepository(raw);
   }
 
   async append(item: ExportedMessageRepositoryItem): Promise<void> {
     const { remoteId } = await this.aui.threadListItem().initialize();
 
     const key = this._messagesKey(remoteId);
-    const raw = await this.storage.getItem(key);
-    const repo: ExportedMessageRepository = raw
-      ? (JSON.parse(raw) as ExportedMessageRepository)
-      : { messages: [] };
+    await this.mutationQueue.run(key, async () => {
+      const raw = await this.storage.getItem(key);
+      const repo = parseStoredMessageRepository(raw);
+
+      const idx = repo.messages.findIndex(
+        (m) => m.message.id === item.message.id,
+      );
+      if (idx >= 0) {
+        repo.messages[idx] = item;
+      } else {
+        repo.messages.push(item);
+      }
+      repo.headId = item.message.id;
 
-    const idx = repo.messages.findIndex(
-      (m) => m.message.id === item.message.id,
-    );
-    if (idx >= 0) {
-      repo.messages[idx] = item;
-    } else {
-      repo.messages.push(item);
-    }
-    repo.headId = item.message.id;
-
-    await this.storage.setItem(key, JSON.stringify(repo));
+      await this.storage.setItem(key, JSON.stringify(repo));
+    });
   }
 }
 
 const createHistoryProvider = (
   storage: AsyncStorageLike,
   prefix: string,
+  mutationQueue: KeyedMutationQueue,
 ): FC<PropsWithChildren> => {
   const Provider: FC<PropsWithChildren> = ({ children }) => {
     const aui = useAui();
     const history = useMemo(
-      () => new AsyncStorageHistoryAdapter(storage, aui, prefix),
+      () => new AsyncStorageHistoryAdapter(storage, aui, prefix, mutationQueue),
       [aui],
     );
     const adapters = useMemo(() => ({ history }), [history]);
@@ -107,10 +346,11 @@ export const createLocalStorageAdapter = (
 
   const threadsKey = `${prefix}threads`;
   const messagesKey = (threadId: string) => `${prefix}messages:${threadId}`;
+  const mutationQueue = getMutationQueue(storage);
 
   const loadThreadMetadata = async (): Promise<StoredThreadMetadata[]> => {
     const raw = await storage.getItem(threadsKey);
-    return raw ? (JSON.parse(raw) as StoredThreadMetadata[]) : [];
+    return parseStoredThreadMetadata(raw);
   };
 
   const saveThreadMetadata = async (
@@ -119,8 +359,22 @@ export const createLocalStorageAdapter = (
     await storage.setItem(threadsKey, JSON.stringify(threads));
   };
 
+  const updateThreadMetadata = async (
+    remoteId: string,
+    update: (thread: StoredThreadMetadata) => void,
+  ): Promise<void> => {
+    await mutationQueue.run(threadsKey, async () => {
+      const threads = await loadThreadMetadata();
+      const thread = threads.find((item) => item.remoteId === remoteId);
+      if (thread) {
+        update(thread);
+        await saveThreadMetadata(threads);
+      }
+    });
+  };
+
   const adapter: RemoteThreadListAdapter = {
-    unstable_Provider: createHistoryProvider(storage, prefix),
+    unstable_Provider: createHistoryProvider(storage, prefix, mutationQueue),
 
     async list(): Promise<RemoteThreadListResponse> {
       const threads = await loadThreadMetadata();
@@ -139,70 +393,66 @@ export const createLocalStorageAdapter = (
       threadId: string,
     ): Promise<RemoteThreadInitializeResponse> {
       const remoteId = threadId;
-      const threads = await loadThreadMetadata();
+      return mutationQueue.run(threadsKey, async () => {
+        const threads = await loadThreadMetadata();
 
-      // Only add if not already present
-      if (!threads.some((t) => t.remoteId === remoteId)) {
-        threads.unshift({
-          remoteId,
-          status: "regular",
-        });
-        await saveThreadMetadata(threads);
-      }
+        // Only add if not already present
+        if (!threads.some((t) => t.remoteId === remoteId)) {
+          threads.unshift({
+            remoteId,
+            status: "regular",
+          });
+          await saveThreadMetadata(threads);
+        }
 
-      return { remoteId, externalId: undefined };
+        return { remoteId, externalId: undefined };
+      });
     },
 
     async rename(remoteId: string, newTitle: string): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const thread = threads.find((t) => t.remoteId === remoteId);
-      if (thread) {
+      await updateThreadMetadata(remoteId, (thread) => {
         thread.title = newTitle;
-        await saveThreadMetadata(threads);
-      }
+      });
     },
 
     async updateCustom(
       remoteId: string,
       custom: Record<string, unknown> | undefined,
     ): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const thread = threads.find((t) => t.remoteId === remoteId);
-      if (thread) {
+      await updateThreadMetadata(remoteId, (thread) => {
         thread.custom = custom;
-        await saveThreadMetadata(threads);
-      }
+      });
     },
 
     async archive(remoteId: string): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const thread = threads.find((t) => t.remoteId === remoteId);
-      if (thread) {
+      await updateThreadMetadata(remoteId, (thread) => {
         thread.status = "archived";
-        await saveThreadMetadata(threads);
-      }
+      });
     },
 
     async unarchive(remoteId: string): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const thread = threads.find((t) => t.remoteId === remoteId);
-      if (thread) {
+      await updateThreadMetadata(remoteId, (thread) => {
         thread.status = "regular";
-        await saveThreadMetadata(threads);
-      }
+      });
     },
 
     async delete(remoteId: string): Promise<void> {
-      const threads = await loadThreadMetadata();
-      const filtered = threads.filter((t) => t.remoteId !== remoteId);
-      await saveThreadMetadata(filtered);
-      await storage.removeItem(messagesKey(remoteId));
+      await mutationQueue.run(threadsKey, async () => {
+        const threads = await loadThreadMetadata();
+        const filtered = threads.filter((t) => t.remoteId !== remoteId);
+        await saveThreadMetadata(filtered);
+      });
+      const key = messagesKey(remoteId);
+      await mutationQueue.run(key, () => storage.removeItem(key));
     },
 
     async fetch(threadId: string): Promise<RemoteThreadMetadata> {
       const threads = await loadThreadMetadata();
       const thread = threads.find((t) => t.remoteId === threadId);
-      if (!thread) throw new Error("Thread not found");
+      if (!thread)
+        throw new Error(
+          `Stored thread "${threadId}" not found while fetching thread metadata.`,
+        );
       return {
         remoteId: thread.remoteId,
         externalId: thread.externalId,
@@ -220,12 +470,9 @@ export const createLocalStorageAdapter = (
         const title = await titleGenerator.generateTitle(messages);
 
         // Update the stored title
-        const threads = await loadThreadMetadata();
-        const thread = threads.find((t) => t.remoteId === remoteId);
-        if (thread) {
+        await updateThreadMetadata(remoteId, (thread) => {
           thread.title = title;
-          await saveThreadMetadata(threads);
-        }
+        });
 
         // Return a stream with a single text part
         return createAssistantStream((controller) => {

package/src/react/index.ts

@@ -50,7 +50,9 @@ export {
 } from "./model-context/provider-tool";
 export {
   defineMcpToolkit,
+  type McpToolkitEntry,
   type McpToolkitDefinition,
+  type McpToolkitToolConfig,
 } from "./model-context/define-mcp-toolkit";
 /**
  * @deprecated Since 2026-06-14 — migrate to the Unstable / Experimental API.

package/src/react/model-context/define-mcp-toolkit.ts

@@ -1,16 +1,49 @@
 import type { McpServerConfig } from "assistant-stream";
 import type { Toolkit } from "./toolbox";
 
-export type McpToolkitDefinition = Record<string, McpServerConfig>;
+export type McpToolkitEntry =
+  | McpServerConfig
+  | {
+      server: McpServerConfig;
+      disabled?: boolean | undefined;
+      /**
+       * Prefix applied to every tool name exposed by this MCP server. Useful
+       * when multiple servers publish the same tool name, such as `search`.
+       */
+      prefix?: string | undefined;
+      tools?: Record<string, McpToolkitToolConfig> | undefined;
+    };
+
+export type McpToolkitToolConfig = {
+  disabled?: boolean | undefined;
+};
+
+export type McpToolkitDefinition = Record<string, McpToolkitEntry>;
 
 /**
- * Defines MCP server tools as a spreadable toolkit fragment.
+ * Defines MCP server tools as a spreadable toolkit fragment. Pass a raw
+ * `McpServerConfig` for always-on servers, or `{ server, disabled, tools }`
+ * when a server should stay configured but not expose all tools for the current
+ * request.
  */
 export function defineMcpToolkit(definition: McpToolkitDefinition): Toolkit {
   return Object.fromEntries(
-    Object.entries(definition).map(([name, server]) => [
-      name,
-      { type: "mcp", server },
-    ]),
+    Object.entries(definition).map(([name, entry]) => {
+      const server = "server" in entry ? entry.server : entry;
+      const disabled = "server" in entry ? entry.disabled : undefined;
+      const prefix = "server" in entry ? entry.prefix : undefined;
+      const tools = "server" in entry ? entry.tools : undefined;
+
+      return [
+        name,
+        {
+          type: "mcp",
+          server,
+          ...(disabled !== undefined && { disabled }),
+          ...(prefix !== undefined && { prefix }),
+          ...(tools !== undefined && { tools }),
+        },
+      ];
+    }),
   ) as Toolkit;
 }

package/src/react/model-context/define-toolkit.test.ts

@@ -1,5 +1,6 @@
 import { describe, it, expect, expectTypeOf } from "vitest";
 import type { AsyncIterableStream } from "assistant-stream/utils";
+import { defineMcpToolkit } from "./define-mcp-toolkit";
 import { defineToolkit } from "./define-toolkit";
 import { hitl, hitlTool, humanTool } from "./human-tool";
 import { providerTool } from "./provider-tool";
@@ -74,12 +75,118 @@ const checkToolkitDefinitionTypes = () => {
 };
 expectTypeOf(checkToolkitDefinitionTypes).toEqualTypeOf<() => void>();
 
+const checkDefineMcpToolkitTypes = () => {
+  defineMcpToolkit({
+    docs: {
+      type: "http",
+      url: "https://example.com/mcp",
+    },
+    prefixedDocs: {
+      server: {
+        type: "http",
+        url: "https://example.com/prefixed-mcp",
+      },
+      disabled: true,
+      prefix: "docs_",
+    },
+    gatedDocs: {
+      server: {
+        type: "http",
+        url: "https://example.com/gated-mcp",
+      },
+      disabled: true,
+      tools: {
+        privateSearch: {
+          disabled: true,
+        },
+      },
+    },
+  });
+};
+expectTypeOf(checkDefineMcpToolkitTypes).toEqualTypeOf<() => void>();
+
 describe("use-generative markers", () => {
   it("defineToolkit returns the toolkit at runtime", () => {
     const toolkit = {};
     expect(defineToolkit(toolkit)).toBe(toolkit);
   });
 
+  it("defineMcpToolkit supports prefixed and disabled MCP entries", () => {
+    expect(
+      defineMcpToolkit({
+        docs: {
+          server: { type: "http", url: "https://example.com/mcp" },
+          disabled: true,
+          prefix: "docs_",
+        },
+      }),
+    ).toEqual({
+      docs: {
+        type: "mcp",
+        server: { type: "http", url: "https://example.com/mcp" },
+        disabled: true,
+        prefix: "docs_",
+      },
+    });
+  });
+
+  it("defineMcpToolkit supports disabled server entries", () => {
+    expect(
+      defineMcpToolkit({
+        docs: {
+          type: "http",
+          url: "https://example.com/mcp",
+        },
+        gatedDocs: {
+          server: {
+            type: "http",
+            url: "https://example.com/gated-mcp",
+          },
+          disabled: true,
+          tools: {
+            privateSearch: {
+              disabled: true,
+            },
+          },
+        },
+      }),
+    ).toEqual({
+      docs: {
+        type: "mcp",
+        server: {
+          type: "http",
+          url: "https://example.com/mcp",
+        },
+      },
+      gatedDocs: {
+        type: "mcp",
+        server: {
+          type: "http",
+          url: "https://example.com/gated-mcp",
+        },
+        disabled: true,
+        tools: {
+          privateSearch: {
+            disabled: true,
+          },
+        },
+      },
+    });
+  });
+
+  it("defineMcpToolkit supports raw MCP server configs", () => {
+    expect(
+      defineMcpToolkit({
+        docs: { type: "http", url: "https://example.com/mcp" },
+      }),
+    ).toEqual({
+      docs: {
+        type: "mcp",
+        server: { type: "http", url: "https://example.com/mcp" },
+      },
+    });
+  });
+
   it("humanTool throws at runtime — it must be stripped by the compiler, never called", () => {
     expect(() => humanTool()).toThrow(/no runtime implementation/);
   });