diff --git a/.changeset/preserve-server-reasoning.md b/.changeset/preserve-server-reasoning.md new file mode 100644 index 0000000000..a99b247723 --- /dev/null +++ b/.changeset/preserve-server-reasoning.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +Preserve reasoning in server chat message history and interrupt snapshots. diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 668c873ed0..1d7cbb969d 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -42,7 +42,11 @@ import { } from './tools/approval-schema' import { maxIterations as maxIterationsStrategy } from './agent-loop-strategies' import { isCancelRequestedReason } from './cancel' -import { convertMessagesToModelMessages, generateMessageId } from './messages' +import { + convertMessagesToModelMessages, + generateMessageId, + modelMessageToUIMessage, +} from './messages' import { MiddlewareRunner } from './middleware/compose' import { getRunDetached } from './middleware/run-store' import { publishRunDetachedSignal } from '../../delivery-detach' @@ -738,6 +742,7 @@ class TextEngine< [] private currentThinkingContent = '' private currentThinkingSignature = '' + private hasSeenReasoningEvents = false private eventOptions?: Record | undefined private eventToolNames?: Array private finishedEvent: RunFinishedEvent | null = null @@ -1152,6 +1157,7 @@ class TextEngine< duration: Date.now() - this.streamStartTime, }) } else { + this.addTerminalReasoningMessage() this.terminalHookCalled = true await this.middlewareRunner.runOnFinish(this.middlewareCtx, { finishReason: this.lastFinishReason, @@ -1282,6 +1288,7 @@ class TextEngine< this.accumulatedThinking = [] this.currentThinkingContent = '' this.currentThinkingSignature = '' + this.hasSeenReasoningEvents = false this.finishedEvent = null this.streamedToolErrorResults.clear() @@ -1533,16 +1540,19 @@ class TextEngine< this.handleStepFinishedEvent(chunk) break + case 'REASONING_MESSAGE_CONTENT': + this.handleReasoningMessageContentEvent(chunk) + break + case 'TOOL_CALL_RESULT': // Tool result is already added to messages in buildToolResultChunks break case 'REASONING_START': case 'REASONING_MESSAGE_START': - case 'REASONING_MESSAGE_CONTENT': case 'REASONING_MESSAGE_END': case 'REASONING_END': - // Reasoning events are handled by StreamProcessor + // No special handling needed break default: @@ -1651,14 +1661,29 @@ class TextEngine< private handleStepFinishedEvent( chunk: Extract, ): void { - if (chunk.delta) { - this.currentThinkingContent += chunk.delta + if (!this.hasSeenReasoningEvents) { + if (chunk.delta) { + this.currentThinkingContent += chunk.delta + } else if (chunk.content) { + if (chunk.content.startsWith(this.currentThinkingContent)) { + this.currentThinkingContent = chunk.content + } else if (!this.currentThinkingContent.startsWith(chunk.content)) { + this.currentThinkingContent += chunk.content + } + } } if (chunk.signature) { this.currentThinkingSignature = chunk.signature } } + private handleReasoningMessageContentEvent( + chunk: Extract, + ): void { + this.hasSeenReasoningEvents = true + this.currentThinkingContent += chunk.delta + } + /** * Tools available for execution this turn. The discovery tool is dropped * from the advertised set (`this.tools`) once every lazy tool is discovered, @@ -2065,6 +2090,30 @@ class TextEngine< this.middlewareCtx.messages = this.messages } + private addTerminalReasoningMessage(): void { + this.finalizeCurrentThinkingStep() + if (this.accumulatedThinking.length === 0) return + + const messages = this.middlewareCtx.messages + const alreadyPresent = messages.some( + (message) => + message.role === 'assistant' && message.id === this.currentMessageId, + ) + if (alreadyPresent) return + + this.messages = [ + ...messages, + { + role: 'assistant', + content: this.accumulatedContent || null, + id: this.currentMessageId ?? undefined, + createdAt: this.currentMessageCreatedAt ?? undefined, + thinking: this.accumulatedThinking, + }, + ] + this.middlewareCtx.messages = this.messages + } + /** * Extract client state (approvals and client tool results) from original messages. * This is called in the constructor BEFORE converting to ModelMessage format, @@ -2254,12 +2303,18 @@ class TextEngine< : message.content === null ? undefined : JSON.stringify(message.content) + const id = + message.id || + `snapshot_${this.runIdOverride ?? this.requestId}_${index}` + const parts = + message.role === 'assistant' && message.thinking?.length + ? modelMessageToUIMessage(message, id).parts + : undefined return { - id: - message.id || - `snapshot_${this.runIdOverride ?? this.requestId}_${index}`, + id, role: message.role, ...(content !== undefined ? { content } : {}), + ...(parts ? { parts } : {}), ...('toolCalls' in message && message.toolCalls ? { toolCalls: message.toolCalls } : {}), diff --git a/packages/ai/tests/chat.test.ts b/packages/ai/tests/chat.test.ts index dac3497b4e..a630a104dd 100644 --- a/packages/ai/tests/chat.test.ts +++ b/packages/ai/tests/chat.test.ts @@ -744,6 +744,68 @@ describe('chat()', () => { }) }) + it('preserves thinking parts on the interrupt MESSAGES_SNAPSHOT', async () => { + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.stepStarted('think-1'), + chunk(EventType.REASONING_MESSAGE_CONTENT, { + messageId: 'reasoning-1', + delta: 'Need to search.', + }), + chunk(EventType.STEP_FINISHED, { + stepName: 'think-1', + stepId: 'think-1', + content: 'Need to search.', + signature: 'sig-think-1', + }), + { + ...ev.toolStart('call_1', 'clientSearch'), + parentMessageId: 'stream-assistant', + }, + ev.toolArgs('call_1', '{"query":"test"}'), + ev.runFinished('tool_calls'), + ], + ], + }) + + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ id: 'user-1', role: 'user', content: 'Search' }], + tools: [clientTool('clientSearch')], + }) as AsyncIterable, + ) + + const snapshot = chunks.find( + (chunk) => chunk.type === EventType.MESSAGES_SNAPSHOT, + ) + expect(snapshot).toMatchObject({ + messages: expect.arrayContaining([ + expect.objectContaining({ + id: 'stream-assistant', + role: 'assistant', + parts: [ + { + type: 'thinking', + content: 'Need to search.', + signature: 'sig-think-1', + }, + { + type: 'tool-call', + id: 'call_1', + name: 'clientSearch', + arguments: '{"query":"test"}', + state: 'input-complete', + input: { query: 'test' }, + }, + ], + }), + ]), + }) + }) + it('should yield an interrupt outcome for client tools', async () => { const { adapter } = createMockAdapter({ iterations: [ @@ -3036,6 +3098,49 @@ describe('chat()', () => { expect((stepChunks[0] as any).stepName).toBeDefined() expect((stepChunks[1] as any).stepName).toBeDefined() }) + + it('should preserve STEP_FINISHED content in terminal message history', async () => { + let messages: ReadonlyArray = [] + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.stepStarted('think-1'), + chunk(EventType.STEP_FINISHED, { + stepName: 'think-1', + stepId: 'think-1', + content: 'Let me think.', + signature: 'sig-think-1', + }), + ev.textStart(), + ev.textContent('Answer!'), + ev.textEnd(), + ev.runFinished('stop'), + ], + ], + }) + + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Think about it' }], + middleware: [ + defineChatMiddleware({ + name: 'capture-message-history', + onFinish(ctx) { + messages = ctx.messages + }, + }), + ], + }) as AsyncIterable, + ) + + expect(messages.at(-1)).toMatchObject({ + role: 'assistant', + content: 'Answer!', + thinking: [{ content: 'Let me think.', signature: 'sig-think-1' }], + }) + }) }) // ========================================================================== @@ -3113,10 +3218,20 @@ describe('chat()', () => { [ ev.runStarted(), ev.stepStarted('think-1'), - { - ...ev.stepFinished('Need inventory.', 'think-1'), + chunk(EventType.REASONING_MESSAGE_CONTENT, { + messageId: 'reasoning-1', + delta: 'Need ', + }), + chunk(EventType.REASONING_MESSAGE_CONTENT, { + messageId: 'reasoning-1', + delta: 'inventory.', + }), + chunk(EventType.STEP_FINISHED, { + stepName: 'think-1', + stepId: 'think-1', + content: 'Need inventory.', signature: 'sig-think-1', - }, + }), ev.toolStart('call_1', 'getInventory'), ev.toolArgs('call_1', '{}'), ev.runFinished('tool_calls'), diff --git a/testing/e2e/fixtures/tools-test/client-tool-reasoning.json b/testing/e2e/fixtures/tools-test/client-tool-reasoning.json new file mode 100644 index 0000000000..43cca4fbb5 --- /dev/null +++ b/testing/e2e/fixtures/tools-test/client-tool-reasoning.json @@ -0,0 +1,28 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[client-tool-reasoning] run test", + "sequenceIndex": 0 + }, + "response": { + "reasoning": "I should show a notification before confirming it was displayed.", + "toolCalls": [ + { + "name": "show_notification", + "arguments": "{\"message\":\"Reasoning notification\",\"type\":\"info\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[client-tool-reasoning] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "The notification has been shown after reasoning." + } + } + ] +} diff --git a/testing/e2e/src/lib/tools-test-tools.ts b/testing/e2e/src/lib/tools-test-tools.ts index c6fa70f43d..319f121a1d 100644 --- a/testing/e2e/src/lib/tools-test-tools.ts +++ b/testing/e2e/src/lib/tools-test-tools.ts @@ -187,6 +187,11 @@ export const SCENARIO_LIST = [ { id: 'text-only', label: 'Text Only (No Tools)', category: 'basic' }, { id: 'server-tool-single', label: 'Single Server Tool', category: 'basic' }, { id: 'client-tool-single', label: 'Single Client Tool', category: 'basic' }, + { + id: 'client-tool-reasoning', + label: 'Client Tool with Reasoning', + category: 'basic', + }, { id: 'approval-tool', label: 'Approval Required Tool', category: 'basic' }, { id: 'sequence-server-client', @@ -287,6 +292,7 @@ export function getToolsForScenario(scenario: string) { return [serverTools.get_weather] case 'client-tool-single': + case 'client-tool-reasoning': return [clientToolDefinitions.show_notification] case 'server-context': diff --git a/testing/e2e/src/routes/api.tools-test.ts b/testing/e2e/src/routes/api.tools-test.ts index 4d9ff14450..d401427a2b 100644 --- a/testing/e2e/src/routes/api.tools-test.ts +++ b/testing/e2e/src/routes/api.tools-test.ts @@ -239,7 +239,12 @@ export const Route = createFileRoute('/api/tools-test')({ const adapterOptions = providerFreeScenarios.has(scenario) ? { adapter: createProviderFreeAdapter(scenario) } - : createTextAdapter('openai', undefined, aimockPort, testId) + : createTextAdapter( + 'openai', + scenario === 'client-tool-reasoning' ? 'gpt-5.2' : undefined, + aimockPort, + testId, + ) const tools = getToolsForScenario(scenario) const runtimeContext: TestRuntimeContext = diff --git a/testing/e2e/tests/tools-test/client-tool.spec.ts b/testing/e2e/tests/tools-test/client-tool.spec.ts index c9f3b58224..5d47f98ad5 100644 --- a/testing/e2e/tests/tools-test/client-tool.spec.ts +++ b/testing/e2e/tests/tools-test/client-tool.spec.ts @@ -6,6 +6,7 @@ import { getMetadata, getEventLog, getToolCalls, + getMessages, } from './helpers' /** @@ -46,6 +47,46 @@ test.describe('Client Tool E2E Tests', () => { expect(startEvents[0]?.toolName).toBe('show_notification') }) + test('client tool preserves reasoning through continuation', async ({ + page, + testId, + aimockPort, + }) => { + await selectScenario(page, 'client-tool-reasoning', testId, aimockPort) + await runTest(page) + await waitForTestComplete(page) + + const messages = await getMessages(page) + const toolCallMessage = messages.find( + (message) => + message.role === 'assistant' && + message.parts.some( + (part) => + part.type === 'tool-call' && part.name === 'show_notification', + ), + ) + + expect(toolCallMessage?.parts).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + type: 'thinking', + content: + 'I should show a notification before confirming it was displayed.', + }), + ]), + ) + + const followUpText = messages + .flatMap((message) => message.parts) + .filter((part) => part.type === 'text') + .map((part) => part.content) + .join(' ') + + expect(followUpText).toContain( + 'The notification has been shown after reasoning.', + ) + }) + test('sequential client tools execute in order', async ({ page, testId,