@assistant-ui/core 0.2.20 → 0.2.21

package/src/model-context/tool.ts

@@ -1,6 +1,19 @@
 import type { Tool } from "assistant-stream";
 
-// TODO re-add the inferrence of the parameters
+type StandardSchemaInput<TSchema> = TSchema extends {
+  readonly "~standard": {
+    readonly types?: { readonly input: infer TInput } | undefined;
+  };
+}
+  ? TInput extends Record<string, unknown>
+    ? TInput
+    : Record<string, unknown>
+  : Record<string, unknown>;
+
+type StandardSchemaParameters = Extract<
+  NonNullable<Extract<Tool<any>, { parameters: unknown }>["parameters"]>,
+  { readonly "~standard": unknown }
+>;
 
 /**
  * Defines a model tool with its argument schema, execution behavior, and
@@ -8,25 +21,34 @@ import type { Tool } from "assistant-stream";
  *
  * This helper keeps reusable tool definitions type-checked and convenient to
  * export for a {@link Toolkit} registered with {@link Tools}.
- * Inference from parameter schemas is currently limited, so provide generic
- * arguments when you need precise args or result types.
+ * When `parameters` is a Standard Schema (for example Zod v4), its input type
+ * is inferred for `execute`, `streamCall`, and model-output callbacks. Plain
+ * JSON Schema objects do not carry TypeScript input types, so provide generic
+ * arguments or annotate callbacks when you need precise args for JSON Schema.
  *
  * @param tool - Tool definition to expose to the assistant model.
  *
  * @example
  * ```ts
- * const getWeather = tool<{ city: string }, string>({
+ * import { z } from "zod";
+ *
+ * const getWeather = tool({
  *   type: "frontend",
  *   description: "Get the weather for a city.",
- *   parameters: {
- *     type: "object",
- *     properties: { city: { type: "string" } },
- *     required: ["city"],
- *   },
+ *   parameters: z.object({ city: z.string() }),
  *   execute: async ({ city }) => `Sunny in ${city}`,
  * });
  * ```
  */
