@assistant-ui/core 0.3.18 → 0.3.19

package/src/store/clients/thread-tasks.ts

@@ -0,0 +1,138 @@
+import { resource } from "@assistant-ui/tap";
+import type { ClientOutput } from "@assistant-ui/store";
+import type {
+  ThreadMessage,
+  ToolCallMessagePart,
+  ToolCallMessagePartStatus,
+} from "../../types/message";
+import {
+  COMPLETE_STATUS,
+  toMessagePartStatus,
+} from "../../utils/normalizePartStatus";
+import type { TaskState } from "../scopes/task";
+
+type TaskEntry = {
+  readonly task: TaskState;
+  readonly part: ToolCallMessagePart;
+  readonly statusType: ToolCallMessagePartStatus["type"];
+  readonly statusReason: string | undefined;
+  readonly statusError: unknown;
+  readonly messages: readonly ThreadMessage[];
+};
+
+const getStatusReason = (status: ToolCallMessagePartStatus) =>
+  "reason" in status ? status.reason : undefined;
+
+const getStatusError = (status: ToolCallMessagePartStatus) =>
+  "error" in status ? status.error : undefined;
+
+const MAX_TASK_DEPTH = 32;
+
+const taskKeys = new WeakMap<TaskState, string>();
+
+/** Lookup key for a task client: its document-order index path, unique by construction where ids from nested payloads are not. */
+export const getTaskKey = (task: TaskState) => taskKeys.get(task) ?? task.id;
+
+const resolveStatus = (
+  message: ThreadMessage,
+  partIndex: number,
+  part: ToolCallMessagePart,
+): ToolCallMessagePartStatus =>
+  "status" in message && message.status
+    ? toMessagePartStatus(message, partIndex, part)
+    : COMPLETE_STATUS;
+
+export const createTaskDeriver = () => {
+  let previous: readonly TaskState[] = [];
+  let previousEntries = new Map<string, TaskEntry>();
+
+  return (messages: readonly ThreadMessage[]): readonly TaskState[] => {
+    const tasks: TaskState[] = [];
+    const entries = new Map<string, TaskEntry>();
+    let allEntriesReused = true;
+
+    const visit = (
+      threadMessages: readonly ThreadMessage[],
+      parentTaskId: string | null,
+      depth: number,
+      path: string,
+    ) => {
+      if (depth > MAX_TASK_DEPTH) return;
+      for (const [messageIndex, message] of threadMessages.entries()) {
+        for (const [partIndex, part] of message.content.entries()) {
+          if (part.type !== "tool-call" || part.messages === undefined)
+            continue;
+
+          const nestedMessages = part.messages;
+          const status = resolveStatus(message, partIndex, part);
+          const statusReason = getStatusReason(status);
+          const statusError = getStatusError(status);
+          const entryKey = `${path}${messageIndex}.${partIndex}`;
+          const previousEntry = previousEntries.get(entryKey);
+          const task =
+            previousEntry?.part === part &&
+            previousEntry.statusType === status.type &&
+            previousEntry.statusReason === statusReason &&
+            Object.is(previousEntry.statusError, statusError) &&
+            previousEntry.messages === nestedMessages &&
+            previousEntry.task.messageId === message.id &&
+            previousEntry.task.parentTaskId === parentTaskId &&
+            previousEntry.task.depth === depth
+              ? previousEntry.task
+              : {
+                  id: part.toolCallId,
+                  toolName: part.toolName,
+                  args: part.args,
+                  result: part.result,
+                  ...(part.isError === undefined
+                    ? undefined
+                    : { isError: part.isError }),
+                  status,
+                  timing: part.timing,
+                  messageId: message.id,
+                  parentTaskId,
+                  depth,
+                  messages: nestedMessages,
+                };
+
+          if (task !== previousEntry?.task) {
+            allEntriesReused = false;
+            taskKeys.set(task, entryKey);
+          }
+          tasks.push(task);
+          entries.set(entryKey, {
+            task,
+            part,
+            statusType: status.type,
+            statusReason,
+            statusError,
+            messages: nestedMessages,
+          });
+          visit(nestedMessages, task.id, depth + 1, `${entryKey}.`);
+        }
+      }
+    };
+
+    visit(messages, null, 0, "");
+
+    const result =
+      allEntriesReused &&
+      tasks.length === previous.length &&
+      tasks.every((task, index) => task === previous[index])
+        ? previous
+        : tasks;
+    previous = result;
+    previousEntries = entries;
+    return result;
+  };
+};
+
+const useTaskClient = ({
+  task,
+}: {
+  task: TaskState;
+}): ClientOutput<"task"> => ({
+  getState: () => task,
+});
+
+export const TaskClient = resource(useTaskClient);

