From f61e07ae58fde0fbf6a315f7617fbd9d3c60f5f6 Mon Sep 17 00:00:00 2001 From: ajnart Date: Sun, 9 Aug 2026 16:37:26 +0200 Subject: [PATCH] fix assistant reasoning mount lifecycle --- apps/docs/docs/management/assistant.mdx | 7 +- .../assistant/assistant-panel.module.css | 16 +- .../components/assistant/assistant-panel.tsx | 23 +-- .../assistant-reasoning-state.spec.tsx | 144 ++++++++++++++++++ .../assistant/assistant-reasoning-state.ts | 29 ++++ 5 files changed, 185 insertions(+), 34 deletions(-) create mode 100644 apps/nextjs/src/components/assistant/assistant-reasoning-state.spec.tsx create mode 100644 apps/nextjs/src/components/assistant/assistant-reasoning-state.ts diff --git a/apps/docs/docs/management/assistant.mdx b/apps/docs/docs/management/assistant.mdx index 8d4f1e073..ff0d83af5 100644 --- a/apps/docs/docs/management/assistant.mdx +++ b/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, diff --git a/apps/nextjs/src/components/assistant/assistant-panel.module.css b/apps/nextjs/src/components/assistant/assistant-panel.module.css index c57a39996..ac8ea0e7b 100644 --- a/apps/nextjs/src/components/assistant/assistant-panel.module.css +++ b/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; diff --git a/apps/nextjs/src/components/assistant/assistant-panel.tsx b/apps/nextjs/src/components/assistant/assistant-panel.tsx index c16dbcccb..23ff73fbf 100644 --- a/apps/nextjs/src/components/assistant/assistant-panel.tsx +++ b/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 ( @@ -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)} > { roundCaps sections={hasContext ? [{ value: contextPercentage, color: getContextColor(contextPercentage) }] : []} /> - {quickLabel && ( - - {quickLabel} - - )} diff --git a/apps/nextjs/src/components/assistant/assistant-reasoning-state.spec.tsx b/apps/nextjs/src/components/assistant/assistant-reasoning-state.spec.tsx new file mode 100644 index 000000000..222bc251a --- /dev/null +++ b/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 }) => {text}; +const ToolPart = ({ toolName }: { toolName: string }) => {toolName}; + +const TestChainOfThought = ({ preferredCollapsed = true }: { preferredCollapsed?: boolean }) => { + const { chainStatus, collapsed } = useAssistantReasoningState(preferredCollapsed); + + return ( + + + + ); +}; + +const TestMessage = () => ( + + + +); + +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 ( + + { + client = aui; + }} + /> + + + ); + }; + + 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()); + }); +}); diff --git a/apps/nextjs/src/components/assistant/assistant-reasoning-state.ts b/apps/nextjs/src/components/assistant/assistant-reasoning-state.ts new file mode 100644 index 000000000..6d5b49dfa --- /dev/null +++ b/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 }; +};