+export function tool<
+  const TSchema extends StandardSchemaParameters,
+  TResult = any,
+>(
+  tool: Tool<StandardSchemaInput<TSchema>, TResult> & { parameters: TSchema },
+): Tool<StandardSchemaInput<TSchema>, TResult>;
+export function tool<TArgs extends Record<string, unknown>, TResult = any>(
+  tool: Tool<TArgs, TResult>,
+): Tool<TArgs, TResult>;
 export function tool<TArgs extends Record<string, unknown>, TResult = any>(
   tool: Tool<TArgs, TResult>,
 ): Tool<TArgs, TResult> {

package/src/react/adapters/LocalStorageThreadListAdapter.test.ts

@@ -0,0 +1,222 @@
+import { describe, expect, it } from "vitest";
+import type { AsyncStorageLike } from "./LocalStorageThreadListAdapter";
+import {
+  createLocalStorageAdapter,
+  parseStoredMessageRepository,
+  parseStoredThreadMetadata,
+} from "./LocalStorageThreadListAdapter";
+
+const storedMessage = (
+  id: string,
+  role: "user" | "assistant" | "system" = "user",
+) => ({
+  id,
+  role,
+  createdAt: "2026-01-01T00:00:00.000Z",
+  content: [],
+  metadata: { custom: {} },
+  ...(role === "user" ? { attachments: [] } : undefined),
+  ...(role === "assistant"
+    ? { status: { type: "complete", reason: "stop" } }
+    : undefined),
+});
+
+const createStorage = (
+  entries: Record<string, string> = {},
+): AsyncStorageLike & { get(key: string): string | undefined } => {
+  const values = new Map(Object.entries(entries));
+  return {
+    get: (key) => values.get(key),
+    getItem: async (key) => values.get(key) ?? null,
+    setItem: async (key, value) => {
+      values.set(key, value);
+    },
+    removeItem: async (key) => {
+      values.delete(key);
+    },
+  };
+};
+
+describe("parseStoredThreadMetadata", () => {
+  it("returns an empty list for invalid JSON", () => {
+    expect(parseStoredThreadMetadata("{not-json")).toEqual([]);
+  });
+
+  it("skips malformed thread records while preserving valid records", () => {
+    const threads = parseStoredThreadMetadata(
+      JSON.stringify([
+        { remoteId: "thread-1", status: "regular", title: "Trip plan" },
+        { remoteId: 123, status: "regular" },
+        { remoteId: "thread-2", status: "archived", custom: { pinned: true } },
+        { remoteId: "thread-3", status: "deleted" },
+      ]),
+    );
+
+    expect(threads).toEqual([
+      { remoteId: "thread-1", status: "regular", title: "Trip plan" },
+      { remoteId: "thread-2", status: "archived", custom: { pinned: true } },
+    ]);
+  });
+
+  it("defaults old thread records without status to regular", () => {
+    expect(
+      parseStoredThreadMetadata(JSON.stringify([{ remoteId: "old" }])),
+    ).toEqual([{ remoteId: "old", status: "regular" }]);
+  });
+});
+
+describe("parseStoredMessageRepository", () => {
+  it("returns empty history for invalid JSON", () => {
+    expect(parseStoredMessageRepository("{not-json")).toEqual({ messages: [] });
+  });
+
+  it("skips malformed message records", () => {
+    const repo = parseStoredMessageRepository(
+      JSON.stringify({
+        headId: "message-2",
+        messages: [
+          {
+            message: storedMessage("message-1"),
+            parentId: null,
+          },
+          { message: { role: "user", content: [] }, parentId: null },
+          {
+            message: storedMessage("message-2", "assistant"),
+            parentId: "message-1",
+          },
+        ],
+      }),
+    );
+
+    expect(repo.messages.map((item) => item.message.id)).toEqual([
+      "message-1",
+      "message-2",
+    ]);
+    expect(repo.headId).toBe("message-2");
+  });
+
+  it("drops a head id that points at a skipped message", () => {
+    const repo = parseStoredMessageRepository(
+      JSON.stringify({
+        headId: "missing",
+        messages: [
+          {
+            message: storedMessage("message-1"),
+            parentId: null,
+          },
+        ],
+      }),
+    );
+
+    expect(repo.headId).toBeUndefined();
+    expect(repo.messages.map((item) => item.message.id)).toEqual(["message-1"]);
+  });
+
+  it("skips messages missing the required thread message shell", () => {
+    const repo = parseStoredMessageRepository(
+      JSON.stringify({
+        messages: [
+          { message: { id: "missing-role" }, parentId: null },
+          {
+            message: {
+              ...storedMessage("missing-content"),
+              content: undefined,
+            },
+            parentId: null,
+          },
+          {
+            message: { ...storedMessage("missing-metadata"), metadata: {} },
+            parentId: null,
+          },
+          {
+            message: storedMessage("valid"),
+            parentId: null,
+          },
+        ],
+      }),
+    );
+
+    expect(repo.messages.map((item) => item.message.id)).toEqual(["valid"]);
+  });
+
+  it("skips messages whose parent is missing, skipped, or appears later", () => {
+    const repo = parseStoredMessageRepository(
+      JSON.stringify({
+        headId: "child-of-root",
+        messages: [
+          {
+            message: storedMessage("child-of-missing"),
+            parentId: "missing",
+          },
+          {
+            message: storedMessage("child-of-invalid"),
+            parentId: "invalid",
+          },
+          {
+            message: { id: "invalid" },
+            parentId: null,
+          },
+          {
+            message: storedMessage("child-before-parent"),
+            parentId: "late-parent",
+          },
+          {
+            message: storedMessage("late-parent"),
+            parentId: null,
+          },
+          {
+            message: storedMessage("child-of-root"),
+            parentId: "late-parent",
+          },
+        ],
+      }),
+    );
+
+    expect(
+      repo.messages.map((item) => ({
+        id: item.message.id,
+        parentId: item.parentId,
+      })),
+    ).toEqual([
+      { id: "late-parent", parentId: null },
+      { id: "child-of-root", parentId: "late-parent" },
+    ]);
+    expect(repo.headId).toBe("child-of-root");
+  });
+});
+
+describe("createLocalStorageAdapter", () => {
+  it("lists no threads when the stored thread list is invalid JSON", async () => {
+    const storage = createStorage({ "@assistant-ui:threads": "{not-json" });
+    const adapter = createLocalStorageAdapter({ storage });
+
+    await expect(adapter.list()).resolves.toEqual({ threads: [] });
+  });
+
+  it("overwrites malformed thread storage when initializing a thread", async () => {
+    const storage = createStorage({ "@assistant-ui:threads": "{not-json" });
+    const adapter = createLocalStorageAdapter({ storage });
+
+    await expect(adapter.initialize("thread-1")).resolves.toEqual({
+      remoteId: "thread-1",
+      externalId: undefined,
+    });
+
+    expect(JSON.parse(storage.get("@assistant-ui:threads") ?? "")).toEqual([
+      { remoteId: "thread-1", status: "regular" },
+    ]);
+  });
+
+  it("includes the thread id when a stored thread cannot be fetched", async () => {
+    const storage = createStorage({
+      "@assistant-ui:threads": JSON.stringify([
+        { remoteId: "thread-1", status: "regular" },
+      ]),
+    });
+    const adapter = createLocalStorageAdapter({ storage });
+
+    await expect(adapter.fetch("missing-thread")).rejects.toThrow(
+      'Stored thread "missing-thread" not found while fetching thread metadata.',
+    );
+  });
+});

package/src/react/adapters/LocalStorageThreadListAdapter.tsx

@@ -8,6 +8,7 @@ import type {
   RemoteThreadMetadata,
   ThreadHistoryAdapter,
   ThreadMessage,
+  RunConfig,
 } from "../../index";
 import type {
   ExportedMessageRepository,
@@ -36,6 +37,212 @@ 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,
@@ -52,8 +259,7 @@ 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> {
@@ -61,9 +267,7 @@ class AsyncStorageHistoryAdapter implements ThreadHistoryAdapter {
 
     const key = this._messagesKey(remoteId);
     const raw = await this.storage.getItem(key);
-    const repo: ExportedMessageRepository = raw
-      ? (JSON.parse(raw) as ExportedMessageRepository)
-      : { messages: [] };
+    const repo = parseStoredMessageRepository(raw);
 
     const idx = repo.messages.findIndex(
       (m) => m.message.id === item.message.id,
@@ -110,7 +314,7 @@ export const createLocalStorageAdapter = (
 
   const loadThreadMetadata = async (): Promise<StoredThreadMetadata[]> => {
     const raw = await storage.getItem(threadsKey);
-    return raw ? (JSON.parse(raw) as StoredThreadMetadata[]) : [];
+    return parseStoredThreadMetadata(raw);
   };
 
   const saveThreadMetadata = async (
@@ -202,7 +406,10 @@ export const createLocalStorageAdapter = (
     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,

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