@assistant-ui/core 0.2.2 → 0.2.4

package/src/react/types/MessagePartComponentTypes.ts

@@ -57,10 +57,21 @@ export type ToolCallMessagePartProps<
   TResult = unknown,
 > = MessagePartState &
   ToolCallMessagePart<TArgs, TResult> & {
+    /**
+     * Sets the result for this tool-call message part.
+     *
+     * Use when the renderer, rather than a tool `execute` function, is the
+     * source of the result.
+     */
     addResult: (result: TResult | ToolResponse<TResult>) => void;
+    /**
+     * Supplies the payload requested by `context.human(...)` and resumes the
+     * paused frontend tool execution.
+     */
     resume: (payload: unknown) => void;
   };
 
+/** Component used to render a tool-call message part. */
 export type ToolCallMessagePartComponent<
   TArgs = any,
   TResult = any,

package/src/runtimes/external-store/external-store-thread-runtime-core.ts

@@ -53,6 +53,7 @@ export class ExternalStoreThreadRuntimeCore
   implements ThreadRuntimeCore
 {
   private _assistantOptimisticId: string | null = null;
+  private _lastSyncedMessageIds = new Set<string>();
 
   private _capabilities: RuntimeCapabilities = {
     switchToBranch: false,
@@ -170,6 +171,7 @@ export class ExternalStoreThreadRuntimeCore
       // Clear and import the message repository
       this.repository.clear();
       this._assistantOptimisticId = null;
+      this._lastSyncedMessageIds = new Set();
       this.repository.import(store.messageRepository);
 
       messages = this.repository.getMessages();
@@ -222,6 +224,12 @@ export class ExternalStoreThreadRuntimeCore
             return newMessage;
           });
 
