fix assistant reasoning mount lifecycle

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