fix assistant reasoning mount lifecycle

apps/docs/docs/management/assistant.mdx

@@ -155,11 +155,12 @@ another message starts. A compact dot matrix distinguishes an initial thinking p
 settles into waiting, success, or error states in the minimized activity bar. The state always keeps its text label,
 and the animation is disabled when the operating system requests reduced motion.
 
-A compact context control beside the model selector gives a quick glance at cumulative conversation tokens and
-cost. Its popover keeps the latest complete context reading while a new response starts and lists every request
+A circular context ring beside the model selector gives a quick glance at cumulative conversation usage without
+expanding the composer toolbar. Its accessible label includes cumulative tokens and cost when available, while
+its popover keeps the latest complete context reading as a new response starts and lists every request
 turn with its model, provider, input, output, cached-input and reasoning tokens, context use, duration, speed,
 agent steps, routing details, and cost when reported. Before usage is available, the control remains a compact
-circle instead of showing a redundant missing-data label. Closing and reopening the popover reads persisted
+ring instead of showing a redundant missing-data label. Closing and reopening the popover reads persisted
 conversation metadata instead of resetting the display to zero.
 
 Each assistant response places a small provider logo in the same stable action toolbar as copy, export,

apps/nextjs/src/components/assistant/assistant-panel.module.css

@@ -1153,14 +1153,13 @@
 
 .composerContext {
   display: flex;
-  min-width: 0;
-  max-width: 14rem;
+  width: 1.75rem;
   height: 1.75rem;
   align-items: center;
-  gap: 0.25rem;
-  padding-inline: 0.25rem 0.45rem;
+  justify-content: center;
+  padding: 0;
   border: 1px solid var(--mantine-color-default-border);
-  border-radius: var(--mantine-radius-sm);
+  border-radius: 50%;
   color: var(--mantine-color-text);
   background: var(--mantine-color-body);
 }
@@ -1169,13 +1168,6 @@
   background: var(--mantine-color-default-hover);
 }
 