+      const nextIds = new Set(messages.map((m) => m.id));
+      for (const prevId of this._lastSyncedMessageIds) {
+        if (!nextIds.has(prevId)) this.repository.deleteMessage(prevId);
+      }
+      this._lastSyncedMessageIds = nextIds;
+
       for (let i = 0; i < messages.length; i++) {
         const message = messages[i]!;
         const parent = messages[i - 1];
@@ -336,6 +344,7 @@ export class ExternalStoreThreadRuntimeCore
       previousMessage.id === messages.at(-1)?.id // ensure the previous message is a leaf node
     ) {
       this.repository.deleteMessage(previousMessage.id);
+      this._lastSyncedMessageIds.delete(previousMessage.id);
       if (!this.composer.text.trim()) {
         this.composer.setText(getThreadMessageText(previousMessage));
       }
@@ -364,6 +373,7 @@ export class ExternalStoreThreadRuntimeCore
   }
 
   public override reset(initialMessages?: readonly ThreadMessageLike[]) {
+    this._lastSyncedMessageIds = new Set();
     const repo = new MessageRepository();
     repo.import(ExportedMessageRepository.fromArray(initialMessages ?? []));
     this.updateMessages(repo.getMessages());
@@ -371,6 +381,7 @@ export class ExternalStoreThreadRuntimeCore
 
   public override import(data: ExportedMessageRepository) {
     this._assistantOptimisticId = null;
+    this._lastSyncedMessageIds = new Set();
 
     super.import(data);
 

package/src/store/clients/model-context-client.test.ts

@@ -0,0 +1,108 @@
+import { describe, it, expect } from "vitest";
+import { createResourceRoot } from "@assistant-ui/tap";
+import type { Tool } from "assistant-stream";
+import { ModelContext } from "./model-context-client";
+import type {
+  ModelContext as ModelContextValue,
+  ModelContextProvider,
+} from "../../model-context/types";
+
+const tick = () => new Promise<void>((resolve) => setTimeout(resolve, 0));
+
+const provider = (ctx: ModelContextValue): ModelContextProvider => ({
+  getModelContext: () => ctx,
+});
+
+const stubTool = (): Tool<any, any> =>
+  ({ description: "", parameters: {} as any }) as unknown as Tool<any, any>;
+
+const render = () => {
+  const root = createResourceRoot();
+  const sub = root.render(ModelContext());
+  return { sub, unmount: () => root.unmount() };
+};
+
+describe("ModelContext", () => {
+  it("starts with undefined modelName and an empty toolNames array", () => {
+    const { sub, unmount } = render();
+    try {
+      const state = sub.getValue().getState();
+      expect(state.modelName).toBeUndefined();
+      expect(state.toolNames).toEqual([]);
+    } finally {
+      unmount();
+    }
+  });
+
+  it("reflects modelName from a registered provider", async () => {
+    const { sub, unmount } = render();
+    try {
+      sub.getValue().register(provider({ config: { modelName: "gpt-4" } }));
+      await tick();
+
+      expect(sub.getValue().getState().modelName).toBe("gpt-4");
+    } finally {
+      unmount();
+    }
+  });
+
+  it("reflects tool names from a registered provider", async () => {
+    const { sub, unmount } = render();
+    try {
+      sub
+        .getValue()
+        .register(provider({ tools: { foo: stubTool(), bar: stubTool() } }));
+      await tick();
+
+      expect(sub.getValue().getState().toolNames).toEqual(["bar", "foo"]);
+    } finally {
+      unmount();
+    }
+  });
+
+  it("keeps the same state reference when an extra provider does not change the merged values", async () => {
+    const { sub, unmount } = render();
+    try {
+      sub.getValue().register(provider({ config: { modelName: "gpt-4" } }));
+      await tick();
+      const before = sub.getValue().getState();
+
+      sub.getValue().register(provider({ config: { modelName: "gpt-4" } }));
+      await tick();
+      const after = sub.getValue().getState();
+
+      expect(after).toBe(before);
+    } finally {
+      unmount();
+    }
+  });
+
+  it("clears modelName when the last contributing provider unsubscribes", async () => {
+    const { sub, unmount } = render();
+    try {
+      const unsubscribe = sub
+        .getValue()
+        .register(provider({ config: { modelName: "gpt-4" } }));
+      await tick();
+      expect(sub.getValue().getState().modelName).toBe("gpt-4");
+
+      unsubscribe();
+      await tick();
+      expect(sub.getValue().getState().modelName).toBeUndefined();
+      expect(sub.getValue().getState().toolNames).toEqual([]);
+    } finally {
+      unmount();
+    }
+  });
+
+  it("reflects modelName synchronously after register without awaiting", () => {
+    const { sub, unmount } = render();
+    try {
+      sub.getValue().register(provider({ config: { modelName: "gpt-4" } }));
+
+      expect(sub.getValue().getState().modelName).toBe("gpt-4");
+    } finally {
+      unmount();
+    }
+  });
+});

package/src/store/clients/model-context-client.ts

@@ -1,18 +1,48 @@
-import { resource, tapMemo, tapState } from "@assistant-ui/tap";
+import { resource, tapEffect, tapMemo, tapState } from "@assistant-ui/tap";
 import type { ClientOutput } from "@assistant-ui/store";
 import { CompositeContextProvider } from "../../utils/composite-context-provider";
 import type { ModelContextState } from "../scopes/model-context";
 
-const version = 1;
+const EMPTY_TOOL_NAMES: readonly string[] = [];
+
+const INITIAL_STATE: ModelContextState = {
+  modelName: undefined,
+  toolNames: EMPTY_TOOL_NAMES,
+};
+
+const toolNamesEqual = (a: readonly string[], b: readonly string[]): boolean =>
+  a === b || (a.length === b.length && a.every((v, i) => v === b[i]));
+
+const deriveState = (
+  composite: CompositeContextProvider,
+  prev: ModelContextState,
+): ModelContextState => {
+  const ctx = composite.getModelContext();
+  const modelName = ctx.config?.modelName;
+  const keys = ctx.tools ? Object.keys(ctx.tools).sort() : EMPTY_TOOL_NAMES;
+  const toolNames = keys.length ? keys : EMPTY_TOOL_NAMES;
+
+  if (modelName === prev.modelName && toolNamesEqual(toolNames, prev.toolNames))
+    return prev;
+
+  return { modelName, toolNames };
+};
 
 export const ModelContext = resource((): ClientOutput<"modelContext"> => {
-  const [state] = tapState<ModelContextState>(
-    () => ({ version: version + 1 }) as unknown as ModelContextState,
-  );
   const composite = tapMemo(() => new CompositeContextProvider(), []);
+  const [state, setState] = tapState<ModelContextState>(() =>
+    deriveState(composite, INITIAL_STATE),
+  );
+
+  tapEffect(() => {
+    setState((prev) => deriveState(composite, prev));
+    return composite.subscribe(() => {
+      setState((prev) => deriveState(composite, prev));
+    });
+  }, [composite]);
 
   return {
-    getState: () => state,
+    getState: () => deriveState(composite, state),
     getModelContext: () => composite.getModelContext(),
     subscribe: (callback) => composite.subscribe(callback),
     register: (provider) => composite.registerModelContextProvider(provider),

package/src/store/scopes/model-context.ts

@@ -1,7 +1,10 @@
 import type { Unsubscribe } from "../../types/unsubscribe";
 import type { ModelContextProvider } from "../../model-context/types";
 
-export type ModelContextState = Record<string, never>;
+export type ModelContextState = {
+  readonly modelName?: string | undefined;
+  readonly toolNames: readonly string[];
+};
 
 export type ModelContextMethods = ModelContextProvider & {
   getState(): ModelContextState;

package/src/tests/external-store-thread-runtime-core.test.ts

@@ -150,3 +150,116 @@ describe("ExternalStoreThreadRuntimeCore - state reference stability", () => {
     expect(runtime.capabilities).toBe(capsBefore);
   });
 });
+
+describe("ExternalStoreThreadRuntimeCore - messages reconciliation", () => {
+  const user = { id: "u", role: "user" as const, content: [] };
+
+  it("drops ids that disappear between syncs (same length, swapped assistant id)", () => {
+    const a1 = { id: "a1", role: "assistant" as const, content: [] };
+    const a2 = { id: "a2", role: "assistant" as const, content: [] };
+
+    const runtime = new ExternalStoreThreadRuntimeCore(
+      mockContextProvider,
+      makeStore({ messages: [user, a1] }),
+    );
+
+    runtime.__internal_setAdapter(makeStore({ messages: [user, a2] }));
+
+    const exported = runtime.export();
+    expect(exported.messages.map((m) => m.message.id)).toEqual(["u", "a2"]);
+    const userChildren = exported.messages
+      .filter((m) => m.parentId === "u")
+      .map((m) => m.message.id);
+    expect(userChildren).toEqual(["a2"]);
+  });
+
+  it("keeps prior ids when they remain in the new sync", () => {
+    const a = { id: "a", role: "assistant" as const, content: [] };
+
+    const runtime = new ExternalStoreThreadRuntimeCore(
+      mockContextProvider,
+      makeStore({ messages: [user, a] }),
+    );
+
+    runtime.__internal_setAdapter(makeStore({ messages: [user, a] }));
+
+    expect(runtime.export().messages.map((m) => m.message.id)).toEqual([
+      "u",
+      "a",
+    ]);
+  });
+
+  it("removes trailing messages dropped from the new sync", () => {
+    const a = { id: "a", role: "assistant" as const, content: [] };
+    const u2 = { id: "u2", role: "user" as const, content: [] };
+
+    const runtime = new ExternalStoreThreadRuntimeCore(
+      mockContextProvider,
+      makeStore({ messages: [user, a, u2] }),
+    );
+
+    runtime.__internal_setAdapter(makeStore({ messages: [user, a] }));
+
+    expect(runtime.export().messages.map((m) => m.message.id)).toEqual([
+      "u",
+      "a",
+    ]);
+  });
+
+  it("does not crash on the next sync after cancelRun removes a leaf user", () => {
+    const userWithText = {
+      id: "u",
+      role: "user" as const,
+      content: [{ type: "text" as const, text: "hi" }],
+    };
+
+    const runtime = new ExternalStoreThreadRuntimeCore(
+      mockContextProvider,
+      makeStore({
+        messages: [userWithText],
+        onCancel: vi.fn(),
+        isRunning: true,
+      }),
+    );
+
+    runtime.cancelRun();
+
+    expect(() => {
+      runtime.__internal_setAdapter(makeStore({ messages: [] }));
+    }).not.toThrow();
+  });
+
+  it("drops phantom sibling when convertMessage swaps the assistant id", () => {
+    type Raw = { id: string; role: "user" | "assistant"; text: string };
+    const rawU: Raw = { id: "u", role: "user", text: "hi" };
+    const rawA1: Raw = { id: "client_id", role: "assistant", text: "" };
+    const rawA2: Raw = { id: "server_id", role: "assistant", text: "" };
+
+    const convertMessage = (m: Raw) => ({
+      id: m.id,
+      role: m.role,
+      content: [{ type: "text" as const, text: m.text }],
+    });
+
+    const runtime = new ExternalStoreThreadRuntimeCore(
+      mockContextProvider,
+      makeStore({
+        messages: [rawU, rawA1] as any,
+        convertMessage: convertMessage as any,
+      }),
+    );
+
+    runtime.__internal_setAdapter(
+      makeStore({
+        messages: [rawU, rawA2] as any,
+        convertMessage: convertMessage as any,
+      }),
+    );
+
+    const userChildren = runtime
+      .export()
+      .messages.filter((m) => m.parentId === "u")
+      .map((m) => m.message.id);
+    expect(userChildren).toEqual(["server_id"]);
+  });
+});

package/src/types/message.ts

@@ -92,18 +92,38 @@ export type ToolCallMessagePart<
   TArgs = ReadonlyJSONObject,
   TResult = unknown,
 > = {
+  /** Identifies this part as a tool call. */
   readonly type: "tool-call";
+  /** Stable identifier for this invocation of the tool. */
   readonly toolCallId: string;
+  /** Name of the tool requested by the model. */
   readonly toolName: string;
+  /**
+   * Arguments supplied by the model. During streaming this is a partial parse:
+   * fields may be missing or incomplete. From a tool-call renderer, use
+   * `useToolArgsStatus` to detect which fields are still arriving.
+   */
   readonly args: TArgs;
+  /** Result returned by the tool, if it has completed. */
   readonly result?: TResult | undefined;
+  /** Whether the result represents a tool execution error. */
   readonly isError?: boolean | undefined;
+  /** Raw JSON argument text streamed by the model. */
   readonly argsText: string;
+  /** UI-only artifact associated with the tool result. */
   readonly artifact?: unknown;
+  /** MCP app metadata associated with this tool call, when present. */
   readonly mcp?: ToolCallMessagePartMcpMetadata;
+  /** Content returned to the model for this tool result. */
   readonly modelContent?: readonly ToolModelContentPart[] | undefined;
+  /** Human-input request that must be resolved before the run can continue. */
   readonly interrupt?: { type: "human"; payload: unknown };
+  /** Parent message-part ID when this part belongs to a nested structure. */
   readonly parentId?: string;
+  /**
+   * Nested thread messages produced by this tool call, for example a sub-agent
+   * conversation.
+   */
   readonly messages?: readonly ThreadMessage[];
 };
 
@@ -143,7 +163,9 @@ export type MessagePartStatus =
 
 export type ToolCallMessagePartStatus =
   | {
+      /** The tool call is waiting for UI or human input before continuing. */
       readonly type: "requires-action";
+      /** Reason the tool call requires action. */
       readonly reason: "interrupt";
     }
   | MessagePartStatus;