From 7f209a6175c2c007b93c31a27050d472a354fd7e Mon Sep 17 00:00:00 2001 From: Shinyaigeek Date: Wed, 3 Jun 2026 22:51:28 +0900 Subject: [PATCH] feat(react-headless,react-ui): render streaming reasoning content Consume AG-UI REASONING_MESSAGE_* events end-to-end. The stream processor now accumulates reasoning deltas into a colocated AssistantMessage.reasoning field, and the Shell / CopilotShell / BottomTray / OpenUIChat surfaces render it inside a collapsible "Behind the scenes" panel via a shared ReasoningSection component. - types: extend AssistantMessage with an optional `reasoning` field and swap the Message union's assistant branch so narrowing exposes it - processStreamedMessage: handle REASONING_MESSAGE_CONTENT/CHUNK, ignoring the reasoning-only messageId (preserved across the TEXT_MESSAGE_START id swap) - identityMessageFormat: strip the client-only `reasoning` field from outbound history so it is not echoed back to the backend each turn - react-ui: new Reasoning/ component dir (ReasoningContent + ReasoningSection), applied to all assistant render surfaces; exported from the package root Tests: reasoning accumulation, id-swap preservation, absence, and the identityMessageFormat strip. Addresses thesysdev/openui#414 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../store/__tests__/createChatStore.test.ts | 80 +++++++++++++++++++ .../src/stream/processStreamedMessage.ts | 13 +++ .../src/types/__tests__/messageFormat.test.ts | 36 +++++++++ packages/react-headless/src/types/message.ts | 21 ++++- .../react-headless/src/types/messageFormat.ts | 19 ++++- .../src/components/BottomTray/Thread.tsx | 14 +++- .../src/components/CopilotShell/Thread.tsx | 14 +++- .../OpenUIChat/GenUIAssistantMessage.tsx | 4 +- .../components/Reasoning/ReasoningContent.tsx | 15 ++++ .../components/Reasoning/ReasoningSection.tsx | 26 ++++++ .../src/components/Reasoning/index.ts | 2 + .../src/components/Reasoning/reasoning.scss | 9 +++ .../react-ui/src/components/Shell/Thread.tsx | 14 +++- packages/react-ui/src/components/index.scss | 1 + packages/react-ui/src/index.ts | 1 + 15 files changed, 260 insertions(+), 9 deletions(-) create mode 100644 packages/react-headless/src/types/__tests__/messageFormat.test.ts create mode 100644 packages/react-ui/src/components/Reasoning/ReasoningContent.tsx create mode 100644 packages/react-ui/src/components/Reasoning/ReasoningSection.tsx create mode 100644 packages/react-ui/src/components/Reasoning/index.ts create mode 100644 packages/react-ui/src/components/Reasoning/reasoning.scss diff --git a/packages/react-headless/src/store/__tests__/createChatStore.test.ts b/packages/react-headless/src/store/__tests__/createChatStore.test.ts index c50d0dd1c..c1df563ad 100644 --- a/packages/react-headless/src/store/__tests__/createChatStore.test.ts +++ b/packages/react-headless/src/store/__tests__/createChatStore.test.ts @@ -396,6 +396,13 @@ describe("createChatStore", () => { beforeEach(() => { fetchSpy = vi.fn(); vi.stubGlobal("fetch", fetchSpy); + // processStreamedMessage debounces updates via requestAnimationFrame, + // which is absent in the node test environment. Run callbacks synchronously. + vi.stubGlobal("requestAnimationFrame", (cb: FrameRequestCallback) => { + cb(0); + return 0; + }); + vi.stubGlobal("cancelAnimationFrame", () => {}); }); it("sends POST to apiUrl with threadId and messages", async () => { @@ -458,6 +465,79 @@ describe("createChatStore", () => { expect((store.getState().messages[1] as any).content).toBe("response text"); }); + it("accumulates REASONING_MESSAGE_CONTENT into the assistant message's reasoning", async () => { + const sseBody = + `data: ${JSON.stringify({ type: "REASONING_MESSAGE_START", messageId: "r1", role: "reasoning" })}\n\n` + + `data: ${JSON.stringify({ type: "REASONING_MESSAGE_CONTENT", messageId: "r1", delta: "let me " })}\n\n` + + `data: ${JSON.stringify({ type: "REASONING_MESSAGE_CONTENT", messageId: "r1", delta: "think" })}\n\n` + + `data: ${JSON.stringify({ type: "REASONING_MESSAGE_END", messageId: "r1" })}\n\n` + + `data: [DONE]\n\n`; + const stream = new ReadableStream({ + start(c) { + c.enqueue(new TextEncoder().encode(sseBody)); + c.close(); + }, + }); + fetchSpy.mockResolvedValue(new Response(stream)); + + const store = createChatStore({ apiUrl: "/api/chat" }); + store.setState({ selectedThreadId: "t1" }); + + await store.getState().processMessage({ role: "user", content: "hello" }); + + const assistant = store.getState().messages[1] as any; + expect(assistant.role).toBe("assistant"); + expect(assistant.reasoning).toBe("let me think"); + }); + + it("leaves reasoning undefined when no reasoning events are streamed", async () => { + const sseBody = `data: ${JSON.stringify({ type: "TEXT_MESSAGE_CONTENT", delta: "just an answer" })}\n\ndata: [DONE]\n\n`; + const stream = new ReadableStream({ + start(c) { + c.enqueue(new TextEncoder().encode(sseBody)); + c.close(); + }, + }); + fetchSpy.mockResolvedValue(new Response(stream)); + + const store = createChatStore({ apiUrl: "/api/chat" }); + store.setState({ selectedThreadId: "t1" }); + + await store.getState().processMessage({ role: "user", content: "hello" }); + + const assistant = store.getState().messages[1] as any; + expect(assistant.role).toBe("assistant"); + expect(assistant.content).toBe("just an answer"); + expect(assistant.reasoning).toBeUndefined(); + }); + + it("preserves reasoning across a TEXT_MESSAGE_START id swap and colocates it with content", async () => { + const sseBody = + `data: ${JSON.stringify({ type: "REASONING_MESSAGE_CONTENT", messageId: "r1", delta: "thinking..." })}\n\n` + + `data: ${JSON.stringify({ type: "TEXT_MESSAGE_START", messageId: "m1", role: "assistant" })}\n\n` + + `data: ${JSON.stringify({ type: "TEXT_MESSAGE_CONTENT", messageId: "m1", delta: "answer" })}\n\n` + + `data: [DONE]\n\n`; + const stream = new ReadableStream({ + start(c) { + c.enqueue(new TextEncoder().encode(sseBody)); + c.close(); + }, + }); + fetchSpy.mockResolvedValue(new Response(stream)); + + const store = createChatStore({ apiUrl: "/api/chat" }); + store.setState({ selectedThreadId: "t1" }); + + await store.getState().processMessage({ role: "user", content: "hello" }); + + const assistants = store.getState().messages.filter((m) => m.role === "assistant"); + expect(assistants).toHaveLength(1); + const assistant = assistants[0] as any; + expect(assistant.id).toBe("m1"); + expect(assistant.reasoning).toBe("thinking..."); + expect(assistant.content).toBe("answer"); + }); + it("throws when neither apiUrl nor processMessage provided", async () => { const store = createChatStore({}); store.setState({ selectedThreadId: "t1" }); diff --git a/packages/react-headless/src/stream/processStreamedMessage.ts b/packages/react-headless/src/stream/processStreamedMessage.ts index 507f28c2a..6797d0817 100644 --- a/packages/react-headless/src/stream/processStreamedMessage.ts +++ b/packages/react-headless/src/stream/processStreamedMessage.ts @@ -95,6 +95,19 @@ export const processStreamedMessage = async ({ } break; + // Reasoning/thinking deltas, colocated on the assistant turn. CHUNK may + // carry a role change but, like TEXT_MESSAGE_CHUNK, we treat it as CONTENT. + // The reasoning's own messageId is ignored — a later TEXT_MESSAGE_START + // swaps in the real assistant id, and the spread preserves `reasoning`. + // (START/END carry no payload we need, so they fall through unhandled.) + case EventType.REASONING_MESSAGE_CHUNK: + case EventType.REASONING_MESSAGE_CONTENT: + currentMessage = { + ...currentMessage, + reasoning: (currentMessage.reasoning || "") + (event.delta ?? ""), + }; + break; + case EventType.TEXT_MESSAGE_START: // Use the ID from the event if it differs from our optimistic ID if (event.messageId !== currentMessage.id) { diff --git a/packages/react-headless/src/types/__tests__/messageFormat.test.ts b/packages/react-headless/src/types/__tests__/messageFormat.test.ts new file mode 100644 index 000000000..7ef5679d3 --- /dev/null +++ b/packages/react-headless/src/types/__tests__/messageFormat.test.ts @@ -0,0 +1,36 @@ +import { describe, expect, it } from "vitest"; +import type { AssistantMessage, Message } from "../message"; +import { identityMessageFormat } from "../messageFormat"; + +describe("identityMessageFormat", () => { + it("strips the client-only reasoning field from assistant messages in toApi", () => { + const messages: Message[] = [ + { id: "u1", role: "user", content: "hi" }, + { id: "a1", role: "assistant", content: "hello", reasoning: "let me think" }, + ]; + + const out = identityMessageFormat.toApi(messages) as Message[]; + + expect(out).toHaveLength(2); + const assistant = out[1] as AssistantMessage; + expect(assistant.content).toBe("hello"); + expect("reasoning" in assistant).toBe(false); + }); + + it("leaves messages without reasoning untouched (same reference)", () => { + const messages: Message[] = [ + { id: "u1", role: "user", content: "hi" }, + { id: "a1", role: "assistant", content: "hello" }, + ]; + + const out = identityMessageFormat.toApi(messages) as Message[]; + + expect(out[0]).toBe(messages[0]); + expect(out[1]).toBe(messages[1]); + }); + + it("fromApi passes data through unchanged", () => { + const data: Message[] = [{ id: "a1", role: "assistant", content: "hello" }]; + expect(identityMessageFormat.fromApi(data)).toBe(data); + }); +}); diff --git a/packages/react-headless/src/types/message.ts b/packages/react-headless/src/types/message.ts index 2458172cf..f43497c1a 100644 --- a/packages/react-headless/src/types/message.ts +++ b/packages/react-headless/src/types/message.ts @@ -1,11 +1,11 @@ +import type { AssistantMessage as AGUIAssistantMessage, Message as AGUIMessage } from "@ag-ui/core"; + export type { ActivityMessage, - AssistantMessage, BinaryInputContent, DeveloperMessage, FunctionCall, InputContent, - Message, ReasoningMessage, SystemMessage, TextInputContent, @@ -13,3 +13,20 @@ export type { ToolMessage, UserMessage, } from "@ag-ui/core"; + +/** + * AG-UI's assistant message, extended with a colocated `reasoning` field so + * thinking/reasoning tokens can live on the same turn as the answer (rendered + * in something like a BehindTheScenes section). + */ +export type AssistantMessage = AGUIAssistantMessage & { + /** Reasoning/thinking tokens colocated with this assistant turn. */ + reasoning?: string; +}; + +/** + * AG-UI's message union with the assistant branch swapped for the extended + * {@link AssistantMessage}, so `role === "assistant"` narrowing exposes + * `reasoning`. + */ +export type Message = Exclude | AssistantMessage; diff --git a/packages/react-headless/src/types/messageFormat.ts b/packages/react-headless/src/types/messageFormat.ts index 9801653df..df14c1662 100644 --- a/packages/react-headless/src/types/messageFormat.ts +++ b/packages/react-headless/src/types/messageFormat.ts @@ -23,10 +23,23 @@ export interface MessageFormat { } /** - * Default identity message format — no conversion. - * Messages are sent and received as-is in AG-UI format. + * `reasoning` is a client-side display field colocated on the assistant turn, + * not part of the wire message shape. Strip it so the identity format doesn't + * echo accumulated thinking back to the backend on every subsequent turn. + */ +function stripReasoning(message: Message): Message { + if (message.role !== "assistant" || message.reasoning === undefined) return message; + const clone = { ...message }; + delete clone.reasoning; + return clone; +} + +/** + * Default identity message format — no conversion (aside from dropping the + * client-only `reasoning` field). Messages are sent and received as-is in + * AG-UI format. */ export const identityMessageFormat: MessageFormat = { - toApi: (messages) => messages, + toApi: (messages) => messages.map(stripReasoning), fromApi: (data) => data as Message[], }; diff --git a/packages/react-ui/src/components/BottomTray/Thread.tsx b/packages/react-ui/src/components/BottomTray/Thread.tsx index 0a9578d13..f88386849 100644 --- a/packages/react-ui/src/components/BottomTray/Thread.tsx +++ b/packages/react-ui/src/components/BottomTray/Thread.tsx @@ -7,6 +7,7 @@ import { ArtifactOverlay } from "../_shared/artifact"; import type { AssistantMessageComponent, UserMessageComponent } from "../_shared/types"; import { MarkDownRenderer } from "../MarkDownRenderer"; import { MessageLoading as MessageLoadingComponent } from "../MessageLoading"; +import { ReasoningSection } from "../Reasoning"; import { ToolCallComponent } from "../ToolCall"; import { ToolResult } from "../ToolResult"; @@ -114,9 +115,11 @@ export const UserMessageContainer = ({ const AssistantMessageContent = ({ message, allMessages, + isStreaming, }: { message: AssistantMessage; allMessages: Message[]; + isStreaming: boolean; }) => { const getToolName = (toolCallId: string) => { const toolCall = message.toolCalls?.find((tc) => tc.id === toolCallId); @@ -138,6 +141,11 @@ const AssistantMessageContent = ({ return ( <> + {message.content && ( - + ); } diff --git a/packages/react-ui/src/components/CopilotShell/Thread.tsx b/packages/react-ui/src/components/CopilotShell/Thread.tsx index ffc103890..fcb64ae3d 100644 --- a/packages/react-ui/src/components/CopilotShell/Thread.tsx +++ b/packages/react-ui/src/components/CopilotShell/Thread.tsx @@ -8,6 +8,7 @@ import { useShellStore } from "../_shared/store"; import type { AssistantMessageComponent, UserMessageComponent } from "../_shared/types"; import { MarkDownRenderer } from "../MarkDownRenderer"; import { MessageLoading as MessageLoadingComponent } from "../MessageLoading"; +import { ReasoningSection } from "../Reasoning"; import { ToolCallComponent } from "../ToolCall"; import { ToolResult } from "../ToolResult"; @@ -127,9 +128,11 @@ export const UserMessageContainer = ({ const AssistantMessageContent = ({ message, allMessages, + isStreaming, }: { message: AssistantMessage; allMessages: Message[]; + isStreaming: boolean; }) => { const getToolName = (toolCallId: string) => { const toolCall = message.toolCalls?.find((tc) => tc.id === toolCallId); @@ -151,6 +154,11 @@ const AssistantMessageContent = ({ return ( <> + {message.content && ( - + ); } diff --git a/packages/react-ui/src/components/OpenUIChat/GenUIAssistantMessage.tsx b/packages/react-ui/src/components/OpenUIChat/GenUIAssistantMessage.tsx index 2a3728137..8623366ba 100644 --- a/packages/react-ui/src/components/OpenUIChat/GenUIAssistantMessage.tsx +++ b/packages/react-ui/src/components/OpenUIChat/GenUIAssistantMessage.tsx @@ -6,6 +6,7 @@ import type { ActionEvent, Library } from "@openuidev/react-lang"; import { BuiltinActionType, Renderer } from "@openuidev/react-lang"; import { useCallback, useMemo } from "react"; import { separateContentAndContext, wrapContent, wrapContext } from "../../utils/contentParser"; +import { ReasoningContent } from "../Reasoning"; import { AssistantMessageContainer } from "../Shell"; import { BehindTheScenes, ToolCallComponent } from "../ToolCall"; import { ToolResult } from "../ToolResult"; @@ -115,8 +116,9 @@ export const GenUIAssistantMessage = ({ return ( - {hasToolActivity && ( + {(hasToolActivity || !!message.reasoning) && ( + {message.reasoning && } {message.toolCalls?.map((toolCall, idx) => ( ( + +); diff --git a/packages/react-ui/src/components/Reasoning/ReasoningSection.tsx b/packages/react-ui/src/components/Reasoning/ReasoningSection.tsx new file mode 100644 index 000000000..36954e106 --- /dev/null +++ b/packages/react-ui/src/components/Reasoning/ReasoningSection.tsx @@ -0,0 +1,26 @@ +import { BehindTheScenes } from "../ToolCall"; +import { ReasoningContent } from "./ReasoningContent"; + +export interface ReasoningSectionProps { + /** Accumulated reasoning/thinking text for the assistant turn, if any. */ + reasoning: string | undefined; + /** True while the assistant turn is still streaming. */ + isStreaming: boolean; + /** True once answer text has started — collapses the section. */ + hasContent: boolean; +} + +/** + * Standalone collapsible reasoning panel for assistant turns whose tool calls + * are rendered outside a BehindTheScenes (the Shell / CopilotShell / BottomTray + * surfaces). Renders nothing when there is no reasoning. The GenUI surface + * instead nests `ReasoningContent` inside its shared tool-call panel. + */ +export const ReasoningSection = ({ reasoning, isStreaming, hasContent }: ReasoningSectionProps) => { + if (!reasoning) return null; + return ( + + + + ); +}; diff --git a/packages/react-ui/src/components/Reasoning/index.ts b/packages/react-ui/src/components/Reasoning/index.ts new file mode 100644 index 000000000..4fffee853 --- /dev/null +++ b/packages/react-ui/src/components/Reasoning/index.ts @@ -0,0 +1,2 @@ +export { ReasoningContent, type ReasoningContentProps } from "./ReasoningContent"; +export { ReasoningSection, type ReasoningSectionProps } from "./ReasoningSection"; diff --git a/packages/react-ui/src/components/Reasoning/reasoning.scss b/packages/react-ui/src/components/Reasoning/reasoning.scss new file mode 100644 index 000000000..48ba0058b --- /dev/null +++ b/packages/react-ui/src/components/Reasoning/reasoning.scss @@ -0,0 +1,9 @@ +@use "../../cssUtils" as cssUtils; + +// ── Reasoning / thinking content ── + +.openui-reasoning { + width: 100%; + padding-left: cssUtils.$space-l; + color: cssUtils.$text-neutral-secondary; +} diff --git a/packages/react-ui/src/components/Shell/Thread.tsx b/packages/react-ui/src/components/Shell/Thread.tsx index df22d6717..e3c6ef847 100644 --- a/packages/react-ui/src/components/Shell/Thread.tsx +++ b/packages/react-ui/src/components/Shell/Thread.tsx @@ -11,6 +11,7 @@ import type { AssistantMessageComponent, UserMessageComponent } from "../_shared import { Callout } from "../Callout"; import { MarkDownRenderer } from "../MarkDownRenderer"; import { MessageLoading as MessageLoadingComponent } from "../MessageLoading"; +import { ReasoningSection } from "../Reasoning"; import { ToolCallComponent } from "../ToolCall"; import { ToolResult } from "../ToolResult"; import { ResizableSeparator } from "./ResizableSeparator"; @@ -189,9 +190,11 @@ export const UserMessageContainer = ({ const AssistantMessageContent = ({ message, allMessages, + isStreaming, }: { message: AssistantMessage; allMessages: Message[]; + isStreaming: boolean; }) => { // Find tool result messages that correspond to this message's tool calls const getToolName = (toolCallId: string) => { @@ -215,6 +218,11 @@ const AssistantMessageContent = ({ return ( <> + {message.content && ( - + ); } diff --git a/packages/react-ui/src/components/index.scss b/packages/react-ui/src/components/index.scss index 84471fec1..d1c88f3ec 100644 --- a/packages/react-ui/src/components/index.scss +++ b/packages/react-ui/src/components/index.scss @@ -37,6 +37,7 @@ @forward "./OpenUIChat/openUIChat.scss"; @forward "./RadioGroup/radioGroup.scss"; @forward "./RadioItem/radioItem.scss"; +@forward "./Reasoning/reasoning.scss"; @forward "./SectionBlock/sectionBlock.scss"; @forward "./Select/select.scss"; @forward "./Separator/separator.scss"; diff --git a/packages/react-ui/src/index.ts b/packages/react-ui/src/index.ts index 0d01180fa..3958c3eea 100644 --- a/packages/react-ui/src/index.ts +++ b/packages/react-ui/src/index.ts @@ -49,6 +49,7 @@ export * from "./components/MessageLoading"; export * from "./components/OpenUIChat"; export * from "./components/RadioGroup"; export * from "./components/RadioItem"; +export * from "./components/Reasoning"; export * from "./components/SectionBlock"; export * from "./components/Select"; export * from "./components/Separator";