package/src/store/index.ts

@@ -48,6 +48,12 @@ export type {
   QueueItemMeta,
   QueueItemClientSchema,
 } from "./scopes/queue-item";
+export type {
+  TaskState,
+  TaskMethods,
+  TaskMeta,
+  TaskClientSchema,
+} from "./scopes/task";
 export type {
   AttachmentState,
   AttachmentMethods,
@@ -96,6 +102,8 @@ export {
 export {
   useExternalMessageConverter,
   convertExternalMessages,
+  createExternalMessageConversionCache,
+  type ExternalMessageConversionCache,
   type JoinStrategy,
 } from "../react/runtimes/external-message-converter";
 export {

package/src/store/primitive-predicates.test.ts

@@ -50,8 +50,12 @@ describe("primitive predicates", () => {
     ).toBe(true);
   });
 
-  it("actionBarReloadDisabled rejects user messages and busy threads", () => {
-    const thread = { isRunning: false, isDisabled: false };
+  it("actionBarReloadDisabled rejects user messages, busy threads, and runtimes without reload", () => {
+    const thread = {
+      isRunning: false,
+      isDisabled: false,
+      capabilities: { reload: true },
+    };
     expect(
       actionBarReloadDisabled(
         state({ thread, message: { role: "assistant" } }),
@@ -68,6 +72,14 @@ describe("primitive predicates", () => {
         }),
       ),
     ).toBe(true);
+    expect(
+      actionBarReloadDisabled(
+        state({
+          thread: { ...thread, capabilities: { reload: false } },
+          message: { role: "assistant" },
+        }),
+      ),
+    ).toBe(true);
   });
 
   it("actionBarCopyDisabled requires settled content with non-empty text", () => {
@@ -109,7 +121,7 @@ describe("primitive predicates", () => {
   it("branch picker predicates respect bounds and run capabilities", () => {
     const thread = {
       isRunning: false,
-      capabilities: { switchBranchDuringRun: false },
+      capabilities: { switchToBranch: true, switchBranchDuringRun: false },
     };
     const message = { branchNumber: 2, branchCount: 3 };
     expect(branchPickerPreviousDisabled(state({ thread, message }))).toBe(
@@ -131,6 +143,16 @@ describe("primitive predicates", () => {
         state({ thread: { ...thread, isRunning: true }, message }),
       ),
     ).toBe(true);
+    const noSwitching = {
+      ...thread,
+      capabilities: { ...thread.capabilities, switchToBranch: false },
+    };
+    expect(
+      branchPickerPreviousDisabled(state({ thread: noSwitching, message })),
+    ).toBe(true);
+    expect(
+      branchPickerNextDisabled(state({ thread: noSwitching, message })),
+    ).toBe(true);
   });
 
   it("suggestionTriggerDisabled gates on send only for queueless runs", () => {
@@ -174,11 +196,32 @@ describe("primitive predicates", () => {
       ),
     ).toBe(true);
 
+    const editable = { optional: { thread: { capabilities: { edit: true } } } };
     expect(
-      actionBarEditDisabled(state({ composer: { isEditing: false } })),
+      actionBarEditDisabled(
+        state({ ...editable, composer: { isEditing: false } }),
+      ),
     ).toBe(false);
     expect(
-      actionBarEditDisabled(state({ composer: { isEditing: true } })),
+      actionBarEditDisabled(
+        state({ ...editable, composer: { isEditing: true } }),
+      ),
+    ).toBe(true);
+    expect(
+      actionBarEditDisabled(
+        state({
+          optional: { thread: { capabilities: { edit: false } } },
+          composer: { isEditing: false },
+        }),
+      ),
     ).toBe(true);
+    expect(
+      actionBarEditDisabled(
+        state({
+          optional: { thread: undefined },
+          composer: { isEditing: false },
+        }),
+      ),
+    ).toBe(false);
   });
 });

package/src/store/primitive-predicates.ts

@@ -20,10 +20,13 @@ export const composerInputDisabled = (s: AssistantState): boolean =>
   s.thread.isDisabled || s.composer.dictation?.inputDisabled === true;
 
 export const actionBarEditDisabled = (s: AssistantState): boolean =>
-  s.composer.isEditing;
+  s.composer.isEditing || s.optional.thread?.capabilities.edit === false;
 
 export const actionBarReloadDisabled = (s: AssistantState): boolean =>
-  s.thread.isRunning || s.thread.isDisabled || s.message.role !== "assistant";
+  s.thread.isRunning ||
+  s.thread.isDisabled ||
+  s.message.role !== "assistant" ||
+  !s.thread.capabilities.reload;
 
 export const actionBarCopyDisabled = (s: AssistantState): boolean =>
   !(
@@ -33,10 +36,12 @@ export const actionBarCopyDisabled = (s: AssistantState): boolean =>
 
 export const branchPickerPreviousDisabled = (s: AssistantState): boolean =>
   s.message.branchNumber <= 1 ||
+  !s.thread.capabilities.switchToBranch ||
   (s.thread.isRunning && !s.thread.capabilities.switchBranchDuringRun);
 
 export const branchPickerNextDisabled = (s: AssistantState): boolean =>
   s.message.branchNumber >= s.message.branchCount ||
+  !s.thread.capabilities.switchToBranch ||
   (s.thread.isRunning && !s.thread.capabilities.switchBranchDuringRun);
 
 export const suggestionTriggerDisabled = (

package/src/store/runtime-clients/thread-runtime-client.ts

@@ -16,6 +16,11 @@ import {
 import { ComposerClient } from "./composer-runtime-client";
 import { MessageClient } from "./message-runtime-client";
 import { ThreadSuggestions } from "../clients/suggestions";
+import {
+  createTaskDeriver,
+  getTaskKey,
+  TaskClient,
+} from "../clients/thread-tasks";
 import { useSubscribable } from "./useSubscribable";
 import type { ThreadState } from "../scopes/thread";
 
@@ -70,6 +75,13 @@ const useThreadClient = ({
       unsubscribers.push(unsubscribe);
     }
 
+    unsubscribers.push(
+      runtime.unstable_on("toolApprovalAnswered", (payload) => {
+        const threadId = runtime.getState()?.threadId || "unknown";
+        emit("thread.toolApprovalAnswered", { threadId, ...payload });
+      }),
+    );
+
     return () => {
       for (const unsub of unsubscribers) unsub();
     };
@@ -102,6 +114,16 @@ const useThreadClient = ({
   const suggestions = useClientResource(
     ThreadSuggestions(runtimeState.suggestions),
   );
+  const taskDeriver = useMemo(() => createTaskDeriver(), []);
+  const tasks = useMemo(
+    () => taskDeriver(runtimeState.messages),
+    [taskDeriver, runtimeState.messages],
+  );
+  const taskClients = useClientLookup(
+    tasks.map((task) =>
+      withKey(getTaskKey(task), TaskClient({ task }), [task]),
+    ),
+  );
   const messages = useClientLookup(
     runtimeState.messages.map((m) =>
       withKey(
@@ -132,13 +154,21 @@ const useThreadClient = ({
 
       composer: composer.state,
       messages: messages.state,
+      tasks,
     };
-  }, [runtimeState, messages, composer.state]);
+  }, [runtimeState, messages, composer.state, tasks]);
 
   return {
     getState: () => state,
     composer: () => composer.methods,
     suggestions: () => suggestions.methods,
+    task: (selector) => {
+      if ("id" in selector) {
+        const task = tasks.find((candidate) => candidate.id === selector.id);
+        return taskClients.get({ key: task ? getTaskKey(task) : selector.id });
+      }
+      return taskClients.get(selector);
+    },
     append: (message) => {
       const appended: Exclude<CreateAppendMessage, string> =
         typeof message === "string"

package/src/store/scope-registration.ts

@@ -10,6 +10,7 @@ import type { SuggestionsClientSchema } from "./scopes/suggestions";
 import type { SuggestionClientSchema } from "./scopes/suggestion";
 import type { ChainOfThoughtClientSchema } from "./scopes/chain-of-thought";
 import type { QueueItemClientSchema } from "./scopes/queue-item";
+import type { TaskClientSchema } from "./scopes/task";
 
 declare module "@assistant-ui/store" {
   interface ScopeRegistry {
@@ -25,5 +26,6 @@ declare module "@assistant-ui/store" {
     suggestion: SuggestionClientSchema;
     chainOfThought: ChainOfThoughtClientSchema;
     queueItem: QueueItemClientSchema;
+    task: TaskClientSchema;
   }
 }

package/src/store/scopes/task.ts

@@ -0,0 +1,34 @@
+import type {
+  ThreadMessage,
+  ToolCallMessagePart,
+  ToolCallMessagePartStatus,
+} from "../../types/message";
+
+/** A tool call that carries a nested conversation, read as one unit of child work; `id` is its toolCallId. */
+export type TaskState = {
+  readonly id: string;
+  readonly toolName: string;
+  readonly args: ToolCallMessagePart["args"];
+  readonly result?: ToolCallMessagePart["result"];
+  readonly isError?: boolean;
+  readonly status: ToolCallMessagePartStatus;
+  readonly timing?: ToolCallMessagePart["timing"];
+  readonly messageId: string;
+  readonly parentTaskId: string | null;
+  readonly depth: number;
+  readonly messages: readonly ThreadMessage[];
+};
+
+export type TaskMethods = {
+  getState(): TaskState;
+};
+
+export type TaskMeta = {
+  source: "thread";
+  query: { type: "index"; index: number } | { type: "id"; id: string };
+};
+
+export type TaskClientSchema = {
+  methods: TaskMethods;
+  meta: TaskMeta;
+};

package/src/store/scopes/thread.ts

@@ -18,6 +18,7 @@ import type { ModelContext } from "../../model-context/types";
 import type { MessageMethods, MessageState } from "./message";
 import type { ComposerMethods, ComposerState } from "./composer";
 import type { SuggestionsMethods } from "./suggestions";
+import type { TaskMethods, TaskState } from "./task";
 
 export type ThreadState = {
   /**
@@ -44,6 +45,10 @@ export type ThreadState = {
    * The messages in the currently selected branch of the thread.
    */
   readonly messages: readonly MessageState[];
+  /**
+   * Child work derived from the thread's tool calls: every tool call that carries a nested conversation, in document order with nested tasks after their parent. The array keeps its identity while no task changed.
+   */
+  readonly tasks: readonly TaskState[];
   /**
    * The thread state.
    * @deprecated This feature is experimental
@@ -76,6 +81,10 @@ export type ThreadMethods = {
    * The suggestions shown for this thread.
    */
   suggestions(): SuggestionsMethods;
+  /**
+   * Access a task by index or toolCallId; an id resolves the first task with that toolCallId in document order.
+   */
+  task(selector: { index: number } | { id: string }): TaskMethods;
   /**
    * Append a new message to the thread.
    *
@@ -144,6 +153,13 @@ export type ThreadMeta = {
 };
 
 export type ThreadEvents = {
+  "thread.toolApprovalAnswered": {
+    threadId: string;
+    messageId: string;
+    toolCallId: string;
+    toolName: string;
+    approved: boolean;
+  };
   /**
    * A run started on this thread. Also observable as `isRunning` flipping to
    * `true` in thread state.

package/src/store/thread-tasks.test.tsx

@@ -0,0 +1,209 @@
+// @vitest-environment jsdom
+
+import { cleanup, render, waitFor } from "@testing-library/react";
+import { useAui, useAuiState } from "@assistant-ui/store";
+import { afterEach, describe, expect, it } from "vitest";
+import type { ThreadMessage } from "../types/message";
+import { AssistantRuntimeProvider } from "../react/AssistantRuntimeProvider";
+import { useExternalStoreRuntime } from "../react/runtimes/useExternalStoreRuntime";
+import type { TaskState } from "./scopes/task";
+
+const taskMessage = (
+  nestedStatus: { type: "running" } | { type: "complete"; reason: "stop" },
+  result?: unknown,
+) =>
+  ({
+    id: "message-1",
+    role: "assistant",
+    createdAt: new Date(0),
+    content: [
+      {
+        type: "tool-call",
+        toolCallId: "delegate-1",
+        toolName: "delegate",
+        args: { topic: "research" },
+        argsText: '{"topic":"research"}',
+        ...(result !== undefined ? { result } : undefined),
+        messages: [
+          {
+            id: "nested-message-1",
+            role: "assistant",
+            createdAt: new Date(0),
+            content: [],
+            status: nestedStatus,
+            metadata: {
+              unstable_state: {},
+              unstable_annotations: [],
+              unstable_data: [],
+              steps: [],
+              custom: {},
+            },
+          },
+        ],
+      },
+    ],
+    status: nestedStatus,
+    metadata: {
+      unstable_state: {},
+      unstable_annotations: [],
+      unstable_data: [],
+      steps: [],
+      custom: {},
+    },
+  }) as ThreadMessage;
+
+afterEach(cleanup);
+
+describe("thread tasks", () => {
+  it("exposes task state and accessors through an external-store runtime", async () => {
+    const captured: {
+      tasks?: readonly TaskState[];
+      byId?: TaskState;
+      byIndex?: TaskState;
+    } = {};
+
+    const Probe = () => {
+      const tasks = useAuiState((s) => s.thread.tasks);
+      const aui = useAui();
+      captured.tasks = tasks;
+      if (tasks.length > 0) {
+        captured.byId = aui.thread.task({ id: "delegate-1" }).getState();
+        captured.byIndex = aui.thread.task({ index: 0 }).getState();
+      }
+      return null;
+    };
+
+    const App = ({ messages }: { messages: readonly ThreadMessage[] }) => {
+      const runtime = useExternalStoreRuntime({
+        messages,
+        convertMessage: (message) => message,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <Probe />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    const view = render(<App messages={[taskMessage({ type: "running" })]} />);
+
+    await waitFor(() => expect(captured.tasks).toHaveLength(1));
+    expect(captured.tasks?.[0]).toMatchObject({
+      status: { type: "running" },
+      toolName: "delegate",
+    });
+    expect(captured.byId).toMatchObject({ messageId: "message-1" });
+    expect(captured.byIndex).toBe(captured.byId);
+
+    view.rerender(
+      <App
+        messages={[taskMessage({ type: "complete", reason: "stop" }, "done")]}
+      />,
+    );
+
+    await waitFor(() =>
+      expect(captured.tasks?.[0]?.status).toEqual({ type: "complete" }),
+    );
+    expect(captured.tasks?.[0]).toMatchObject({
+      toolName: "delegate",
+      messageId: "message-1",
+      result: "done",
+    });
+    expect(captured.tasks?.[0]?.messages).toHaveLength(1);
+  });
+
+  it("survives nested payloads that repeat ids and resolves the first task in document order", async () => {
+    const nested = (id: string, toolCallId: string) =>
+      ({
+        id,
+        role: "assistant",
+        createdAt: new Date(0),
+        content: [
+          {
+            type: "tool-call",
+            toolCallId,
+            toolName: "search",
+            args: {},
+            argsText: "{}",
+            result: "ok",
+            messages: [],
+          },
+        ],
+        status: { type: "complete", reason: "stop" },
+        metadata: {
+          unstable_state: {},
+          unstable_annotations: [],
+          unstable_data: [],
+          steps: [],
+          custom: {},
+        },
+      }) as ThreadMessage;
+    const message = {
+      id: "message-1",
+      role: "assistant",
+      createdAt: new Date(0),
+      content: [
+        {
+          type: "tool-call",
+          toolCallId: "delegate-a",
+          toolName: "delegate",
+          args: {},
+          argsText: "{}",
+          result: "ok",
+          messages: [nested("nested", "call_1")],
+        },
+        {
+          type: "tool-call",
+          toolCallId: "delegate-b",
+          toolName: "delegate",
+          args: {},
+          argsText: "{}",
+          result: "ok",
+          messages: [nested("nested", "call_1")],
+        },
+      ],
+      status: { type: "complete", reason: "stop" },
+      metadata: {
+        unstable_state: {},
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+    } as ThreadMessage;
+    const captured: { ids?: readonly string[]; first?: TaskState } = {};
+    const Probe = () => {
+      const tasks = useAuiState((s) => s.thread.tasks);
+      const aui = useAui();
+      captured.ids = tasks.map((task) => task.id);
+      if (tasks.length === 4) {
+        captured.first = aui.thread.task({ id: "call_1" }).getState();
+      }
+      return null;
+    };
+    const App = () => {
+      const runtime = useExternalStoreRuntime({
+        messages: [message],
+        convertMessage: (m) => m,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <Probe />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    render(<App />);
+
+    await waitFor(() => expect(captured.ids).toHaveLength(4));
+    expect(captured.ids).toEqual([
+      "delegate-a",
+      "call_1",
+      "delegate-b",
+      "call_1",
+    ]);
+    expect(captured.first).toMatchObject({ parentTaskId: "delegate-a" });
+  });
+});