chore: persist reasoning content on events for consitent replay

.changeset/model-message-reasoning-content.md

@@ -0,0 +1,5 @@
+---
+"@truefoundry/trueforge-core": patch
+---
+
+Persist exact streamed reasoning_content on model.message session events (omit from thread context).

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');

packages/trueforge-core/src/core/llm/VercelAILLM.ts

@@ -1135,6 +1135,7 @@ export async function* mapStreamToChunks({
   const toolCallStates = new Map<string, ToolCallState>();
   let nextToolIndex = 0;
   let accumulatedText = '';
+  let accumulatedReasoning = '';
   const accumulatedThinking: ThinkingBlock[] = [];
   const thinkingByReasoningItem = new Map<string, ThinkingBlock>();
   let currentThinkingBlock: ThinkingBlock | null = null;
@@ -1195,6 +1196,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(),
@@ -1356,6 +1358,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 } : {}),
   };

packages/trueforge-core/src/core/llm/toOpenAIChatMessage.ts

@@ -4,9 +4,10 @@ import type { LLMUserMessage } from './LLMTypes';
 
 /**
  * Maps harness context messages to the OpenAI chat message param shape.
- * For assistants: spreads the message (preserving thinking_blocks, reasoning_content,
- * and tool-call provider_specific_fields), strips only `tool_info`, and omits empty
- * `tool_calls` (OpenAI rejects `tool_calls: []`).
+ * For assistants: spreads the message (preserving thinking_blocks and tool-call
+ * provider_specific_fields), strips only `tool_info`, and omits empty `tool_calls`
+ * (OpenAI rejects `tool_calls: []`). AgentThread context omits display-only
+ * `reasoning_content`; if present on a message it is forwarded unchanged.
  * User content is rebuilt field-by-field for exactOptionalPropertyTypes.
  */
 export function toOpenAIChatMessage(msg: LLMContextMessage): ChatCompletionMessageParam {
@@ -27,7 +28,7 @@ export function toOpenAIChatMessage(msg: LLMContextMessage): ChatCompletionMessa
     const { tool_calls: _omit, ...rest } = msg;
     void _omit;
     // Assert: OpenAI SDK message type omits gateway thinking-block / provider extensions
-    // (thinking_blocks, reasoning_content, …); we intentionally forward them for replay.
+    // (thinking_blocks, …); we intentionally forward them for LLM replay.
     return rest as ChatCompletionMessageParam;
   }
 

packages/trueforge-core/src/core/runtime/AgentThread.ts

@@ -225,6 +225,15 @@ function buildSandboxCreatedEvent(info: SandboxInfo): SandboxCreatedEvent {
   };
 }
 
+/** Reasoning content is display-only; omit from context so LLM replay uses thinking_blocks alone. */
+function toContextAssistantMessage(
+  assistantMessage: InternalEnrichedAssistantMessage,
+): InternalEnrichedAssistantMessage {
+  const { reasoning_content, ...forContext } = assistantMessage;
+  void reasoning_content;
+  return forContext;
+}
+
 function buildModelMessageEvent({
   assistantMessage,
   threadId,
@@ -239,6 +248,7 @@ function buildModelMessageEvent({
   id: string;
 }): ModelMessageEvent {
   // `thinking_blocks` / `source` stay on the context message for replay; strip them from the client event.
+  // `reasoning_content` stays on the event for UI replay (exact streamed concat).
   const { role, tool_calls, thinking_blocks, source, content, ...rest } = assistantMessage;
   void role;
   void thinking_blocks;
@@ -1077,6 +1087,7 @@ export class AgentThread {
       // Hence, we resolve the underlying tool to get the tool information.
       resolveUnderlyingTool: true,
     });
+    const contextAssistantMessage = toContextAssistantMessage(assistantMessage);
     const finishReason = result.value.finish_reason;
     const agentAssistantMessage = buildModelMessageEvent({
       assistantMessage: await enrichAssistantMessage({
@@ -1115,7 +1126,7 @@ export class AgentThread {
     }
 
     yield* this.appendToContext({
-      context: [assistantMessage],
+      context: [contextAssistantMessage],
       output: [agentAssistantMessage],
       currentContextUsage: currentContextUsageFromCompletion(result.value.usage),
       usage: result.value.usage,

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 () => {

packages/trueforge-core/tests/core/runtime/modelMessageReasoningContent.test.ts

@@ -0,0 +1,95 @@
+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',
+      reasoning_content: 'step one',
+      thinking_blocks: [{ type: 'thinking', thinking: 'step one' }],
+    },
+    usage: getEmptyUsage(),
+    finish_reason: 'stop',
+  };
+}
+
+describe('AgentThread model.message reasoning_content', () => {
+  it('persists reasoning_content on the event and omits it from 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?.context[0]).not.toHaveProperty('reasoning_content');
+    expect(append?.context[0]).toMatchObject({
+      thinking_blocks: [{ type: 'thinking', thinking: 'step one' }],
+    });
+
+    const modelMessage = append?.output.find(e => e.type === 'model.message');
+    expect(modelMessage).toMatchObject({ reasoning_content: 'step one' });
+    expect(modelMessage).not.toHaveProperty('thinking_blocks');
+  });
+});