diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/assistant-message-projection.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/assistant-message-projection.ts index 4ae38bce..6e49a1a5 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/assistant-message-projection.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/assistant-message-projection.ts @@ -28,6 +28,9 @@ function toSessionMessageSegments(segments: SessionViewSegment[]): SessionMessag result.push({ kind: "text", text: segment.text }); break; } + case "reasoning": { + break; + } case "tool_use": { result.push({ argsText: segment.argsText, diff --git a/apps/api/tests/runtime-final-output-ingestion.test.ts b/apps/api/tests/runtime-final-output-ingestion.test.ts index 98d80868..62bff293 100644 --- a/apps/api/tests/runtime-final-output-ingestion.test.ts +++ b/apps/api/tests/runtime-final-output-ingestion.test.ts @@ -521,6 +521,58 @@ describe("runtime final output ingestion", () => { }); }); + test("omits live-only reasoning from stored final assistant segments", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + const privateReasoningText = "Private reasoning should stay out of stored history."; + const events = [ + runtimeEvent({ + kind: "thought.started", + payload: { messageId: finalMessageId }, + sourceEventId: "reasoning:started", + }), + runtimeEvent({ + kind: "thought.delta", + payload: { contentDelta: privateReasoningText, messageId: finalMessageId }, + sourceEventId: "reasoning:delta", + }), + runtimeEvent({ + kind: "thought.completed", + payload: { messageId: finalMessageId }, + sourceEventId: "reasoning:completed", + }), + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "reasoning:final", + text: FINAL_TEXT, + }), + runtimeEvent({ + kind: "run.completed", + payload: { + finalMessageId, + finalMessageText: FINAL_TEXT, + stopReason: "end_turn", + }, + sourceEventId: "reasoning:run-completed", + }), + ]; + + await pushFreshController(bindings, events); + + const persistedMessage = await database + .prepare("SELECT content_text, segments_json FROM session_message WHERE id = ?") + .bind(finalMessageId) + .first<{ content_text: string; segments_json: string }>(); + + expect(persistedMessage?.content_text).toBe(FINAL_TEXT); + expect(JSON.parse(persistedMessage?.segments_json ?? "[]")).toEqual([ + { kind: "text", text: FINAL_TEXT }, + ]); + expect(persistedMessage?.segments_json).not.toContain(privateReasoningText); + }); + test("fails closed when a cross-boot replay conflicts with the persisted final snapshot", async () => { const database = await createPublicHttpContractDatabase(); await insertRuntimeFixture(database); diff --git a/apps/web/src/features/session-chat/assistant-ui/convert-session-message.ts b/apps/web/src/features/session-chat/assistant-ui/convert-session-message.ts index d9bab3ae..4d01ef65 100644 --- a/apps/web/src/features/session-chat/assistant-ui/convert-session-message.ts +++ b/apps/web/src/features/session-chat/assistant-ui/convert-session-message.ts @@ -6,6 +6,7 @@ import type { SessionViewMessage, SessionViewSegment } from "@mosoo/ag-ui-sessio // is resolved separately in the renderer (see session-permission-context). type AssistantContentPart = | { type: "text"; text: string } + | { type: "reasoning"; text: string } | { type: "tool-call"; toolCallId: string; @@ -37,6 +38,18 @@ function sessionSegmentsToParts(segments: readonly SessionViewSegment[]): Assist continue; } + if (segment.kind === "reasoning") { + const last = parts.at(-1); + + if (last?.type === "reasoning") { + last.text += segment.text; + } else { + parts.push({ type: "reasoning", text: segment.text }); + } + + continue; + } + if (segment.kind === "tool_use") { const existingIndex = indexByCallId.get(segment.toolCallId); diff --git a/apps/web/src/features/session-chat/assistant-ui/session-message-parts.tsx b/apps/web/src/features/session-chat/assistant-ui/session-message-parts.tsx index 0658f12d..4c49a4a8 100644 --- a/apps/web/src/features/session-chat/assistant-ui/session-message-parts.tsx +++ b/apps/web/src/features/session-chat/assistant-ui/session-message-parts.tsx @@ -1,7 +1,8 @@ import { ActionBarPrimitive, MessagePrimitive, useMessage } from "@assistant-ui/react"; -import type { TextMessagePartComponent } from "@assistant-ui/react"; -import { Copy } from "lucide-react"; +import type { ReasoningMessagePartComponent, TextMessagePartComponent } from "@assistant-ui/react"; +import { ChevronRight, Copy } from "lucide-react"; import type { ReactElement } from "react"; +import { useState } from "react"; import { Markdown } from "@/shared/ui/markdown"; @@ -43,6 +44,43 @@ function ThinkingIndicator({ status }: { status: { readonly type: string } }): R ); } +// Reasoning stream from thinking models (e.g. DeepSeek V4 Pro reasons for +// minutes before the first visible token). Rendered as a collapsed disclosure +// so the user sees live progress — shimmering "Thinking" plus a growing +// character count — instead of dead air, and can expand the raw thought text. +const AssistantReasoning: ReasoningMessagePartComponent = ({ status, text }) => { + const [expanded, setExpanded] = useState(false); + const running = status.type === "running"; + const charCount = text.length; + + return ( +
+ + {expanded ? ( +
+ {text} +
+ ) : null} +
+ ); +}; + export function UserMessage(): ReactElement { return ( @@ -72,6 +110,7 @@ export function AssistantMessage(): ReactElement { unstable_showEmptyOnNonTextEnd={false} components={{ Empty: ThinkingIndicator, + Reasoning: AssistantReasoning, Text: AssistantText, tools: { Override: SessionToolPart }, }} diff --git a/pkgs/ag-ui-session/src/live-state-message-reasoning.reducer.ts b/pkgs/ag-ui-session/src/live-state-message-reasoning.reducer.ts new file mode 100644 index 00000000..cad93fc1 --- /dev/null +++ b/pkgs/ag-ui-session/src/live-state-message-reasoning.reducer.ts @@ -0,0 +1,72 @@ +import type { SessionLiveState } from "./live-state"; +import { createLiveStateMessage } from "./live-state-message-core.reducer"; +import { touchSessionLiveState } from "./live-state.reducer-core"; + +// Reasoning streams (REASONING_MESSAGE_*) get their own assistant message keyed +// by the reasoning messageId, holding a single growing `reasoning` segment. +// Thinking models can reason for minutes before the first visible token; without +// this the whole phase renders as dead air. The text is intentionally NOT added +// to `content`, so copy actions and text fallbacks ignore it. +export function appendReasoningDelta( + state: SessionLiveState, + input: { delta: string; messageId: string }, +): SessionLiveState { + const messages = [...state.messages]; + const index = messages.findIndex((message) => message.id === input.messageId); + const current = index === -1 ? undefined : messages[index]; + + if (current === undefined) { + messages.push( + createLiveStateMessage({ + content: "", + id: input.messageId, + role: "assistant", + segments: [{ kind: "reasoning", text: input.delta }], + }), + ); + + return touchSessionLiveState({ + ...state, + messages, + }); + } + + if (input.delta.length === 0) { + return state; + } + + const segments = [...current.segments]; + const lastSegment = segments.at(-1); + + if (lastSegment?.kind === "reasoning") { + segments[segments.length - 1] = { + kind: "reasoning", + text: `${lastSegment.text}${input.delta}`, + }; + } else { + segments.push({ kind: "reasoning", text: input.delta }); + } + + messages[index] = { + ...current, + segments, + }; + + return touchSessionLiveState({ + ...state, + messages, + }); +} + +export function startReasoning( + state: SessionLiveState, + input: { messageId: string }, +): SessionLiveState { + const exists = state.messages.some((message) => message.id === input.messageId); + + if (exists) { + return state; + } + + return appendReasoningDelta(state, { delta: "", messageId: input.messageId }); +} diff --git a/pkgs/ag-ui-session/src/live-state-message.reducer.ts b/pkgs/ag-ui-session/src/live-state-message.reducer.ts index 19a34fc9..b043949b 100644 --- a/pkgs/ag-ui-session/src/live-state-message.reducer.ts +++ b/pkgs/ag-ui-session/src/live-state-message.reducer.ts @@ -3,6 +3,7 @@ export { normalizeMessagePlan, upsertMessage, } from "./live-state-message-core.reducer"; +export { appendReasoningDelta, startReasoning } from "./live-state-message-reasoning.reducer"; export { appendTextDelta } from "./live-state-message-text.reducer"; export { completePendingToolUses, diff --git a/pkgs/ag-ui-session/src/live-state.reducer.ts b/pkgs/ag-ui-session/src/live-state.reducer.ts index 5f79214e..abeeba80 100644 --- a/pkgs/ag-ui-session/src/live-state.reducer.ts +++ b/pkgs/ag-ui-session/src/live-state.reducer.ts @@ -7,6 +7,7 @@ import type { SessionLiveState, SessionViewMessage } from "./live-state"; import { updateCustomState } from "./live-state-custom.reducer"; import { applyJsonPatch } from "./live-state-json-patch.reducer"; import { + appendReasoningDelta, appendTextDelta, appendToolArgs, appendToolResult, @@ -14,6 +15,7 @@ import { completePendingToolUses, completeToolUse, createSessionLiveStateMessage, + startReasoning, upsertMessage, } from "./live-state-message.reducer"; import { @@ -279,14 +281,21 @@ function applyEvent(state: SessionLiveState, event: AgUiEvent): SessionLiveState }, }); } + case EventType.REASONING_MESSAGE_START: { + return startReasoning(currentState, { messageId: event.messageId }); + } + case EventType.REASONING_MESSAGE_CONTENT: { + return appendReasoningDelta(currentState, { + delta: event.delta, + messageId: event.messageId, + }); + } case EventType.ACTIVITY_DELTA: case EventType.ACTIVITY_SNAPSHOT: case EventType.REASONING_ENCRYPTED_VALUE: case EventType.REASONING_END: case EventType.REASONING_MESSAGE_CHUNK: - case EventType.REASONING_MESSAGE_CONTENT: case EventType.REASONING_MESSAGE_END: - case EventType.REASONING_MESSAGE_START: case EventType.REASONING_START: case EventType.STEP_FINISHED: case EventType.STEP_STARTED: diff --git a/pkgs/ag-ui-session/src/live-state.ts b/pkgs/ag-ui-session/src/live-state.ts index cb5c712b..d0207a43 100644 --- a/pkgs/ag-ui-session/src/live-state.ts +++ b/pkgs/ag-ui-session/src/live-state.ts @@ -9,6 +9,10 @@ export type SessionViewSegment = kind: "text"; text: string; } + | { + kind: "reasoning"; + text: string; + } | { argsText: string; kind: "tool_use"; diff --git a/pkgs/ag-ui-session/src/session-live-state-schema.ts b/pkgs/ag-ui-session/src/session-live-state-schema.ts index ea1bedde..3614c606 100644 --- a/pkgs/ag-ui-session/src/session-live-state-schema.ts +++ b/pkgs/ag-ui-session/src/session-live-state-schema.ts @@ -17,6 +17,10 @@ export const SessionViewSegmentSchema = type.or( kind: '"text"', text: "string", }), + type({ + kind: '"reasoning"', + text: "string", + }), type({ argsText: "string", kind: '"tool_use"', diff --git a/pkgs/ag-ui-session/tests/live-state.reducer.test.ts b/pkgs/ag-ui-session/tests/live-state.reducer.test.ts index df795517..b14543ef 100644 --- a/pkgs/ag-ui-session/tests/live-state.reducer.test.ts +++ b/pkgs/ag-ui-session/tests/live-state.reducer.test.ts @@ -111,6 +111,55 @@ describe("session live-state transcript reducer", () => { ]); }); + test("streams reasoning deltas into a collapsed reasoning segment", () => { + const nextState = applyAgUiEventsToSessionLiveState(baseState(), [ + { messageId: "prompt-1:thought", role: "reasoning", type: "REASONING_MESSAGE_START" }, + { + delta: "The user asks about MCP.", + messageId: "prompt-1:thought", + type: "REASONING_MESSAGE_CONTENT", + }, + { + delta: " Check the configured servers.", + messageId: "prompt-1:thought", + type: "REASONING_MESSAGE_CONTENT", + }, + { messageId: "prompt-1:thought", type: "REASONING_MESSAGE_END" }, + ]); + + expect(nextState.messages).toHaveLength(1); + expect(nextState.messages[0]?.role).toBe("assistant"); + expect(nextState.messages[0]?.content).toBe(""); + expect(nextState.messages[0]?.segments).toEqual([ + { kind: "reasoning", text: "The user asks about MCP. Check the configured servers." }, + ]); + }); + + test("reasoning start alone creates the message so the UI can show progress", () => { + const nextState = applyAgUiEventsToSessionLiveState(baseState(), [ + { messageId: "prompt-1:thought", role: "reasoning", type: "REASONING_MESSAGE_START" }, + ]); + + expect(nextState.messages).toHaveLength(1); + expect(nextState.messages[0]?.segments).toEqual([{ kind: "reasoning", text: "" }]); + }); + + test("keeps reasoning separate from assistant text in the same turn", () => { + const nextState = applyAgUiEventsToSessionLiveState(baseState(), [ + { messageId: "prompt-1:thought", role: "reasoning", type: "REASONING_MESSAGE_START" }, + { delta: "thinking...", messageId: "prompt-1:thought", type: "REASONING_MESSAGE_CONTENT" }, + { messageId: "assistant-1", role: "assistant", type: "TEXT_MESSAGE_START" }, + { delta: "Answer.", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }, + { messageId: "prompt-1:thought", type: "REASONING_MESSAGE_END" }, + { messageId: "assistant-1", type: "TEXT_MESSAGE_END" }, + ]); + + expect(nextState.messages).toHaveLength(2); + expect(nextState.messages[0]?.segments).toEqual([{ kind: "reasoning", text: "thinking..." }]); + expect(nextState.messages[1]?.segments).toEqual([{ kind: "text", text: "Answer." }]); + expect(nextState.messages[1]?.content).toBe("Answer."); + }); + test("does not merge streamed text across intervening transcript events", () => { const nextState = applyAgUiEventsToSessionLiveState(baseState(), [ { messageId: "assistant-1", role: "assistant", type: "TEXT_MESSAGE_START" },