From 2dedab2a8b9daa48ee135abfcb5143c5110da4f7 Mon Sep 17 00:00:00 2001 From: jaberjaber23 Date: Fri, 1 May 2026 13:01:55 +0300 Subject: [PATCH] think persist --- crates/openfang-memory/src/session.rs | 2 +- crates/openfang-runtime/src/agent_loop.rs | 171 +++++++- .../openfang-runtime/src/drivers/anthropic.rs | 271 +++++++++++- crates/openfang-runtime/src/drivers/gemini.rs | 119 +++++- crates/openfang-runtime/src/drivers/openai.rs | 402 +++++++++++++----- crates/openfang-types/src/message.rs | 142 ++++++- 6 files changed, 969 insertions(+), 138 deletions(-) diff --git a/crates/openfang-memory/src/session.rs b/crates/openfang-memory/src/session.rs index a4403534..f0d4c373 100644 --- a/crates/openfang-memory/src/session.rs +++ b/crates/openfang-memory/src/session.rs @@ -586,7 +586,7 @@ impl SessionStore { ContentBlock::Image { media_type, .. } => { text_parts.push(format!("[image: {media_type}]")); } - ContentBlock::Thinking { thinking } => { + ContentBlock::Thinking { thinking, .. } => { text_parts.push(format!( "[thinking: {}]", openfang_types::truncate_str(thinking, 200) diff --git a/crates/openfang-runtime/src/agent_loop.rs b/crates/openfang-runtime/src/agent_loop.rs index b01e6b66..7f5f05ed 100644 --- a/crates/openfang-runtime/src/agent_loop.rs +++ b/crates/openfang-runtime/src/agent_loop.rs @@ -109,6 +109,70 @@ fn append_tool_error_guidance(tool_result_blocks: &mut Vec) { } } +/// Build an assistant message that preserves Thinking blocks alongside the +/// final visible text. +/// +/// Issue #1098 — thinking-model state preservation. When the LLM response +/// contains `ContentBlock::Thinking` (Anthropic extended thinking with +/// signatures, Gemini 2.5+ thoughts, OpenAI-compat reasoning_content, +/// MiniMax/Qwen inline `` blocks), the prior code stored only the +/// final text via `Message::assistant(text)` — discarding all reasoning +/// state. On the next turn the model re-derived its answer from scratch +/// and quality degraded. +/// +/// This helper preserves the full block list whenever any Thinking block is +/// present, otherwise returns the legacy `Message::assistant(text)` form so +/// downstream consumers (channel formatters, JSONL mirrors, embeddings) keep +/// working without changes. +/// +/// Note: we deliberately replace any visible Text blocks in `response_blocks` +/// with `final_text` so that any post-processing the agent loop applied +/// (phantom-action recovery, accumulated_text fallback, EmptyResponse guard +/// stub) is reflected in the persisted message. +fn build_assistant_message_preserving_thinking( + response_blocks: &[ContentBlock], + final_text: &str, +) -> Message { + let has_thinking = response_blocks + .iter() + .any(|b| matches!(b, ContentBlock::Thinking { .. })); + if !has_thinking { + return Message::assistant(final_text.to_string()); + } + + // Preserve order: Thinking blocks first (in original order), then a + // single Text block carrying `final_text`. Tool blocks aren't expected + // here (StopReason::EndTurn path), but copy them through if present so + // we don't drop information. + let mut blocks: Vec = Vec::with_capacity(response_blocks.len() + 1); + let mut emitted_text = false; + for b in response_blocks { + match b { + ContentBlock::Thinking { .. } => blocks.push(b.clone()), + ContentBlock::Text { .. } if !emitted_text => { + blocks.push(ContentBlock::Text { + text: final_text.to_string(), + provider_metadata: None, + }); + emitted_text = true; + } + ContentBlock::Text { .. } => { + // Drop additional text blocks — final_text already captures + // the canonical visible message. + } + other => blocks.push(other.clone()), + } + } + if !emitted_text && !final_text.is_empty() { + blocks.push(ContentBlock::Text { + text: final_text.to_string(), + provider_metadata: None, + }); + } + + Message::assistant_with_blocks(blocks) +} + /// Strip a provider prefix from a model ID before sending to the API. /// /// Many models are stored as `provider/org/model` (e.g. `openrouter/google/gemini-2.5-flash`) @@ -605,7 +669,16 @@ pub async fn run_agent_loop( }; final_response = text.clone(); - session.messages.push(Message::assistant(text)); + // Issue #1098: persist Thinking blocks alongside the text so + // reasoning models retain state across turns. When the + // response carries any Thinking content (Anthropic extended + // thinking, Gemini 2.5 thought signatures, DeepSeek-R1/Qwen3 + // `reasoning_content`, MiniMax inline ``), save the + // full content blocks; otherwise fall back to the legacy + // Text shape so existing sessions/snapshots stay readable. + let assistant_msg = + build_assistant_message_preserving_thinking(&response.content, &text); + session.messages.push(assistant_msg); // Prune NO_REPLY heartbeat turns to save context budget crate::session_repair::prune_heartbeat_turns(&mut session.messages, 10); @@ -1798,7 +1871,13 @@ pub async fn run_agent_loop_streaming( text }; final_response = text.clone(); - session.messages.push(Message::assistant(text)); + // Issue #1098: preserve Thinking blocks (with Anthropic + // signatures / Gemini thought signatures / inline-think / + // reasoning_content) on the persisted assistant turn. See + // build_assistant_message_preserving_thinking for details. + let assistant_msg = + build_assistant_message_preserving_thinking(&response.content, &text); + session.messages.push(assistant_msg); // Prune NO_REPLY heartbeat turns to save context budget crate::session_repair::prune_heartbeat_turns(&mut session.messages, 10); @@ -3092,6 +3171,94 @@ mod tests { assert_eq!(MAX_ITERATIONS, 50); } + /// Issue #1098: when a response carries Thinking blocks, the persisted + /// assistant turn must keep them so the next turn round-trips reasoning + /// state to the model. + #[test] + fn test_build_assistant_message_preserves_thinking() { + let response_blocks = vec![ + ContentBlock::Thinking { + thinking: "Let me reason carefully...".to_string(), + signature: Some("sig_anthropic_xyz".to_string()), + provider_metadata: Some(serde_json::json!({ + "format": "anthropic_extended_thinking" + })), + }, + ContentBlock::Text { + text: "Initial response text".to_string(), + provider_metadata: None, + }, + ]; + // Final text might differ from the original Text block (phantom-action + // recovery / synthesis fallback rewrites it). The helper should adopt + // final_text into the persisted Text block. + let final_text = "Initial response text"; + let msg = build_assistant_message_preserving_thinking(&response_blocks, final_text); + assert_eq!(msg.role, Role::Assistant); + let blocks = match &msg.content { + MessageContent::Blocks(b) => b, + other => panic!("expected blocks, got {other:?}"), + }; + assert_eq!(blocks.len(), 2, "must preserve thinking + text"); + match &blocks[0] { + ContentBlock::Thinking { + thinking, signature, .. + } => { + assert_eq!(thinking, "Let me reason carefully..."); + assert_eq!(signature.as_deref(), Some("sig_anthropic_xyz")); + } + _ => panic!("expected Thinking first"), + } + match &blocks[1] { + ContentBlock::Text { text, .. } => assert_eq!(text, "Initial response text"), + _ => panic!("expected Text second"), + } + } + + /// Without thinking, fall back to the legacy `Message::assistant(text)` + /// shape so existing JSONL mirrors and embeddings keep working. + #[test] + fn test_build_assistant_message_no_thinking_is_plain_text() { + let response_blocks = vec![ContentBlock::Text { + text: "Hi.".to_string(), + provider_metadata: None, + }]; + let msg = build_assistant_message_preserving_thinking(&response_blocks, "Hi."); + match msg.content { + MessageContent::Text(t) => assert_eq!(t, "Hi."), + _ => panic!("expected plain text content for non-thinking responses"), + } + } + + /// Final text supplied by the loop (e.g. recovery stub) must replace + /// the original text part — the persisted message reflects what was + /// actually returned to the user, not the raw LLM output. + #[test] + fn test_build_assistant_message_final_text_replaces_original_text() { + let response_blocks = vec![ + ContentBlock::Thinking { + thinking: "deliberation".to_string(), + signature: None, + provider_metadata: Some(serde_json::json!({"format": "inline_think"})), + }, + ContentBlock::Text { + text: "raw LLM output".to_string(), + provider_metadata: None, + }, + ]; + let final_text = "[Task completed — recovered after empty response.]"; + let msg = build_assistant_message_preserving_thinking(&response_blocks, final_text); + let blocks = match &msg.content { + MessageContent::Blocks(b) => b, + _ => panic!("expected blocks"), + }; + let saved_text = blocks.iter().find_map(|b| match b { + ContentBlock::Text { text, .. } => Some(text.as_str()), + _ => None, + }); + assert_eq!(saved_text, Some(final_text)); + } + #[test] fn test_retry_constants() { assert_eq!(MAX_RETRIES, 3); diff --git a/crates/openfang-runtime/src/drivers/anthropic.rs b/crates/openfang-runtime/src/drivers/anthropic.rs index 6bac8b31..e6f10fe5 100644 --- a/crates/openfang-runtime/src/drivers/anthropic.rs +++ b/crates/openfang-runtime/src/drivers/anthropic.rs @@ -84,6 +84,15 @@ enum ApiContentBlock { #[serde(skip_serializing_if = "std::ops::Not::not")] is_error: bool, }, + /// Extended-thinking block echoed back to the API. + /// + /// Anthropic requires the original `signature` to be returned verbatim + /// alongside the `thinking` text on subsequent turns; otherwise the + /// model loses its prior reasoning state. Without `signature` the API + /// rejects the block, so we omit thinking blocks that arrive without + /// one (e.g. legacy sessions saved before this field was tracked). + #[serde(rename = "thinking")] + Thinking { thinking: String, signature: String }, } #[derive(Debug, Serialize)] @@ -120,8 +129,16 @@ enum ResponseContentBlock { name: String, input: serde_json::Value, }, + /// Extended-thinking block from Anthropic. The `signature` is opaque + /// to us but MUST be persisted and echoed back on the next request, + /// otherwise the API rejects the resubmitted thinking block and the + /// model loses its reasoning state. #[serde(rename = "thinking")] - Thinking { thinking: String }, + Thinking { + thinking: String, + #[serde(default)] + signature: Option, + }, } #[derive(Debug, Deserialize)] @@ -144,7 +161,14 @@ struct ApiErrorDetail { /// Accumulator for content blocks during streaming. enum ContentBlockAccum { Text(String), - Thinking(String), + /// Extended thinking — text plus an opaque signature delivered as + /// `signature_delta` events (or as a single field on `content_block_stop` + /// for older API versions). The signature is required to round-trip + /// thinking blocks on subsequent turns. + Thinking { + thinking: String, + signature: String, + }, ToolUse { id: String, name: String, @@ -412,7 +436,16 @@ impl LlmDriver for AnthropicDriver { }); } "thinking" => { - blocks.push(ContentBlockAccum::Thinking(String::new())); + // Some API versions ship the signature on + // content_block_start instead of as a delta. + let initial_sig = block["signature"] + .as_str() + .unwrap_or("") + .to_string(); + blocks.push(ContentBlockAccum::Thinking { + thinking: String::new(), + signature: initial_sig, + }); } _ => {} } @@ -452,11 +485,33 @@ impl LlmDriver for AnthropicDriver { } } "thinking_delta" => { - if let Some(thinking) = delta["thinking"].as_str() { - if let Some(ContentBlockAccum::Thinking(ref mut t)) = - blocks.get_mut(block_idx) + if let Some(t) = delta["thinking"].as_str() { + if let Some(ContentBlockAccum::Thinking { + thinking: ref mut buf, + .. + }) = blocks.get_mut(block_idx) { - t.push_str(thinking); + buf.push_str(t); + } + // Forward to UI as ThinkingDelta event so dashboards can show reasoning. + let _ = tx + .send(StreamEvent::ThinkingDelta { + text: t.to_string(), + }) + .await; + } + } + "signature_delta" => { + // Anthropic streams the thinking signature + // as its own delta type; concatenate any + // partial pieces into the accumulator. + if let Some(sig) = delta["signature"].as_str() { + if let Some(ContentBlockAccum::Thinking { + ref mut signature, + .. + }) = blocks.get_mut(block_idx) + { + signature.push_str(sig); } } } @@ -512,8 +567,26 @@ impl LlmDriver for AnthropicDriver { provider_metadata: None, }); } - ContentBlockAccum::Thinking(thinking) => { - content.push(ContentBlock::Thinking { thinking }); + ContentBlockAccum::Thinking { + thinking, + signature, + } => { + // Drop empty thinking blocks (rare, but happens if the + // stream is interrupted mid-block). Always keep the + // signature when present — it's required to round-trip. + if !thinking.is_empty() || !signature.is_empty() { + content.push(ContentBlock::Thinking { + thinking, + signature: if signature.is_empty() { + None + } else { + Some(signature) + }, + provider_metadata: Some(serde_json::json!({ + "format": "anthropic_extended_thinking" + })), + }); + } } ContentBlockAccum::ToolUse { id, @@ -615,7 +688,28 @@ fn convert_message(msg: &Message) -> ApiMessage { content: content.clone(), is_error: *is_error, }), - ContentBlock::Thinking { .. } => None, + ContentBlock::Thinking { + thinking, + signature, + .. + } => { + // Anthropic's extended-thinking spec requires the + // verbatim `signature` to accompany any thinking block + // resubmitted in conversation history. Without one, + // the API rejects the request, so we silently drop + // legacy thinking blocks (saved before signature + // tracking) instead of round-tripping them. + signature.as_ref().and_then(|sig| { + if sig.is_empty() { + None + } else { + Some(ApiContentBlock::Thinking { + thinking: thinking.clone(), + signature: sig.clone(), + }) + } + }) + } ContentBlock::Unknown => None, }) .collect(); @@ -651,8 +745,17 @@ fn convert_response(api: ApiResponse) -> CompletionResponse { }); tool_calls.push(ToolCall { id, name, input }); } - ResponseContentBlock::Thinking { thinking } => { - content.push(ContentBlock::Thinking { thinking }); + ResponseContentBlock::Thinking { + thinking, + signature, + } => { + content.push(ContentBlock::Thinking { + thinking, + signature, + provider_metadata: Some(serde_json::json!({ + "format": "anthropic_extended_thinking" + })), + }); } } } @@ -772,4 +875,148 @@ mod tests { panic!("Expected Blocks content"); } } + + /// Issue #1098: Anthropic extended-thinking blocks must round-trip + /// through the driver — the inbound response carries a `signature` that + /// MUST be echoed verbatim on the next request, otherwise the API + /// rejects the resubmitted thinking block and the model loses prior + /// reasoning state. + #[test] + fn test_thinking_block_signature_round_trip() { + // Step 1: API delivers a thinking block with signature + let api_response = ApiResponse { + content: vec![ + ResponseContentBlock::Thinking { + thinking: "Let me carefully consider this problem...".to_string(), + signature: Some("WaUjzkypQ2mUEVM36O2TxuC".to_string()), + }, + ResponseContentBlock::Text { + text: "The answer is 42.".to_string(), + }, + ], + stop_reason: "end_turn".to_string(), + usage: ApiUsage { + input_tokens: 100, + output_tokens: 50, + }, + }; + let response = convert_response(api_response); + assert_eq!(response.content.len(), 2); + + // Step 2: Verify the signature reached the ContentBlock + let thinking_block = &response.content[0]; + match thinking_block { + ContentBlock::Thinking { + thinking, + signature, + .. + } => { + assert_eq!(thinking, "Let me carefully consider this problem..."); + assert_eq!(signature.as_deref(), Some("WaUjzkypQ2mUEVM36O2TxuC")); + } + _ => panic!("expected Thinking content block"), + } + + // Step 3: Now feed the assistant turn back into the driver as if + // it were prior conversation history (next user turn). The signature + // must survive into the outbound API request. + let assistant_msg = Message { + role: Role::Assistant, + content: MessageContent::Blocks(response.content.clone()), + }; + let api_msg = convert_message(&assistant_msg); + let blocks = match api_msg.content { + ApiContent::Blocks(b) => b, + _ => panic!("expected Blocks content"), + }; + + // The Thinking block must appear in the outbound payload with its signature. + let mut found_thinking = false; + for block in &blocks { + if let ApiContentBlock::Thinking { + thinking, + signature, + } = block + { + assert_eq!(thinking, "Let me carefully consider this problem..."); + assert_eq!(signature, "WaUjzkypQ2mUEVM36O2TxuC"); + found_thinking = true; + } + } + assert!( + found_thinking, + "outbound API request must include the thinking block with signature" + ); + + // Step 4: Verify on-the-wire JSON shape (`type=thinking`, `signature` present). + let outbound_json = serde_json::to_value(&blocks).unwrap(); + let arr = outbound_json.as_array().unwrap(); + let thinking_json = arr + .iter() + .find(|v| v["type"] == "thinking") + .expect("thinking block in JSON"); + assert_eq!(thinking_json["signature"], "WaUjzkypQ2mUEVM36O2TxuC"); + assert_eq!( + thinking_json["thinking"], + "Let me carefully consider this problem..." + ); + } + + /// Legacy thinking blocks saved before signature tracking should NOT + /// be replayed — Anthropic rejects thinking blocks without signatures. + #[test] + fn test_thinking_block_without_signature_dropped_outbound() { + let assistant_msg = Message { + role: Role::Assistant, + content: MessageContent::Blocks(vec![ + ContentBlock::Thinking { + thinking: "old reasoning from before sig tracking".to_string(), + signature: None, + provider_metadata: None, + }, + ContentBlock::Text { + text: "Hello.".to_string(), + provider_metadata: None, + }, + ]), + }; + let api_msg = convert_message(&assistant_msg); + let blocks = match api_msg.content { + ApiContent::Blocks(b) => b, + _ => panic!("expected Blocks content"), + }; + // The legacy thinking block must be dropped (no sig = API would 400). + for block in &blocks { + assert!( + !matches!(block, ApiContentBlock::Thinking { .. }), + "thinking block without signature must be dropped" + ); + } + // The text part is still preserved. + assert!(blocks + .iter() + .any(|b| matches!(b, ApiContentBlock::Text { .. }))); + } + + /// Streaming path: signature_delta events accumulate into the final block. + #[test] + fn test_thinking_block_serde_with_signature_field() { + // Verify the API response wire format is parsed correctly. + let json = serde_json::json!({ + "type": "thinking", + "thinking": "step 1, step 2", + "signature": "abc123" + }); + let block: ResponseContentBlock = serde_json::from_value(json).unwrap(); + match block { + ResponseContentBlock::Thinking { + thinking, + signature, + } => { + assert_eq!(thinking, "step 1, step 2"); + assert_eq!(signature.as_deref(), Some("abc123")); + } + _ => panic!("expected Thinking response block"), + } + } } diff --git a/crates/openfang-runtime/src/drivers/gemini.rs b/crates/openfang-runtime/src/drivers/gemini.rs index cac4c3a3..9efc617b 100644 --- a/crates/openfang-runtime/src/drivers/gemini.rs +++ b/crates/openfang-runtime/src/drivers/gemini.rs @@ -321,7 +321,39 @@ fn convert_messages( }, }); } - ContentBlock::Thinking { .. } => {} + ContentBlock::Thinking { + thinking, + provider_metadata, + .. + } => { + // Issue #1098: preserve Gemini 2.5+ thought parts + // when the upstream model originally emitted them. + // Most Gemini state actually rides on the + // thoughtSignature attached to text/tool_use + // parts above, but we round-trip the visible + // thinking text + sig as a `Thought` part too + // so the model's internal state is fully + // preserved. Other providers' thinking blocks + // are dropped here (they have their own + // outbound paths in the OpenAI/Anthropic + // drivers). + let format = provider_metadata + .as_ref() + .and_then(|m| m.get("format")) + .and_then(|v| v.as_str()); + if format == Some("gemini_thought") && !thinking.is_empty() { + let sig = provider_metadata + .as_ref() + .and_then(|m| m.get("thought_signature")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + parts.push(GeminiPart::Thought { + text: thinking.clone(), + thought: true, + thought_signature: sig, + }); + } + } _ => {} } } @@ -561,12 +593,29 @@ fn convert_response(resp: GeminiResponse) -> Result { + GeminiPart::Thought { + text, + thought_signature, + .. + } => { // Gemini 2.5+ thinking parts — internal reasoning. // Store as Thinking content block so the UI can - // optionally display it (like blocks). + // optionally display it. Issue #1098: preserve the + // part-level `thoughtSignature` in `provider_metadata` + // (and on subsequent text/tool_use parts) so the + // model retains state across turns. if !text.is_empty() { - content.push(ContentBlock::Thinking { thinking: text }); + let provider_metadata = thought_signature.map(|sig| { + serde_json::json!({ + "format": "gemini_thought", + "thought_signature": sig, + }) + }); + content.push(ContentBlock::Thinking { + thinking: text, + signature: None, + provider_metadata, + }); } } GeminiPart::InlineData { .. } | GeminiPart::FunctionResponse { .. } => { @@ -788,6 +837,10 @@ impl LlmDriver for GeminiDriver { let mut text_content = String::new(); // Thought signature for accumulated text content (last one wins) let mut text_thought_sig: Option = None; + // Accumulated thought (Gemini 2.5+) text + signature, for + // round-tripping reasoning state across turns (issue #1098). + let mut thought_text = String::new(); + let mut thought_sig: Option = None; // Track function calls: (name, args_json, thought_signature) let mut fn_calls: Vec<(String, serde_json::Value, Option)> = Vec::new(); let mut finish_reason: Option = None; @@ -894,17 +947,26 @@ impl LlmDriver for GeminiDriver { thought_signature.clone(), )); } - GeminiPart::Thought { ref text, .. } => { + GeminiPart::Thought { + ref text, + ref thought_signature, + .. + } => { // Gemini 2.5+ thinking chunk — emit as // thinking delta so UIs can optionally - // show it; do NOT mix into text_content. + // show it; accumulate the text + sig + // for later persistence (issue #1098). if !text.is_empty() { + thought_text.push_str(text); let _ = tx .send(StreamEvent::ThinkingDelta { text: text.clone(), }) .await; } + if thought_signature.is_some() { + thought_sig = thought_signature.clone(); + } } GeminiPart::InlineData { .. } | GeminiPart::FunctionResponse { .. } => {} @@ -985,14 +1047,24 @@ impl LlmDriver for GeminiDriver { thought_signature.clone(), )); } - GeminiPart::Thought { ref text, .. } - if !text.is_empty() => + GeminiPart::Thought { + ref text, + ref thought_signature, + .. + } if !text.is_empty() + || thought_signature.is_some() => { - let _ = tx - .send(StreamEvent::ThinkingDelta { - text: text.clone(), - }) - .await; + if !text.is_empty() { + thought_text.push_str(text); + let _ = tx + .send(StreamEvent::ThinkingDelta { + text: text.clone(), + }) + .await; + } + if thought_signature.is_some() { + thought_sig = thought_signature.clone(); + } } _ => {} } @@ -1034,6 +1106,25 @@ impl LlmDriver for GeminiDriver { let mut content = Vec::new(); let mut tool_calls = Vec::new(); + // Issue #1098: persist any accumulated Thought parts (Gemini + // 2.5+ thinking) so reasoning state round-trips on the next + // turn. The thoughtSignature also rides on text/tool_use + // parts below; this Thinking block carries the human-readable + // reasoning text for UI display + audit. + if !thought_text.is_empty() || thought_sig.is_some() { + let provider_metadata = thought_sig.as_ref().map(|sig| { + serde_json::json!({ + "format": "gemini_thought", + "thought_signature": sig, + }) + }); + content.push(ContentBlock::Thinking { + thinking: thought_text, + signature: None, + provider_metadata, + }); + } + if !text_content.is_empty() { let provider_metadata = text_thought_sig.map(|sig| serde_json::json!({ "thought_signature": sig })); @@ -1994,7 +2085,7 @@ mod tests { // Should have a Thinking block and a Text block assert_eq!(completion.content.len(), 2); match &completion.content[0] { - ContentBlock::Thinking { thinking } => { + ContentBlock::Thinking { thinking, .. } => { assert_eq!(thinking, "Let me reason..."); } _ => panic!("Expected Thinking block, got {:?}", completion.content[0]), diff --git a/crates/openfang-runtime/src/drivers/openai.rs b/crates/openfang-runtime/src/drivers/openai.rs index 967c26b9..554e14b5 100644 --- a/crates/openfang-runtime/src/drivers/openai.rs +++ b/crates/openfang-runtime/src/drivers/openai.rs @@ -289,6 +289,135 @@ fn strip_trailing_empty_assistant(messages: &mut Vec) { } } +/// Assemble an outbound assistant `OaiMessage` from `ContentBlock`s, replaying +/// any `Thinking` blocks in the format the upstream model originally emitted. +/// +/// This is the fix for issue #1098 — thinking-model state preservation. +/// Without this, `...` and `reasoning_content` are stripped on +/// the next turn so the model loses its prior reasoning trace and re-derives +/// the answer (degrading quality). We honour `provider_metadata.format`: +/// +/// - `"reasoning_content"` → emitted on the OpenAI `reasoning_content` field +/// (DeepSeek-R1, Qwen3, MiniMax M2 via LM Studio/Ollama) +/// - `"inline_think"` → wrapped in `...` and prepended to +/// the visible content (MiniMax M2.5, Llama-3.3-think variants) +/// - missing/other → fall back to the legacy Moonshot/Kimi behaviour +/// (only emit `reasoning_content` when `needs_reasoning_content()` is true) +fn assemble_assistant_message( + blocks: &[ContentBlock], + model: &str, + driver: &OpenAIDriver, +) -> OaiMessage { + let mut text_parts: Vec = Vec::new(); + let mut tool_calls: Vec = Vec::new(); + let mut reasoning_field: Option = None; + let mut inline_think: Option = None; + + for block in blocks { + match block { + ContentBlock::Text { text, .. } => text_parts.push(text.clone()), + ContentBlock::ToolUse { + id, name, input, .. + } => { + tool_calls.push(OaiToolCall { + id: id.clone(), + call_type: "function".to_string(), + function: OaiFunction { + name: name.clone(), + arguments: serde_json::to_string(input).unwrap_or_default(), + }, + }); + } + ContentBlock::Thinking { + thinking, + provider_metadata, + .. + } => { + if thinking.is_empty() { + continue; + } + let format = provider_metadata + .as_ref() + .and_then(|m| m.get("format")) + .and_then(|v| v.as_str()); + match format { + Some("inline_think") => { + // MiniMax / models trained to expect `` in + // historical assistant messages. Concatenate + // multiple thinking blocks if present. + let entry = format!("{thinking}"); + match &mut inline_think { + Some(existing) => existing.push_str(&entry), + None => inline_think = Some(entry), + } + } + Some("reasoning_content") => { + // DeepSeek-R1 / Qwen3 / OpenAI-compat servers that + // expose a separate `reasoning_content` field. + match &mut reasoning_field { + Some(existing) => existing.push_str(thinking), + None => reasoning_field = Some(thinking.clone()), + } + } + _ => { + // Unknown format — preserve as inline_think since it's + // safe (visible to the model as ordinary text). The + // legacy Moonshot path overrides this below. + let entry = format!("{thinking}"); + match &mut inline_think { + Some(existing) => existing.push_str(&entry), + None => inline_think = Some(entry), + } + } + } + } + _ => {} + } + } + + // Build the visible content by prepending inline_think (if any). + let mut visible = String::new(); + if let Some(it) = inline_think.as_ref() { + visible.push_str(it); + } + if !text_parts.is_empty() { + visible.push_str(&text_parts.join("")); + } + + let has_tool_calls = !tool_calls.is_empty(); + let needs_reasoning = driver.needs_reasoning_content(model); + + // Final reasoning_content field: the per-block format hint wins; otherwise + // fall back to legacy Moonshot/Kimi behaviour (empty string when needed). + let reasoning_content = if reasoning_field.is_some() { + reasoning_field + } else if needs_reasoning { + Some(String::new()) + } else { + None + }; + + OaiMessage { + role: "assistant".to_string(), + content: if visible.is_empty() { + if has_tool_calls { + Some(OaiMessageContent::Text(String::new())) + } else { + None + } + } else { + Some(OaiMessageContent::Text(visible)) + }, + tool_calls: if tool_calls.is_empty() { + None + } else { + Some(tool_calls) + }, + tool_call_id: None, + reasoning_content, + } +} + #[async_trait] impl LlmDriver for OpenAIDriver { async fn complete(&self, request: CompletionRequest) -> Result { @@ -384,59 +513,8 @@ impl LlmDriver for OpenAIDriver { } } (Role::Assistant, MessageContent::Blocks(blocks)) => { - let mut text_parts = Vec::new(); - let mut tool_calls = Vec::new(); - let mut reasoning_text = String::new(); - for block in blocks { - match block { - ContentBlock::Text { text, .. } => text_parts.push(text.clone()), - ContentBlock::ToolUse { - id, name, input, .. - } => { - tool_calls.push(OaiToolCall { - id: id.clone(), - call_type: "function".to_string(), - function: OaiFunction { - name: name.clone(), - arguments: serde_json::to_string(input).unwrap_or_default(), - }, - }); - } - ContentBlock::Thinking { thinking, .. } => { - reasoning_text = thinking.clone(); - } - _ => {} - } - } - let has_tool_calls = !tool_calls.is_empty(); - let needs_reasoning = self.needs_reasoning_content(&request.model); - oai_messages.push(OaiMessage { - role: "assistant".to_string(), - content: if text_parts.is_empty() { - if has_tool_calls { - Some(OaiMessageContent::Text(String::new())) - } else { - None - } - } else { - Some(OaiMessageContent::Text(text_parts.join(""))) - }, - tool_calls: if tool_calls.is_empty() { - None - } else { - Some(tool_calls) - }, - tool_call_id: None, - reasoning_content: if needs_reasoning { - Some(if reasoning_text.is_empty() { - String::new() - } else { - reasoning_text - }) - } else { - None - }, - }); + let assembled = assemble_assistant_message(blocks, &request.model, self); + oai_messages.push(assembled); } _ => {} } @@ -650,8 +728,15 @@ impl LlmDriver for OpenAIDriver { len = reasoning.len(), "Captured reasoning_content from response" ); + // Mark the format so the outbound path knows to re-emit + // this as a `reasoning_content` field rather than as + // inline `` tags. Issue #1098. content.push(ContentBlock::Thinking { thinking: reasoning.clone(), + signature: None, + provider_metadata: Some(serde_json::json!({ + "format": "reasoning_content" + })), }); } } @@ -664,8 +749,14 @@ impl LlmDriver for OpenAIDriver { if let Some(think_text) = thinking { // Only add if we didn't already get reasoning_content if choice.message.reasoning_content.is_none() { + // Mark the format so we re-emit as inline `` + // tags on the next turn (MiniMax/M2.5 style). content.push(ContentBlock::Thinking { thinking: think_text, + signature: None, + provider_metadata: Some(serde_json::json!({ + "format": "inline_think" + })), }); } } @@ -692,7 +783,7 @@ impl LlmDriver for OpenAIDriver { let thinking_text = content .iter() .find_map(|b| match b { - ContentBlock::Thinking { thinking } => Some(thinking.as_str()), + ContentBlock::Thinking { thinking, .. } => Some(thinking.as_str()), _ => None, }) .unwrap_or(""); @@ -841,59 +932,8 @@ impl LlmDriver for OpenAIDriver { } } (Role::Assistant, MessageContent::Blocks(blocks)) => { - let mut text_parts = Vec::new(); - let mut tool_calls_out = Vec::new(); - let mut reasoning_text = String::new(); - for block in blocks { - match block { - ContentBlock::Text { text, .. } => text_parts.push(text.clone()), - ContentBlock::ToolUse { - id, name, input, .. - } => { - tool_calls_out.push(OaiToolCall { - id: id.clone(), - call_type: "function".to_string(), - function: OaiFunction { - name: name.clone(), - arguments: serde_json::to_string(input).unwrap_or_default(), - }, - }); - } - ContentBlock::Thinking { thinking, .. } => { - reasoning_text = thinking.clone(); - } - _ => {} - } - } - let has_tool_calls = !tool_calls_out.is_empty(); - let needs_reasoning = self.needs_reasoning_content(&request.model); - oai_messages.push(OaiMessage { - role: "assistant".to_string(), - content: if text_parts.is_empty() { - if has_tool_calls { - Some(OaiMessageContent::Text(String::new())) - } else { - None - } - } else { - Some(OaiMessageContent::Text(text_parts.join(""))) - }, - tool_calls: if tool_calls_out.is_empty() { - None - } else { - Some(tool_calls_out) - }, - tool_call_id: None, - reasoning_content: if needs_reasoning { - Some(if reasoning_text.is_empty() { - String::new() - } else { - reasoning_text - }) - } else { - None - }, - }); + let assembled = assemble_assistant_message(blocks, &request.model, self); + oai_messages.push(assembled); } _ => {} } @@ -1292,8 +1332,15 @@ impl LlmDriver for OpenAIDriver { // Add reasoning/thinking content if present if !reasoning_content.is_empty() { + // Mark format so outbound path replays this as + // `reasoning_content` (DeepSeek-R1, Qwen3, MiniMax via + // LM Studio/Ollama). Issue #1098. content.push(ContentBlock::Thinking { thinking: reasoning_content.clone(), + signature: None, + provider_metadata: Some(serde_json::json!({ + "format": "reasoning_content" + })), }); } @@ -1303,8 +1350,14 @@ impl LlmDriver for OpenAIDriver { if let Some(think_text) = thinking { // Only add if we didn't already get reasoning_content if reasoning_content.is_empty() { + // Mark as inline-think so the next outbound turn + // re-emits the content wrapped in `...`. content.push(ContentBlock::Thinking { thinking: think_text, + signature: None, + provider_metadata: Some(serde_json::json!({ + "format": "inline_think" + })), }); } } @@ -1329,7 +1382,7 @@ impl LlmDriver for OpenAIDriver { let thinking_text = content .iter() .find_map(|b| match b { - ContentBlock::Thinking { thinking } => Some(thinking.as_str()), + ContentBlock::Thinking { thinking, .. } => Some(thinking.as_str()), _ => None, }) .unwrap_or(""); @@ -1894,4 +1947,137 @@ mod tests { let url = driver.chat_url("moonshot-v1-128k"); assert_eq!(url, "https://api.moonshot.ai/v1/chat/completions"); } + + // ── issue #1098: thinking-block round-trip ──────────────────────── + + /// Inline `` blocks captured on ingress must be re-emitted in + /// historical assistant turns so MiniMax-style models retain reasoning + /// state across turns. + #[test] + fn test_assemble_assistant_replays_inline_think() { + let driver = OpenAIDriver::new( + "test".to_string(), + "https://api.minimax.chat/v1".to_string(), + ); + let blocks = vec![ + ContentBlock::Thinking { + thinking: "step-by-step reasoning".to_string(), + signature: None, + provider_metadata: Some(serde_json::json!({"format": "inline_think"})), + }, + ContentBlock::Text { + text: "Hello, user.".to_string(), + provider_metadata: None, + }, + ]; + let msg = assemble_assistant_message(&blocks, "minimax-m2.5", &driver); + let content = match msg.content { + Some(OaiMessageContent::Text(t)) => t, + _ => panic!("expected text content"), + }; + assert_eq!( + content, "step-by-step reasoningHello, user.", + "inline_think must be re-emitted as wrapping prepended to text" + ); + // No reasoning_content field should be set for non-Moonshot models. + assert!(msg.reasoning_content.is_none()); + } + + /// `reasoning_content`-flavoured Thinking blocks must re-emit on the + /// `reasoning_content` field, NOT inline (DeepSeek-R1, Qwen3, MiniMax M2 + /// via LM Studio/Ollama). + #[test] + fn test_assemble_assistant_replays_reasoning_content_field() { + let driver = OpenAIDriver::new( + "test".to_string(), + "https://api.deepseek.com/v1".to_string(), + ); + let blocks = vec![ + ContentBlock::Thinking { + thinking: "internal chain-of-thought".to_string(), + signature: None, + provider_metadata: Some(serde_json::json!({"format": "reasoning_content"})), + }, + ContentBlock::Text { + text: "answer".to_string(), + provider_metadata: None, + }, + ]; + let msg = assemble_assistant_message(&blocks, "deepseek-reasoner", &driver); + let content = match msg.content { + Some(OaiMessageContent::Text(t)) => t, + _ => panic!("expected text content"), + }; + assert_eq!(content, "answer", "visible content must not include "); + assert_eq!( + msg.reasoning_content.as_deref(), + Some("internal chain-of-thought"), + "reasoning_content field must carry the reasoning text" + ); + } + + /// Without thinking blocks, the outbound message should be a plain + /// assistant message — preserve the legacy shape. + #[test] + fn test_assemble_assistant_no_thinking_is_plain() { + let driver = + OpenAIDriver::new("test".to_string(), "https://api.openai.com/v1".to_string()); + let blocks = vec![ContentBlock::Text { + text: "Hi.".to_string(), + provider_metadata: None, + }]; + let msg = assemble_assistant_message(&blocks, "gpt-4o", &driver); + match msg.content { + Some(OaiMessageContent::Text(t)) => assert_eq!(t, "Hi."), + _ => panic!("expected text content"), + } + assert!(msg.reasoning_content.is_none()); + } + + /// Issue #1098 round-trip: parse a wire response with `reasoning_content`, + /// then feed the parsed assistant turn back through the outbound path + /// and confirm the reasoning is replayed. + #[test] + fn test_reasoning_content_full_round_trip() { + // Step 1: parse server response shape. + let json = serde_json::json!({ + "content": "Final answer.", + "reasoning_content": "I considered options A, B, and C…", + "tool_calls": null + }); + let server_msg: OaiResponseMessage = serde_json::from_value(json).unwrap(); + assert_eq!(server_msg.content.as_deref(), Some("Final answer.")); + assert_eq!( + server_msg.reasoning_content.as_deref(), + Some("I considered options A, B, and C…") + ); + + // Step 2: simulate the driver building blocks (mirrors the live + // path in `complete()`). + let mut content = Vec::new(); + if let Some(ref reasoning) = server_msg.reasoning_content { + content.push(ContentBlock::Thinking { + thinking: reasoning.clone(), + signature: None, + provider_metadata: Some(serde_json::json!({"format": "reasoning_content"})), + }); + } + content.push(ContentBlock::Text { + text: server_msg.content.unwrap(), + provider_metadata: None, + }); + + // Step 3: replay through the outbound path. + let driver = OpenAIDriver::new( + "test".to_string(), + "https://api.deepseek.com/v1".to_string(), + ); + let outbound = assemble_assistant_message(&content, "deepseek-reasoner", &driver); + // The reasoning_content field must round-trip verbatim. + assert_eq!( + outbound.reasoning_content.as_deref(), + Some("I considered options A, B, and C…"), + "issue #1098 regression: reasoning was stripped on resubmission" + ); + } } diff --git a/crates/openfang-types/src/message.rs b/crates/openfang-types/src/message.rs index 00c8312a..67955a44 100644 --- a/crates/openfang-types/src/message.rs +++ b/crates/openfang-types/src/message.rs @@ -85,10 +85,26 @@ pub enum ContentBlock { is_error: bool, }, /// Extended thinking content block (model's reasoning trace). + /// + /// Preserved across turns so reasoning models retain state. Anthropic's + /// extended thinking requires the `signature` to be echoed on resubmission; + /// other providers (Gemini thought signatures, DeepSeek/Qwen + /// `reasoning_content`, MiniMax inline ``) round-trip via + /// `provider_metadata` or by inlining into the assistant message body. #[serde(rename = "thinking")] Thinking { /// The thinking/reasoning text. thinking: String, + /// Provider-issued signature required to resubmit thinking blocks + /// (Anthropic extended thinking). `None` for providers that don't + /// emit a signature. + #[serde(default, skip_serializing_if = "Option::is_none")] + signature: Option, + /// Provider-specific metadata (e.g. `{"format": "reasoning_content"}` + /// or `{"format": "inline_think"}` so the outbound driver knows how + /// the upstream model originally delivered the reasoning). + #[serde(default, skip_serializing_if = "Option::is_none")] + provider_metadata: Option, }, /// Catch-all for unrecognized content block types (forward compatibility). #[serde(other)] @@ -141,7 +157,7 @@ impl MessageContent { .map(|b| match b { ContentBlock::Text { text, .. } => text.len(), ContentBlock::ToolResult { content, .. } => content.len(), - ContentBlock::Thinking { thinking } => thinking.len(), + ContentBlock::Thinking { thinking, .. } => thinking.len(), ContentBlock::ToolUse { name, input, .. } => { name.len() + input.to_string().len() } @@ -199,6 +215,17 @@ impl Message { content: MessageContent::Text(content.into()), } } + + /// Create an assistant message with structured content blocks. + /// + /// Used to preserve `Thinking` blocks (with signatures and reasoning text) + /// across persistence so reasoning models retain state between turns. + pub fn assistant_with_blocks(blocks: Vec) -> Self { + Self { + role: Role::Assistant, + content: MessageContent::Blocks(blocks), + } + } } /// Why the LLM stopped generating. @@ -311,6 +338,119 @@ mod tests { assert!(matches!(block, ContentBlock::Unknown)); } + #[test] + fn test_thinking_block_roundtrip_preserves_signature() { + // Anthropic extended thinking — the signature MUST round-trip through + // serde so it can be echoed on the next request. + let block = ContentBlock::Thinking { + thinking: "Let me reason about this carefully...".to_string(), + signature: Some("sig_abc123_anthropic_extended_thinking".to_string()), + provider_metadata: None, + }; + let json = serde_json::to_value(&block).unwrap(); + assert_eq!(json["type"], "thinking"); + assert_eq!(json["signature"], "sig_abc123_anthropic_extended_thinking"); + + // Round-trip through serialize → deserialize (e.g. SQLite session blob) + let serialized = serde_json::to_string(&block).unwrap(); + let restored: ContentBlock = serde_json::from_str(&serialized).unwrap(); + match restored { + ContentBlock::Thinking { + thinking, signature, .. + } => { + assert_eq!(thinking, "Let me reason about this carefully..."); + assert_eq!( + signature.as_deref(), + Some("sig_abc123_anthropic_extended_thinking") + ); + } + _ => panic!("expected Thinking block"), + } + } + + #[test] + fn test_thinking_block_roundtrip_with_provider_metadata() { + // OpenAI-compat models (DeepSeek-R1, Qwen3, MiniMax) — record the + // wire format so the outbound driver knows whether to re-emit as + // `reasoning_content` or inline `` tags. + let block = ContentBlock::Thinking { + thinking: "step-by-step analysis".to_string(), + signature: None, + provider_metadata: Some(serde_json::json!({"format": "inline_think"})), + }; + let serialized = serde_json::to_string(&block).unwrap(); + let restored: ContentBlock = serde_json::from_str(&serialized).unwrap(); + match restored { + ContentBlock::Thinking { + thinking, + signature, + provider_metadata, + } => { + assert_eq!(thinking, "step-by-step analysis"); + assert!(signature.is_none()); + let meta = provider_metadata.expect("provider_metadata preserved"); + assert_eq!(meta["format"], "inline_think"); + } + _ => panic!("expected Thinking block"), + } + } + + #[test] + fn test_thinking_block_legacy_deser() { + // Existing sessions on disk only have `{"type": "thinking", "thinking": "..."}`. + // The new fields must be optional so old payloads still load. + let json = serde_json::json!({"type": "thinking", "thinking": "old session reasoning"}); + let block: ContentBlock = serde_json::from_value(json).unwrap(); + match block { + ContentBlock::Thinking { + thinking, + signature, + provider_metadata, + } => { + assert_eq!(thinking, "old session reasoning"); + assert!(signature.is_none()); + assert!(provider_metadata.is_none()); + } + _ => panic!("expected Thinking block"), + } + } + + #[test] + fn test_assistant_with_blocks_preserves_thinking() { + // A complete round-trip: build an assistant turn that mixes Thinking + + // Text (the shape we'll store after fix) and confirm the Thinking + // block survives serialization (msgpack is what session.rs uses; JSON + // exercises the same serde path). + let msg = Message::assistant_with_blocks(vec![ + ContentBlock::Thinking { + thinking: "Internal reasoning".to_string(), + signature: Some("sig_xyz".to_string()), + provider_metadata: None, + }, + ContentBlock::Text { + text: "Hello!".to_string(), + provider_metadata: None, + }, + ]); + let bytes = rmp_serde::to_vec_named(&msg).expect("msgpack encode"); + let restored: Message = rmp_serde::from_slice(&bytes).expect("msgpack decode"); + match restored.content { + MessageContent::Blocks(blocks) => { + assert_eq!(blocks.len(), 2); + match &blocks[0] { + ContentBlock::Thinking { + thinking, signature, .. + } => { + assert_eq!(thinking, "Internal reasoning"); + assert_eq!(signature.as_deref(), Some("sig_xyz")); + } + _ => panic!("expected Thinking first"), + } + } + _ => panic!("expected Blocks content"), + } + } + #[test] fn test_user_with_blocks() { let blocks = vec![