From 64e7481b78e39ad7aeb060f24e5049df293c7123 Mon Sep 17 00:00:00 2001 From: Srajan Asthana Date: Fri, 11 Sep 2026 13:35:23 +0530 Subject: [PATCH 1/3] chore: Persist model.message reasoning_content on session events --- .../pre/model-message-reasoning-content.md | 5 + .../src/core/runtime/AgentThread.ts | 24 +++- .../modelMessageReasoningContent.test.ts | 110 ++++++++++++++++++ 3 files changed, 138 insertions(+), 1 deletion(-) create mode 100644 .changeset/pre/model-message-reasoning-content.md create mode 100644 packages/trueforge-core/tests/core/runtime/modelMessageReasoningContent.test.ts diff --git a/.changeset/pre/model-message-reasoning-content.md b/.changeset/pre/model-message-reasoning-content.md new file mode 100644 index 000000000..c5b957c88 --- /dev/null +++ b/.changeset/pre/model-message-reasoning-content.md @@ -0,0 +1,5 @@ +--- +"@truefoundry/trueforge-core": patch +--- + +Persist model.message reasoning_content on session events by deriving it from thinking_blocks so reasoning survives reload without storing it on thread context. diff --git a/packages/trueforge-core/src/core/runtime/AgentThread.ts b/packages/trueforge-core/src/core/runtime/AgentThread.ts index 3e2c6785b..f79c3ed0d 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThread.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThread.ts @@ -239,16 +239,19 @@ function buildModelMessageEvent({ id: string; }): ModelMessageEvent { // `thinking_blocks` / `source` stay on the context message for replay; strip them from the client event. + // Derive display `reasoning_content` from thinking text so /events reload matches a folded stream + // without storing reasoning on thread context. const { role, tool_calls, thinking_blocks, source, content, ...rest } = assistantMessage; void role; - void thinking_blocks; void source; + const reasoningContent = reasoningContentFromThinkingBlocks(thinking_blocks); const event: ModelMessageEvent = { ...rest, // Tool-only completions store content: null on context (OpenAI replay) but the // SSE placeholder omits the field. Drop null/empty here so listTurnEvents JSON // matches a folded stream. ...(content && { content }), + ...(reasoningContent && { reasoning_content: reasoningContent }), tool_calls: tool_calls?.map(toEnrichedToolCall), type: EventType.MODEL_MESSAGE, id, @@ -260,6 +263,25 @@ function buildModelMessageEvent({ return event; } +/** Join non-empty `thinking` blocks for client display; skip redacted / empty blocks. */ +function reasoningContentFromThinkingBlocks( + thinking_blocks: InternalEnrichedAssistantMessage['thinking_blocks'], +): string | undefined { + if (!thinking_blocks?.length) { + return undefined; + } + const parts: string[] = []; + for (const block of thinking_blocks) { + if (block.type === 'thinking' && block.thinking.length > 0) { + parts.push(block.thinking); + } + } + if (parts.length === 0) { + return undefined; + } + return parts.join(''); +} + function validateUserMessage( message: { content: AgentInputUserMessage['content'] }, blockingOpenToolCallIds: Set, diff --git a/packages/trueforge-core/tests/core/runtime/modelMessageReasoningContent.test.ts b/packages/trueforge-core/tests/core/runtime/modelMessageReasoningContent.test.ts new file mode 100644 index 000000000..da4937eeb --- /dev/null +++ b/packages/trueforge-core/tests/core/runtime/modelMessageReasoningContent.test.ts @@ -0,0 +1,110 @@ +import type { ILLM } from '../../../src/core/llm/ILLM'; +import type { ExtendedChatCompletionChunk, RawAssistantMessageWithUsage } from '../../../src/core/llm/LLMTypes'; +import { getEmptyUsage } from '../../../src/core/llm/LLMTypes'; +import { AgentThread } from '../../../src/core/runtime/AgentThread'; +import { InternalEventType, type AgentThreadAppendContext } from '../../../src/core/runtime/AgentThread.types'; +import { NOOP_AGENT_TRACING } from '../../../src/core/tracing/NoopAgentTracing'; +import '../harnessMocks'; +import { makeSilentLogger } from '../harnessMocks'; + +const silentLogger = makeSilentLogger(); + +// eslint-disable-next-line @typescript-eslint/require-await -- async generator fixture, not awaiting I/O +async function* reasoningAnswerStream(): AsyncGenerator< + ExtendedChatCompletionChunk, + RawAssistantMessageWithUsage, + unknown +> { + yield { + id: 'chunk-1', + object: 'chat.completion.chunk', + created: 0, + model: 'test-model', + choices: [ + { + index: 0, + delta: { reasoning_content: 'step one' }, + finish_reason: null, + logprobs: null, + }, + ], + }; + yield { + id: 'chunk-2', + object: 'chat.completion.chunk', + created: 0, + model: 'test-model', + choices: [ + { + index: 0, + delta: { content: 'hello', role: 'assistant' }, + finish_reason: 'stop', + logprobs: null, + }, + ], + }; + + return { + output: { + role: 'assistant', + content: 'hello', + thinking_blocks: [ + { type: 'thinking', thinking: 'step one' }, + { type: 'thinking', thinking: 'step two' }, + { type: 'redacted_thinking', data: 'opaque' }, + ], + }, + usage: getEmptyUsage(), + finish_reason: 'stop', + }; +} + +describe('AgentThread model.message reasoning_content', () => { + it('derives event reasoning_content from thinking_blocks and leaves it off context', async () => { + const modelClient: ILLM = { + create: jest.fn().mockImplementation(() => reasoningAnswerStream()), + createNonStream: jest.fn(), + }; + + const thread = new AgentThread({ + threadId: 'main', + title: 'Main', + tracing: NOOP_AGENT_TRACING, + logger: silentLogger, + definition: { + modelClient, + instruction: 'test', + toolSets: [], + }, + context: [{ role: 'user', content: 'hi' }], + }); + + let append: AgentThreadAppendContext | undefined; + for await (const event of thread.execute({ signal: new AbortController().signal })) { + if (event.type === InternalEventType.AGENT_CONTEXT_APPEND) { + append = event; + } + } + + expect(append).toBeDefined(); + const contextMessage = append?.context[0]; + expect(contextMessage).toMatchObject({ + role: 'assistant', + content: 'hello', + thinking_blocks: [ + { type: 'thinking', thinking: 'step one' }, + { type: 'thinking', thinking: 'step two' }, + { type: 'redacted_thinking', data: 'opaque' }, + ], + }); + expect(contextMessage).not.toHaveProperty('reasoning_content'); + + const modelMessage = append?.output.find(e => e.type === 'model.message'); + expect(modelMessage).toMatchObject({ + type: 'model.message', + content: 'hello', + reasoning_content: 'step onestep two', + }); + expect(modelMessage).not.toHaveProperty('thinking_blocks'); + }); +}); From 7b22d518da4502b4d01826a6153bb4a76930d0e4 Mon Sep 17 00:00:00 2001 From: Srajan Asthana Date: Fri, 11 Sep 2026 13:57:39 +0530 Subject: [PATCH 2/3] nit --- .../pre/model-message-reasoning-content.md | 2 +- .../trueforge-core/src/core/llm/LLMTypes.ts | 4 +-- .../src/core/llm/VercelAILLM.ts | 3 ++ .../src/core/runtime/AgentThread.ts | 29 ++++--------------- .../tests/core/llm/VercelAILLM.stream.test.ts | 1 + .../modelMessageReasoningContent.test.ts | 29 +++++-------------- 6 files changed, 19 insertions(+), 49 deletions(-) diff --git a/.changeset/pre/model-message-reasoning-content.md b/.changeset/pre/model-message-reasoning-content.md index c5b957c88..af8a2af57 100644 --- a/.changeset/pre/model-message-reasoning-content.md +++ b/.changeset/pre/model-message-reasoning-content.md @@ -2,4 +2,4 @@ "@truefoundry/trueforge-core": patch --- -Persist model.message reasoning_content on session events by deriving it from thinking_blocks so reasoning survives reload without storing it on thread context. +Persist exact streamed reasoning_content on model.message session events (omit from thread context). diff --git a/packages/trueforge-core/src/core/llm/LLMTypes.ts b/packages/trueforge-core/src/core/llm/LLMTypes.ts index 065aea70e..66e880aa8 100644 --- a/packages/trueforge-core/src/core/llm/LLMTypes.ts +++ b/packages/trueforge-core/src/core/llm/LLMTypes.ts @@ -86,7 +86,7 @@ export const RawAssistantMessageSchema = ChatCompletionAssistantMessageParamSche .extend({ tool_calls: z.array(RawToolCallSchema).optional(), thinking_blocks: z.array(ThinkingBlockUnionSchema).optional(), - /** Plain-text thinking content streamed incrementally for frontend display; redundant with thinking_blocks[].thinking. */ + /** Plain-text thinking for frontend display; exact concat of streamed deltas when set on the assembled message. */ reasoning_content: z.string().optional(), /** Source of the message: which provider/model sent it (`provider_type/provider_name/model_name`). */ source: z.string().optional(), @@ -120,7 +120,7 @@ export const ExtendedChunkDeltaSchema = ChatCompletionChunkDeltaSchema.omit({ to tool_calls: z.array(ExtendedChunkDeltaToolCallSchema).optional(), /** Structured thinking blocks from the gateway; accumulated into complete blocks (with signatures) for multi-turn replay. */ thinking_blocks: z.array(ThinkingBlockUnionSchema).optional(), - /** Plain-text thinking content streamed incrementally for frontend display; not stored — redundant with thinking_blocks[].thinking. */ + /** Plain-text thinking fragment for frontend display; assembled message stores the full concat separately. */ reasoning_content: z.string().optional(), }) .openapi('ExtendedChunkDelta'); diff --git a/packages/trueforge-core/src/core/llm/VercelAILLM.ts b/packages/trueforge-core/src/core/llm/VercelAILLM.ts index da7e1a3b6..65148cc08 100644 --- a/packages/trueforge-core/src/core/llm/VercelAILLM.ts +++ b/packages/trueforge-core/src/core/llm/VercelAILLM.ts @@ -1128,6 +1128,7 @@ export async function* mapStreamToChunks({ const toolCallStates = new Map(); let nextToolIndex = 0; let accumulatedText = ''; + let accumulatedReasoning = ''; const accumulatedThinking: ThinkingBlock[] = []; const thinkingByReasoningItem = new Map(); let currentThinkingBlock: ThinkingBlock | null = null; @@ -1188,6 +1189,7 @@ export async function* mapStreamToChunks({ // signature, breaking replay while the text still streams out. currentThinkingBlock ??= openThinkingBlock(part.providerMetadata); currentThinkingBlock.thinking += part.text; + accumulatedReasoning += part.text; applyReasoningSignature({ block: currentThinkingBlock, providerMetadata: part.providerMetadata }); yield { ...makeBase(), @@ -1349,6 +1351,7 @@ export async function* mapStreamToChunks({ const output: RawAssistantMessage = { role: 'assistant', content: accumulatedText || null, + ...(accumulatedReasoning.length > 0 ? { reasoning_content: accumulatedReasoning } : {}), ...(accumulatedThinking.length > 0 ? { thinking_blocks: accumulatedThinking } : {}), ...(toolCalls.length > 0 ? { tool_calls: toolCalls } : {}), }; diff --git a/packages/trueforge-core/src/core/runtime/AgentThread.ts b/packages/trueforge-core/src/core/runtime/AgentThread.ts index f79c3ed0d..6185bb779 100644 --- a/packages/trueforge-core/src/core/runtime/AgentThread.ts +++ b/packages/trueforge-core/src/core/runtime/AgentThread.ts @@ -239,19 +239,16 @@ function buildModelMessageEvent({ id: string; }): ModelMessageEvent { // `thinking_blocks` / `source` stay on the context message for replay; strip them from the client event. - // Derive display `reasoning_content` from thinking text so /events reload matches a folded stream - // without storing reasoning on thread context. const { role, tool_calls, thinking_blocks, source, content, ...rest } = assistantMessage; void role; + void thinking_blocks; void source; - const reasoningContent = reasoningContentFromThinkingBlocks(thinking_blocks); const event: ModelMessageEvent = { ...rest, // Tool-only completions store content: null on context (OpenAI replay) but the // SSE placeholder omits the field. Drop null/empty here so listTurnEvents JSON // matches a folded stream. ...(content && { content }), - ...(reasoningContent && { reasoning_content: reasoningContent }), tool_calls: tool_calls?.map(toEnrichedToolCall), type: EventType.MODEL_MESSAGE, id, @@ -263,25 +260,6 @@ function buildModelMessageEvent({ return event; } -/** Join non-empty `thinking` blocks for client display; skip redacted / empty blocks. */ -function reasoningContentFromThinkingBlocks( - thinking_blocks: InternalEnrichedAssistantMessage['thinking_blocks'], -): string | undefined { - if (!thinking_blocks?.length) { - return undefined; - } - const parts: string[] = []; - for (const block of thinking_blocks) { - if (block.type === 'thinking' && block.thinking.length > 0) { - parts.push(block.thinking); - } - } - if (parts.length === 0) { - return undefined; - } - return parts.join(''); -} - function validateUserMessage( message: { content: AgentInputUserMessage['content'] }, blockingOpenToolCallIds: Set, @@ -1099,6 +1077,9 @@ export class AgentThread { // Hence, we resolve the underlying tool to get the tool information. resolveUnderlyingTool: true, }); + // Display-only; keep off thread context so replay uses thinking_blocks alone. + const { reasoning_content: _omitReasoning, ...assistantMessageForContext } = assistantMessage; + void _omitReasoning; const finishReason = result.value.finish_reason; const agentAssistantMessage = buildModelMessageEvent({ assistantMessage: await enrichAssistantMessage({ @@ -1137,7 +1118,7 @@ export class AgentThread { } yield* this.appendToContext({ - context: [assistantMessage], + context: [assistantMessageForContext], output: [agentAssistantMessage], currentContextUsage: currentContextUsageFromCompletion(result.value.usage), usage: result.value.usage, diff --git a/packages/trueforge-core/tests/core/llm/VercelAILLM.stream.test.ts b/packages/trueforge-core/tests/core/llm/VercelAILLM.stream.test.ts index bb15d4aed..6435bfcdd 100644 --- a/packages/trueforge-core/tests/core/llm/VercelAILLM.stream.test.ts +++ b/packages/trueforge-core/tests/core/llm/VercelAILLM.stream.test.ts @@ -269,6 +269,7 @@ describe('mapStreamToChunks', () => { expect(reasoningChunks).toHaveLength(1); expect(reasoningChunks[0]?.choices[0]?.delta.reasoning_content).toBe('step one'); expect(final.output.thinking_blocks).toEqual([{ type: 'thinking', thinking: 'step one' }]); + expect(final.output.reasoning_content).toBe('step one'); }); it('attaches signature from providerMetadata.*.reasoningEncryptedContent on reasoning-end (OpenAI)', async () => { diff --git a/packages/trueforge-core/tests/core/runtime/modelMessageReasoningContent.test.ts b/packages/trueforge-core/tests/core/runtime/modelMessageReasoningContent.test.ts index da4937eeb..4010f813d 100644 --- a/packages/trueforge-core/tests/core/runtime/modelMessageReasoningContent.test.ts +++ b/packages/trueforge-core/tests/core/runtime/modelMessageReasoningContent.test.ts @@ -48,11 +48,8 @@ async function* reasoningAnswerStream(): AsyncGenerator< output: { role: 'assistant', content: 'hello', - thinking_blocks: [ - { type: 'thinking', thinking: 'step one' }, - { type: 'thinking', thinking: 'step two' }, - { type: 'redacted_thinking', data: 'opaque' }, - ], + reasoning_content: 'step one', + thinking_blocks: [{ type: 'thinking', thinking: 'step one' }], }, usage: getEmptyUsage(), finish_reason: 'stop', @@ -60,7 +57,7 @@ async function* reasoningAnswerStream(): AsyncGenerator< } describe('AgentThread model.message reasoning_content', () => { - it('derives event reasoning_content from thinking_blocks and leaves it off context', async () => { + it('persists reasoning_content on the event and omits it from context', async () => { const modelClient: ILLM = { create: jest.fn().mockImplementation(() => reasoningAnswerStream()), createNonStream: jest.fn(), @@ -86,25 +83,13 @@ describe('AgentThread model.message reasoning_content', () => { } } - expect(append).toBeDefined(); - const contextMessage = append?.context[0]; - expect(contextMessage).toMatchObject({ - role: 'assistant', - content: 'hello', - thinking_blocks: [ - { type: 'thinking', thinking: 'step one' }, - { type: 'thinking', thinking: 'step two' }, - { type: 'redacted_thinking', data: 'opaque' }, - ], + expect(append?.context[0]).not.toHaveProperty('reasoning_content'); + expect(append?.context[0]).toMatchObject({ + thinking_blocks: [{ type: 'thinking', thinking: 'step one' }], }); - expect(contextMessage).not.toHaveProperty('reasoning_content'); const modelMessage = append?.output.find(e => e.type === 'model.message'); - expect(modelMessage).toMatchObject({ - type: 'model.message', - content: 'hello', - reasoning_content: 'step onestep two', - }); + expect(modelMessage).toMatchObject({ reasoning_content: 'step one' }); expect(modelMessage).not.toHaveProperty('thinking_blocks'); }); }); From 5bd78d309705d104d3959b2e65117a94971e5b73 Mon Sep 17 00:00:00 2001 From: Srajan Asthana Date: Fri, 11 Sep 2026 14:31:33 +0530 Subject: [PATCH 3/3] nit --- .changeset/{pre => }/model-message-reasoning-content.md | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename .changeset/{pre => }/model-message-reasoning-content.md (100%) diff --git a/.changeset/pre/model-message-reasoning-content.md b/.changeset/model-message-reasoning-content.md similarity index 100% rename from .changeset/pre/model-message-reasoning-content.md rename to .changeset/model-message-reasoning-content.md