-.composerContext[data-compact] {
-  width: 1.75rem;
-  justify-content: center;
-  padding: 0;
-  border-radius: 50%;
-}
-
 .composerContext:focus-visible {
   outline: 2px solid var(--mantine-primary-color-filled);
   outline-offset: 2px;

apps/nextjs/src/components/assistant/assistant-panel.tsx

@@ -133,6 +133,7 @@ import type { AssistantConversationTurnUsage } from "./assistant-message-metadat
 import type { AssistantPendingAction } from "./assistant-pending-action";
 import { AssistantQuestionPortalProvider } from "./assistant-question-portal";
 import type { AssistantReasoningMode, AssistantRuntimeModelOption } from "./assistant-preferences";
+import { useAssistantReasoningState } from "./assistant-reasoning-state";
 import { getNearestTriggerScrollTop } from "./assistant-trigger-scroll";
 import { getToolResultPresentation } from "./assistant-tool-result";
 import { getSafeAssistantHttpUrl } from "./assistant-url";
@@ -961,20 +962,9 @@ const ChainOfThoughtLayout = ({ children }: { children?: ReactNode }) => (
 const AssistantChainOfThought = () => {
   const t = useScopedI18n("common.assistant");
   const contentId = useId();
-  const aui = useAui();
-  const chainStatus = useAuiState((state) => state.chainOfThought.status);
-  const collapsed = useAuiState((state) => state.chainOfThought.collapsed);
-  const running = chainStatus.type === "running";
   const { collapsed: preferredCollapsed, setCollapsed: setPreferredCollapsed } = useContext(ReasoningVisibilityContext);
-
-  useLayoutEffect(() => {
-    aui.chainOfThought().setCollapsed(preferredCollapsed);
-  }, [aui, preferredCollapsed]);
-
-  useEffect(() => {
-    if (chainStatus.type !== "requires-action") return;
-    aui.chainOfThought().setCollapsed(false);
-  }, [aui, chainStatus.type]);
+  const { chainStatus, collapsed } = useAssistantReasoningState(preferredCollapsed);
+  const running = chainStatus.type === "running";
 
   return (
     <ChainOfThoughtPrimitive.Root className={classes.reasoning} data-opened={!collapsed || undefined}>
@@ -1568,9 +1558,9 @@ const ConversationContext = () => {
           className={classes.composerContext}
           type="button"
           aria-label={quickLabel ? `${t("usage.contextWindow")}: ${quickLabel}` : t("usage.contextWindow")}
+          title={quickLabel || t("usage.contextWindow")}
           aria-expanded={opened}
           aria-haspopup="dialog"
-          data-compact={!quickLabel || undefined}
           onClick={() => setOpened((value) => !value)}
         >
           <RingProgress
@@ -1579,11 +1569,6 @@ const ConversationContext = () => {
             roundCaps
             sections={hasContext ? [{ value: contextPercentage, color: getContextColor(contextPercentage) }] : []}
           />
-          {quickLabel && (
-            <Text size="xs" fw={650} truncate>
-              {quickLabel}
-            </Text>
-          )}
         </UnstyledButton>
       </Popover.Target>
       <Popover.Dropdown className={classes.conversationContextPopover}>

apps/nextjs/src/components/assistant/assistant-reasoning-state.spec.tsx

@@ -0,0 +1,144 @@
+// @vitest-environment jsdom
+
+import { act, createElement, useEffect } from "react";
+import { createRoot } from "react-dom/client";
+import type { AssistantClient, ChatModelAdapter, ThreadMessageLike } from "@assistant-ui/react";
+import {
+  AssistantRuntimeProvider,
+  ChainOfThoughtPrimitive,
+  MessagePrimitive,
+  ThreadPrimitive,
+  useAui,
+  useLocalRuntime,
+} from "@assistant-ui/react";
+import { afterEach, describe, expect, test } from "vitest";
+
+import { useAssistantReasoningState } from "./assistant-reasoning-state";
+
+(globalThis as typeof globalThis & { IS_REACT_ACT_ENVIRONMENT: boolean }).IS_REACT_ACT_ENVIRONMENT = true;
+
+const flushEffects = () => new Promise((resolve) => setTimeout(resolve, 0));
+const containers: HTMLDivElement[] = [];
+
+afterEach(() => {
+  for (const container of containers.splice(0)) container.remove();
+});
+
+const ReasoningPart = ({ text }: { text: string }) => <span data-testid="reasoning-part">{text}</span>;
+const ToolPart = ({ toolName }: { toolName: string }) => <span data-testid="tool-part">{toolName}</span>;
+
+const TestChainOfThought = ({ preferredCollapsed = true }: { preferredCollapsed?: boolean }) => {
+  const { chainStatus, collapsed } = useAssistantReasoningState(preferredCollapsed);
+
+  return (
+    <ChainOfThoughtPrimitive.Root
+      data-testid="chain-of-thought"
+      data-status={chainStatus.type}
+      data-collapsed={String(collapsed)}
+    >
+      <ChainOfThoughtPrimitive.Parts components={{ Reasoning: ReasoningPart, tools: { Fallback: ToolPart } }} />
+    </ChainOfThoughtPrimitive.Root>
+  );
+};
+
+const TestMessage = () => (
+  <MessagePrimitive.Root>
+    <MessagePrimitive.Parts components={{ ChainOfThought: TestChainOfThought }} />
+  </MessagePrimitive.Root>
+);
+
+const RuntimeProbe = ({ onReady }: { onReady: (aui: AssistantClient) => void }) => {
+  const aui = useAui();
+
+  useEffect(() => {
+    onReady(aui);
+  }, [aui, onReady]);
+  return null;
+};
+
+const renderAssistant = async (adapter: ChatModelAdapter, initialMessages: ThreadMessageLike[] = []) => {
+  const container = document.createElement("div");
+  containers.push(container);
+  document.body.append(container);
+  const root = createRoot(container);
+  let client: AssistantClient | undefined;
+
+  const TestAssistant = () => {
+    const runtime = useLocalRuntime(adapter, { initialMessages });
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        <RuntimeProbe
+          onReady={(aui) => {
+            client = aui;
+          }}
+        />
+        <ThreadPrimitive.Messages components={{ Message: TestMessage }} />
+      </AssistantRuntimeProvider>
+    );
+  };
+
+  await act(async () => {
+    root.render(createElement(TestAssistant));
+    await flushEffects();
+  });
+
+  if (!client) throw new Error("The assistant runtime was not registered");
+  return { client, container, root };
+};
+
+describe("useAssistantReasoningState", () => {
+  test("mounts and renders streamed reasoning without updating the chain resource before mount", async () => {
+    const adapter: ChatModelAdapter = {
+      async *run() {
+        yield { content: [{ type: "reasoning", text: "Inspecting the dashboard" }] };
+        await flushEffects();
+        yield {
+          content: [
+            { type: "reasoning", text: "Inspecting the dashboard" },
+            { type: "text", text: "Done" },
+          ],
+        };
+      },
+    };
+    const { client, container, root } = await renderAssistant(adapter);
+
+    await act(async () => {
+      client.composer().setText("Inspect my dashboard");
+      client.composer().send();
+      await flushEffects();
+      await flushEffects();
+    });
+
+    expect(container.querySelector('[data-testid="reasoning-part"]')?.textContent).toBe("Inspecting the dashboard");
+    expect(container.querySelector('[data-testid="chain-of-thought"]')?.getAttribute("data-status")).toBe("complete");
+
+    await act(async () => root.unmount());
+  });
+
+  test("opens a pending tool call even when completed reasoning is normally collapsed", async () => {
+    const pendingToolMessage: ThreadMessageLike = {
+      role: "assistant",
+      content: [
+        { type: "reasoning", text: "I need approval" },
+        {
+          type: "tool-call",
+          toolCallId: "approval-1",
+          toolName: "app_create",
+          args: { name: "Wikipedia" },
+          argsText: JSON.stringify({ name: "Wikipedia" }),
+        },
+      ],
+      status: { type: "requires-action", reason: "tool-calls" },
+    };
+    const { container, root } = await renderAssistant({ async *run() {} }, [pendingToolMessage]);
+    await act(async () => flushEffects());
+
+    expect(container.querySelector('[data-testid="tool-part"]')?.textContent).toBe("app_create");
+    expect(container.querySelector('[data-testid="chain-of-thought"]')?.getAttribute("data-status")).toBe(
+      "requires-action",
+    );
+    expect(container.querySelector('[data-testid="chain-of-thought"]')?.getAttribute("data-collapsed")).toBe("false");
+
+    await act(async () => root.unmount());
+  });
+});

apps/nextjs/src/components/assistant/assistant-reasoning-state.ts

@@ -0,0 +1,29 @@
+"use client";
+
+import { useEffect } from "react";
+import { useAui, useAuiState } from "@assistant-ui/react";
+
+/**
+ * Keeps the per-message chain of thought aligned with the conversation preference.
+ *
+ * assistant-ui commits its chain resource in a passive effect. Updating it from a
+ * layout effect races that commit and throws `Resource updated before mount` when
+ * reasoning or tool parts first appear. Keep all programmatic chain updates in
+ * passive effects so the resource is mounted before it receives an update.
+ */
+export const useAssistantReasoningState = (preferredCollapsed: boolean) => {
+  const aui = useAui();
+  const chainStatus = useAuiState((state) => state.chainOfThought.status);
+  const collapsed = useAuiState((state) => state.chainOfThought.collapsed);
+
+  useEffect(() => {
+    aui.chainOfThought().setCollapsed(preferredCollapsed);
+  }, [aui, preferredCollapsed]);
+
+  useEffect(() => {
+    if (chainStatus.type !== "requires-action") return;
+    aui.chainOfThought().setCollapsed(false);
+  }, [aui, chainStatus.type]);
+
+  return { chainStatus, collapsed };
+};