admin/openfang
1
0
Fork 0
mirror of https://github.com/RightNow-AI/openfang.git synced 2026-07-29 22:22:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

think persist

Browse Source
This commit is contained in:
jaberjaber23
2026-05-01 13:01:55 +03:00
parent 8642c4d442
commit 2dedab2a8b
6 changed files with 969 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/openfang-memory/src/session.rs
+1 -1
View File
@@ -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)
crates/openfang-runtime/src/agent_loop.rs
+169 -2
View File
@@ -109,6 +109,70 @@ fn append_tool_error_guidance(tool_result_blocks: &mut Vec<ContentBlock>) {
}
}
/// 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 `<think>` 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<ContentBlock> = 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 `<think>`), 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);
crates/openfang-runtime/src/drivers/anthropic.rs
+259 -12
View File
@@ -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<String>,
},
}
#[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"),
}
}
}
crates/openfang-runtime/src/drivers/gemini.rs
+105 -14
View File
@@ -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<CompletionResponse, LlmError
input: function_call.args,
});
}
GeminiPart::Thought { text, .. } => {
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 <think> 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<String> = 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<String> = None;
// Track function calls: (name, args_json, thought_signature)
let mut fn_calls: Vec<(String, serde_json::Value, Option<String>)> = Vec::new();
let mut finish_reason: Option<String> = 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]),
crates/openfang-runtime/src/drivers/openai.rs
+294 -108
View File
@@ -289,6 +289,135 @@ fn strip_trailing_empty_assistant(messages: &mut Vec<OaiMessage>) {
}
}
/// 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, `<think>...</think>` 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 `<think>...</think>` 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<String> = Vec::new();
let mut tool_calls: Vec<OaiToolCall> = Vec::new();
let mut reasoning_field: Option<String> = None;
let mut inline_think: Option<String> = 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 `<think>` in
// historical assistant messages. Concatenate
// multiple thinking blocks if present.
let entry = format!("<think>{thinking}</think>");
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!("<think>{thinking}</think>");
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<CompletionResponse, LlmError> {
@@ -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 `<think>` 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 `<think>`
// 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 `<think>...</think>`.
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 `<think>` 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, "<think>step-by-step reasoning</think>Hello, user.",
"inline_think must be re-emitted as <think> 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 <think>");
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"
);
}
}
crates/openfang-types/src/message.rs
+141 -1
View File
@@ -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 `<think>`) 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<String>,
/// 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<serde_json::Value>,
},
/// 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<ContentBlock>) -> 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 `<think>` 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![