@assistant-ui/core 0.2.20 → 0.2.21

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

package/src/react/primitive-hooks/useMessageError.test.ts

@@ -0,0 +1,64 @@
+import { describe, it, expect, vi } from "vitest";
+
+const { mockUseAuiState } = vi.hoisted(() => ({ mockUseAuiState: vi.fn() }));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAuiState: ((selector: (s: unknown) => unknown) =>
+    mockUseAuiState(
+      selector,
+    )) as typeof import("@assistant-ui/store").useAuiState,
+}));
+
+import { useMessageError } from "./useMessageError";
+
+const against = (message: unknown) =>
+  mockUseAuiState.mockImplementationOnce((selector: (s: unknown) => unknown) =>
+    selector({ message }),
+  );
+
+describe("useMessageError", () => {
+  it("returns the message from a structured AssistantError", () => {
+    against({
+      status: {
+        type: "incomplete",
+        reason: "error",
+        error: { code: "network", message: "offline" },
+      },
+    });
+    expect(useMessageError()).toBe("offline");
+  });
+
+  it("passes a legacy string error through", () => {
+    against({
+      status: {
+        type: "incomplete",
+        reason: "error",
+        error: "legacy failure",
+      },
+    });
+    expect(useMessageError()).toBe("legacy failure");
+  });
+
+  it("returns message from an object that fails the full guard", () => {
+    against({
+      status: {
+        type: "incomplete",
+        reason: "error",
+        error: {
+          code: "unknown",
+          message: "foreign severity stream",
+          severity: "fatal",
+        },
+      },
+    });
+    expect(useMessageError()).toBe("foreign severity stream");
+  });
+
+  it("returns undefined when there is no error status", () => {
+    against({
+      status: { type: "complete", reason: "stop" },
+    });
+    expect(useMessageError()).toBeUndefined();
+  });
+});

package/src/react/primitive-hooks/useMessageError.ts

@@ -1,10 +1,23 @@
 import { useAuiState } from "@assistant-ui/store";
 
 export const useMessageError = () => {
-  return useAuiState((s) =>
-    s.message.status?.type === "incomplete" &&
-    s.message.status.reason === "error"
-      ? (s.message.status.error ?? "An error occurred")
-      : undefined,
-  );
+  return useAuiState((s) => {
+    if (
+      s.message.status?.type !== "incomplete" ||
+      s.message.status.reason !== "error"
+    ) {
+      return undefined;
+    }
+    const error = s.message.status.error;
+    if (typeof error === "string") return error;
+    if (
+      typeof error === "object" &&
+      error !== null &&
+      "message" in error &&
+      typeof error.message === "string"
+    ) {
+      return error.message;
+    }
+    return error ?? "An error occurred";
+  });
 };

package/src/react/runtimes/RemoteThreadListHookInstanceManager.tsx

@@ -109,7 +109,9 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     const updateRuntime = useCallback(() => {
       const aliveThread = this.instances.get(threadId);
       if (!aliveThread)
-        throw new Error("Thread not found. This is a bug in assistant-ui.");
+        throw new Error(
+          `Thread "${threadId}" runtime binding not found. This is a bug in assistant-ui.`,
+        );
 
       aliveThread.runtime = threadBinding.getState();
       this._notifySubscribers();

package/src/react/runtimes/RemoteThreadListThreadListRuntimeCore.tsx

@@ -22,6 +22,18 @@ import { AssistantMessageStream } from "assistant-stream";
 import type { ModelContextProvider } from "../../model-context/types";
 import { RuntimeAdapterProvider } from "./RuntimeAdapterProvider";
 
+const threadNotFoundError = (threadIdOrRemoteId: string, action: string) =>
+  new Error(`Thread "${threadIdOrRemoteId}" not found while ${action}.`);
+
+const threadStatusError = (
+  threadIdOrRemoteId: string,
+  status: RemoteThreadData["status"],
+  action: string,
+) =>
+  new Error(
+    `Thread "${threadIdOrRemoteId}" has status "${status}", so it cannot ${action}.`,
+  );
+
 export class RemoteThreadListThreadListRuntimeCore
   extends BaseSubscribable
   implements ThreadListRuntimeCore
@@ -280,10 +292,14 @@ export class RemoteThreadListThreadListRuntimeCore
 
   public getThreadRuntimeCore(threadIdOrRemoteId: string) {
     const data = this.getItemById(threadIdOrRemoteId);
-    if (!data) throw new Error("Thread not found");
+    if (!data)
+      throw threadNotFoundError(threadIdOrRemoteId, "getting its runtime");
 
     const result = this._hookManager.getThreadRuntimeCore(data.id);
-    if (!result) throw new Error("Thread not found");
+    if (!result)
+      throw new Error(
+        `Runtime for thread "${threadIdOrRemoteId}" not found while getting its runtime.`,
+      );
     return result;
   }
 
@@ -354,7 +370,7 @@ export class RemoteThreadListThreadListRuntimeCore
       data = this.getItemById(threadIdOrRemoteId);
     }
 
-    if (!data) throw new Error("Thread not found");
+    if (!data) throw threadNotFoundError(threadIdOrRemoteId, "switching to it");
     if (this._mainThreadId === data.id) return;
 
     const task = this._hookManager.startThreadRuntime(data.id);
@@ -417,8 +433,9 @@ export class RemoteThreadListThreadListRuntimeCore
   public initialize = async (threadId: string) => {
     if (this._state.value.newThreadId !== threadId) {
       const data = this.getItemById(threadId);
-      if (!data) throw new Error("Thread not found");
-      if (data.status === "new") throw new Error("Unexpected new state");
+      if (!data) throw threadNotFoundError(threadId, "initializing it");
+      if (data.status === "new")
+        throw threadStatusError(threadId, data.status, "be initialized here");
       return data.initializeTask;
     }
 
@@ -469,8 +486,9 @@ export class RemoteThreadListThreadListRuntimeCore
 
   public generateTitle = async (threadId: string) => {
     const data = this.getItemById(threadId);
-    if (!data) throw new Error("Thread not found");
-    if (data.status === "new") throw new Error("Thread is not yet initialized");
+    if (!data) throw threadNotFoundError(threadId, "generating its title");
+    if (data.status === "new")
+      throw threadStatusError(threadId, data.status, "generate a title");
 
     const { remoteId } = await data.initializeTask;
 
@@ -503,8 +521,9 @@ export class RemoteThreadListThreadListRuntimeCore
 
   public rename(threadIdOrRemoteId: string, newTitle: string): Promise<void> {
     const data = this.getItemById(threadIdOrRemoteId);
-    if (!data) throw new Error("Thread not found");
-    if (data.status === "new") throw new Error("Thread is not yet initialized");
+    if (!data) throw threadNotFoundError(threadIdOrRemoteId, "renaming it");
+    if (data.status === "new")
+      throw threadStatusError(threadIdOrRemoteId, data.status, "be renamed");
 
     return this._state.optimisticUpdate({
       execute: async () => {
@@ -534,8 +553,17 @@ export class RemoteThreadListThreadListRuntimeCore
     custom: Record<string, unknown> | undefined,
   ): Promise<void> {
     const data = this.getItemById(threadIdOrRemoteId);
-    if (!data) throw new Error("Thread not found");
-    if (data.status === "new") throw new Error("Thread is not yet initialized");
+    if (!data)
+      throw threadNotFoundError(
+        threadIdOrRemoteId,
+        "updating its custom metadata",
+      );
+    if (data.status === "new")
+      throw threadStatusError(
+        threadIdOrRemoteId,
+        data.status,
+        "update custom metadata",
+      );
 
     if (!this._options.adapter.updateCustom) {
       throw new Error(
@@ -583,9 +611,9 @@ export class RemoteThreadListThreadListRuntimeCore
 
   public async archive(threadIdOrRemoteId: string) {
     const data = this.getItemById(threadIdOrRemoteId);
-    if (!data) throw new Error("Thread not found");
+    if (!data) throw threadNotFoundError(threadIdOrRemoteId, "archiving it");
     if (data.status !== "regular")
-      throw new Error("Thread is not yet initialized or already archived");
+      throw threadStatusError(threadIdOrRemoteId, data.status, "be archived");
 
     await this._ensureThreadIsNotMain(data.id);
 
@@ -602,8 +630,9 @@ export class RemoteThreadListThreadListRuntimeCore
 
   public unarchive(threadIdOrRemoteId: string): Promise<void> {
     const data = this.getItemById(threadIdOrRemoteId);
-    if (!data) throw new Error("Thread not found");
-    if (data.status !== "archived") throw new Error("Thread is not archived");
+    if (!data) throw threadNotFoundError(threadIdOrRemoteId, "unarchiving it");
+    if (data.status !== "archived")
+      throw threadStatusError(threadIdOrRemoteId, data.status, "be unarchived");
 
     return this._state.optimisticUpdate({
       execute: async () => {
@@ -623,9 +652,9 @@ export class RemoteThreadListThreadListRuntimeCore
 
   public async delete(threadIdOrRemoteId: string) {
     const data = this.getItemById(threadIdOrRemoteId);
-    if (!data) throw new Error("Thread not found");
+    if (!data) throw threadNotFoundError(threadIdOrRemoteId, "deleting it");
     if (data.status !== "regular" && data.status !== "archived")
-      throw new Error("Thread is not yet initialized");
+      throw threadStatusError(threadIdOrRemoteId, data.status, "be deleted");
 
     await this._ensureThreadIsNotMain(data.id);
     this._hookManager.stopThreadRuntime(data.id);
@@ -643,9 +672,9 @@ export class RemoteThreadListThreadListRuntimeCore
 
   public async detach(threadIdOrRemoteId: string): Promise<void> {
     const data = this.getItemById(threadIdOrRemoteId);
-    if (!data) throw new Error("Thread not found");
+    if (!data) throw threadNotFoundError(threadIdOrRemoteId, "detaching it");
     if (data.status !== "regular" && data.status !== "archived")
-      throw new Error("Thread is not yet initialized");
+      throw threadStatusError(threadIdOrRemoteId, data.status, "be detached");
 
     await this._ensureThreadIsNotMain(data.id);
     this._hookManager.stopThreadRuntime(data.id);

package/src/react/runtimes/cloud/CloudFileAttachmentAdapter.test.ts

@@ -0,0 +1,111 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type { AssistantCloud } from "assistant-cloud";
+import type { PendingAttachment } from "../../../types/attachment";
+import { CloudFileAttachmentAdapter } from "./CloudFileAttachmentAdapter";
+
+const makeCloud = () =>
+  ({
+    files: {
+      generatePresignedUploadUrl: vi.fn().mockResolvedValue({
+        signedUrl: "https://storage.example/upload",
+        publicUrl: "https://cdn.example/file.png",
+      }),
+    },
+  }) as unknown as AssistantCloud;
+
+const makeFile = () =>
+  new File([new Uint8Array([1, 2, 3])], "pixel.png", { type: "image/png" });
+
+const drain = async (adapter: CloudFileAttachmentAdapter) => {
+  const yields: PendingAttachment[] = [];
+  for await (const value of adapter.add({ file: makeFile() })) {
+    yields.push(value);
+  }
+  return yields;
+};
+
+describe("CloudFileAttachmentAdapter", () => {
+  afterEach(() => {
+    vi.unstubAllGlobals();
+    vi.restoreAllMocks();
+  });
+
+  it("marks the attachment ready when the upload succeeds", async () => {
+    vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: true }));
+    const adapter = new CloudFileAttachmentAdapter(makeCloud());
+
+    const yields = await drain(adapter);
+
+    expect(yields.at(0)?.status).toEqual({
+      type: "running",
+      reason: "uploading",
+      progress: 0,
+    });
+    expect(yields.at(-1)?.status).toEqual({
+      type: "requires-action",
+      reason: "composer-send",
+    });
+
+    const complete = await adapter.send(yields.at(-1)!);
+    expect(complete.status).toEqual({ type: "complete" });
+    expect(complete.content).toEqual([
+      {
+        type: "image",
+        image: "https://cdn.example/file.png",
+        filename: "pixel.png",
+      },
+    ]);
+  });
+
+  it("marks the attachment incomplete when the upload returns an HTTP error", async () => {
+    const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
+    vi.stubGlobal(
+      "fetch",
+      vi.fn().mockResolvedValue({
+        ok: false,
+        status: 403,
+        statusText: "Forbidden",
+      }),
+    );
+    const adapter = new CloudFileAttachmentAdapter(makeCloud());
+
+    const yields = await drain(adapter);
+
+    expect(yields.at(-1)?.status).toEqual({
+      type: "incomplete",
+      reason: "error",
+      message: "Failed to upload file: 403 Forbidden",
+    });
+    expect(errorSpy).toHaveBeenCalledWith(
+      "[assistant-ui] Failed to upload attachment:",
+      expect.objectContaining({
+        message: "Failed to upload file: 403 Forbidden",
+      }),
+    );
+    await expect(adapter.send(yields.at(-1)!)).rejects.toThrow(
+      "Attachment not uploaded",
+    );
+  });
+
+  it("marks the attachment incomplete when the upload throws", async () => {
+    const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
+    const uploadError = new Error("network down");
+    vi.stubGlobal("fetch", vi.fn().mockRejectedValue(uploadError));
+    const adapter = new CloudFileAttachmentAdapter(makeCloud());
+
+    const yields = await drain(adapter);
+
+    expect(yields.at(-1)?.status).toEqual({
+      type: "incomplete",
+      reason: "error",
+      message: "network down",
+    });
+    expect(errorSpy).toHaveBeenCalledWith(
+      "[assistant-ui] Failed to upload attachment:",
+      uploadError,
+    );
+    await expect(adapter.send(yields.at(-1)!)).rejects.toThrow(
+      "Attachment not uploaded",
+    );
+  });
+});

package/src/react/runtimes/cloud/CloudFileAttachmentAdapter.ts

@@ -44,7 +44,7 @@ export class CloudFileAttachmentAdapter implements AttachmentAdapter {
         await this.cloud.files.generatePresignedUploadUrl({
           filename: file.name,
         });
-      await fetch(signedUrl, {
+      const res = await fetch(signedUrl, {
         method: "PUT",
         body: file,
         headers: {
@@ -52,16 +52,26 @@ export class CloudFileAttachmentAdapter implements AttachmentAdapter {
         },
         mode: "cors",
       });
+      if (!res.ok) {
+        throw new Error(
+          `Failed to upload file: ${res.status} ${res.statusText}`,
+        );
+      }
       this.uploadedUrls.set(id, publicUrl);
       attachment = {
         ...attachment,
         status: { type: "requires-action", reason: "composer-send" },
       };
       yield attachment;
-    } catch {
+    } catch (error) {
+      console.error("[assistant-ui] Failed to upload attachment:", error);
       attachment = {
         ...attachment,
-        status: { type: "incomplete", reason: "error" },
+        status: {
+          type: "incomplete",
+          reason: "error",
+          message: error instanceof Error ? error.message : String(error),
+        },
       };
       yield attachment;
     }

package/src/react/runtimes/cloud/auiV0.ts

@@ -1,8 +1,14 @@
 import type {
+  DataMessagePart,
+  FileMessagePart,
+  ImageMessagePart,
   MessageStatus,
   SourceProviderMetadata,
   ThreadMessage,
+  TextMessagePart,
+  Unstable_AudioMessagePart,
 } from "../../../types/message";
+import type { CompleteAttachment } from "../../../types/attachment";
 import { fromThreadMessageLike } from "../../../runtime/utils/thread-message-like";
 import type { CloudMessage } from "assistant-cloud";
 import { isJSONValue } from "../../../utils/json/is-json";
@@ -65,10 +71,27 @@ type AuiV0MessagePart =
       readonly filename?: string;
     };
 
+type AuiV0AttachmentPart =
+  | TextMessagePart
+  | ImageMessagePart
+  | FileMessagePart
+  | Unstable_AudioMessagePart
+  | DataMessagePart<ReadonlyJSONValue>;
+
+type AuiV0Attachment = {
+  readonly id: string;
+  readonly type: CompleteAttachment["type"];
+  readonly name: string;
+  readonly contentType?: string;
+  readonly status: CompleteAttachment["status"];
+  readonly content: readonly AuiV0AttachmentPart[];
+};
+
 type AuiV0Message = {
   readonly role: "assistant" | "user" | "system";
   readonly status?: MessageStatus;
   readonly content: readonly AuiV0MessagePart[];
+  readonly attachments?: readonly AuiV0Attachment[];
   readonly metadata: {
     readonly unstable_state?: ReadonlyJSONValue;
     readonly unstable_annotations: readonly ReadonlyJSONValue[];
@@ -83,14 +106,59 @@ type AuiV0Message = {
   };
 };
 
+const encodeAttachmentPart = (
+  part: CompleteAttachment["content"][number],
+): AuiV0AttachmentPart => {
+  const type = part.type;
+  switch (type) {
+    case "text":
+    case "image":
+    case "file":
+    case "audio":
+      return part;
+
+    case "data": {
+      if (!isJSONValue(part.data)) {
+        console.warn(`attachment data is not JSON! ${JSON.stringify(part)}`);
+      }
+      return { ...part, data: part.data as ReadonlyJSONValue };
+    }
+
+    default: {
+      const unhandledType: never = type;
+      throw new Error(
+        `Attachment part type not supported by aui/v0: ${unhandledType}`,
+      );
+    }
+  }
+};
+
+const encodeAttachments = (
+  message: ThreadMessage,
+): readonly AuiV0Attachment[] | undefined => {
+  if (message.role !== "user" || message.attachments.length === 0) {
+    return undefined;
+  }
+
+  return message.attachments.map(
+    ({ id, type, name, contentType, status, content }) => ({
+      id,
+      type,
+      name,
+      status,
+      ...(contentType != null ? { contentType } : undefined),
+      content: content.map(encodeAttachmentPart),
+    }),
+  );
+};
+
 export function auiV0Encode(message: ThreadMessage): AuiV0Message {
-  // TODO attachments are currently intentionally ignored
   // info: ID and createdAt are ignored (we use the server value instead)
-
   const status: MessageStatus | undefined =
     message.status?.type === "running"
       ? { type: "incomplete", reason: "cancelled" }
       : message.status;
+  const attachments = encodeAttachments(message);
 
   return {
     role: message.role,
@@ -172,6 +240,7 @@ export function auiV0Encode(message: ThreadMessage): AuiV0Message {
     }),
     metadata: message.metadata as AuiV0Message["metadata"],
     ...(status ? { status } : undefined),
+    ...(attachments ? { attachments } : undefined),
   };
 }