fix(cerebras): preserve thinking inline for models that reject reasoning_content (#10774)

This commit is contained in:
Aaron Alaniz
2026-07-30 16:39:44 -05:00
committed by GitHub
parent ee61c7c499
commit ee5e5f17b2
11 changed files with 410 additions and 25 deletions
+61
View File
@@ -2,6 +2,8 @@ use async_trait::async_trait;
use futures::Stream;
use rmcp::model::Tool;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::pin::Pin;
use crate::{
@@ -216,6 +218,21 @@ impl ConfigKey {
}
}
/// How a model's thinking is replayed back to the provider on subsequent turns.
///
/// Cerebras rejects requests that replay `messages[].reasoning_content`, so such models
/// declare an inline `content` form instead.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum ThinkingPreservationFormat {
/// Prepend the thinking to the message content as plain text.
ContentPrepend,
/// Prepend the thinking to the message content wrapped in `<think>` tags.
ContentXml,
/// Replay in the separate `reasoning_content` field, the OpenAI-compatible default.
ReasoningContent,
}
/// Information about a model's capabilities
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ModelInfo {
@@ -237,6 +254,11 @@ pub struct ModelInfo {
/// Whether this model supports reasoning/thinking controls
#[serde(default)]
pub reasoning: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub thinking_preservation_format: Option<ThinkingPreservationFormat>,
/// Static params merged into the request body for this model.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_params: Option<HashMap<String, Value>>,
}
impl ModelInfo {
@@ -251,6 +273,8 @@ impl ModelInfo {
currency: None,
supports_cache_control: None,
reasoning: false,
thinking_preservation_format: None,
request_params: None,
}
}
@@ -270,6 +294,8 @@ impl ModelInfo {
currency: Some("$".to_string()),
supports_cache_control: None,
reasoning: false,
thinking_preservation_format: None,
request_params: None,
}
}
}
@@ -315,6 +341,8 @@ pub fn model_info_for_provider_model(provider_name: &str, model_name: &str) -> M
currency: None,
supports_cache_control: None,
reasoning,
thinking_preservation_format: None,
request_params: None,
}
}
@@ -775,6 +803,8 @@ mod tests {
currency: None,
supports_cache_control: None,
reasoning: false,
thinking_preservation_format: None,
request_params: None,
};
assert_eq!(info.context_limit, 1000);
@@ -788,6 +818,8 @@ mod tests {
currency: None,
supports_cache_control: None,
reasoning: false,
thinking_preservation_format: None,
request_params: None,
};
assert_eq!(info, info2);
@@ -801,10 +833,39 @@ mod tests {
currency: None,
supports_cache_control: None,
reasoning: false,
thinking_preservation_format: None,
request_params: None,
};
assert_ne!(info, info3);
}
#[test]
fn test_model_info_deserializes_thinking_preservation_and_request_params() {
let info: ModelInfo = serde_json::from_str(
r#"{
"name": "zai-glm-4.7",
"context_limit": 131072,
"thinking_preservation_format": "content_xml",
"request_params": {"reasoning_format": "parsed"}
}"#,
)
.unwrap();
assert_eq!(
info.thinking_preservation_format,
Some(ThinkingPreservationFormat::ContentXml)
);
assert_eq!(
info.request_params.unwrap().get("reasoning_format"),
Some(&serde_json::json!("parsed"))
);
let bare: ModelInfo =
serde_json::from_str(r#"{"name": "gpt-4o", "context_limit": 128000}"#).unwrap();
assert_eq!(bare.thinking_preservation_format, None);
assert_eq!(bare.request_params, None);
}
#[test]
fn test_model_info_with_cost() {
let info = ModelInfo::with_cost("gpt-4o", 128000, 0.0000025, 0.00001);
@@ -1,3 +1,4 @@
use crate::base::ThinkingPreservationFormat;
use crate::conversation::message::{Message, MessageContentBlock, ProviderMetadata};
use crate::conversation::token_usage::{CostSource, ProviderUsage, Usage};
use crate::errors::ProviderError;
@@ -48,13 +49,14 @@ fn describe_json_value(value: &Value) -> &'static str {
}
}
fn is_reserved_request_param_key(key: &str) -> bool {
pub fn is_reserved_request_param_key(key: &str) -> bool {
matches!(key, "messages" | "model" | "stream" | "stream_options")
}
#[derive(Debug, Clone, Copy, Default)]
pub struct OpenAiFormatOptions {
pub preserve_thinking_context: bool,
pub thinking_preservation_format: Option<ThinkingPreservationFormat>,
}
fn merge_reasoning_text(prefix: &str, suffix: &str) -> String {
@@ -187,6 +189,7 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<
image_format,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
)
}
@@ -500,9 +503,43 @@ pub fn format_messages_with_options(
}
merge_split_tool_call_messages(&mut messages_spec);
if let Some(format) = options.thinking_preservation_format {
inline_reasoning_content(&mut messages_spec, format);
}
messages_spec
}
/// Rewrites `reasoning_content` into the message `content` for models that reject a
/// separate reasoning field on replay.
///
/// Must run after `merge_split_tool_call_messages`, which relies on `reasoning_content`
/// to identify messages split from the same assistant turn.
fn inline_reasoning_content(messages: &mut [Value], format: ThinkingPreservationFormat) {
let wrap: fn(&str) -> String = match format {
ThinkingPreservationFormat::ReasoningContent => return,
ThinkingPreservationFormat::ContentPrepend => |text| format!("{text}\n\n"),
ThinkingPreservationFormat::ContentXml => |text| format!("<think>\n{text}\n</think>\n\n"),
};
for message in messages {
let Some(object) = message.as_object_mut() else {
continue;
};
let Some(Value::String(reasoning)) = object.remove("reasoning_content") else {
continue;
};
let prefix = wrap(&reasoning);
match object.entry("content").or_insert(Value::Null) {
Value::String(content) => content.insert_str(0, &prefix),
Value::Array(blocks) => blocks.insert(0, json!({"type": "text", "text": prefix})),
content => *content = json!(prefix.trim_end()),
}
}
}
/// The agent splits a single assistant response with N tool_calls into N
/// interleaved `asst(TC)/tool` pairs, cloning `reasoning_content` onto each.
/// This function merges them back into one assistant message with all tool_calls,
@@ -1389,6 +1426,7 @@ pub fn create_request(
for_streaming,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
)
}
@@ -3331,6 +3369,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
@@ -3389,6 +3428,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: false,
..Default::default()
},
);
@@ -3441,6 +3481,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
@@ -3471,6 +3512,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
@@ -3499,6 +3541,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
@@ -3523,6 +3566,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
@@ -3559,6 +3603,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
@@ -3617,6 +3662,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
@@ -3664,6 +3710,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
@@ -3934,6 +3981,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
assert_eq!(spec.len(), 1);
@@ -3968,6 +4016,7 @@ data: [DONE]"#;
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
assert_eq!(spec.len(), 1);
@@ -4419,4 +4468,111 @@ data: [DONE]"#;
}
}
}
fn format_with_preservation(
messages: &[Message],
format: ThinkingPreservationFormat,
) -> Vec<Value> {
format_messages_with_options(
messages,
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
thinking_preservation_format: Some(format),
},
)
}
#[test]
fn test_thinking_preservation_content_prepend() {
let message = Message::assistant()
.with_thinking("Thinking process", "")
.with_text("Hello");
let spec = format_with_preservation(
std::slice::from_ref(&message),
ThinkingPreservationFormat::ContentPrepend,
);
assert_eq!(spec.len(), 1);
assert_eq!(spec[0]["content"], json!("Thinking process\n\nHello"));
assert!(spec[0].get("reasoning_content").is_none());
}
#[test]
fn test_thinking_preservation_content_xml() {
let message = Message::assistant()
.with_thinking("Thinking process", "")
.with_text("Hello");
let spec = format_with_preservation(
std::slice::from_ref(&message),
ThinkingPreservationFormat::ContentXml,
);
assert_eq!(spec.len(), 1);
assert_eq!(
spec[0]["content"],
json!("<think>\nThinking process\n</think>\n\nHello")
);
assert!(spec[0].get("reasoning_content").is_none());
}
#[test]
fn test_thinking_preservation_reasoning_content_is_unchanged() {
let message = Message::assistant()
.with_thinking("Thinking process", "")
.with_text("Hello");
let spec = format_with_preservation(
std::slice::from_ref(&message),
ThinkingPreservationFormat::ReasoningContent,
);
assert_eq!(spec.len(), 1);
assert_eq!(spec[0]["content"], json!("Hello"));
assert_eq!(spec[0]["reasoning_content"], json!("Thinking process"));
}
#[test]
fn test_thinking_preservation_runs_after_split_tool_call_merge() {
// Split tool-call messages are reunited by matching reasoning_content, so
// inlining must happen afterwards or the merge silently stops working.
let messages = vec![
Message::assistant().with_thinking("reasoning", ""),
Message::assistant()
.with_thinking("reasoning", "")
.with_tool_request(
"tool1",
Ok(CallToolRequestParams::new("tool_a").with_arguments(object!({}))),
),
Message::user().with_tool_response(
"tool1",
Ok(rmcp::model::CallToolResult::success(vec![
ContentBlock::text("result1"),
])),
),
Message::assistant()
.with_thinking("reasoning", "")
.with_tool_request(
"tool2",
Ok(CallToolRequestParams::new("tool_b").with_arguments(object!({}))),
),
];
let spec = format_with_preservation(&messages, ThinkingPreservationFormat::ContentXml);
let assistant: Vec<_> = spec
.iter()
.filter(|m| m.get("role") == Some(&json!("assistant")))
.collect();
assert_eq!(assistant.len(), 1);
assert_eq!(assistant[0]["tool_calls"].as_array().unwrap().len(), 2);
assert!(assistant[0].get("reasoning_content").is_none());
assert_eq!(
assistant[0]["content"],
json!("<think>\nreasoning\n</think>")
);
}
}
@@ -363,6 +363,8 @@ fn model_info_for_deployment(deployment_name: &str, model_name: &str) -> ModelIn
reasoning: canonical
.and_then(|model| model.reasoning)
.unwrap_or_else(|| ModelConfig::new(model_name).is_reasoning_model()),
thinking_preservation_format: None,
request_params: None,
}
}
+2
View File
@@ -469,6 +469,8 @@ impl DatabricksProvider {
currency: None,
supports_cache_control: None,
reasoning,
thinking_preservation_format: None,
request_params: None,
}
}
@@ -8,11 +8,27 @@
"models": [
{
"name": "gpt-oss-120b",
"context_limit": 131072
"context_limit": 131072,
"thinking_preservation_format": "content_prepend",
"request_params": {
"reasoning_format": "parsed"
}
},
{
"name": "zai-glm-4.7",
"context_limit": 131072
"context_limit": 131072,
"thinking_preservation_format": "content_xml",
"request_params": {
"reasoning_format": "parsed"
}
},
{
"name": "gemma-4-31b",
"context_limit": 131072,
"thinking_preservation_format": "content_prepend",
"request_params": {
"reasoning_format": "parsed"
}
}
],
"supports_streaming": true
+132 -15
View File
@@ -8,7 +8,8 @@ use crate::declarative::{DeclarativeProviderConfig, KeyResolver};
use crate::errors::ProviderError;
use crate::formats::openai::is_openai_responses_model;
use crate::formats::openai::{
create_request_with_options, get_cost, get_usage, response_to_message, OpenAiFormatOptions,
create_request_with_options, get_cost, get_usage, is_reserved_request_param_key,
response_to_message, OpenAiFormatOptions,
};
use crate::formats::openai_responses::{
create_responses_request_for_model, get_responses_usage, responses_api_to_message,
@@ -142,7 +143,7 @@ pub struct OpenAiProvider {
custom_headers: Option<HashMap<String, String>>,
supports_streaming: bool,
name: String,
custom_models: Option<Vec<String>>,
custom_models: Option<Vec<ModelInfo>>,
dynamic_models: Option<bool>,
skip_canonical_filtering: bool,
preserve_thinking_context: bool,
@@ -163,7 +164,7 @@ pub struct OpenAiProviderBuilder {
custom_headers: Option<HashMap<String, String>>,
supports_streaming: bool,
name: String,
custom_models: Option<Vec<String>>,
custom_models: Option<Vec<ModelInfo>>,
dynamic_models: Option<bool>,
skip_canonical_filtering: bool,
preserve_thinking_context: bool,
@@ -234,7 +235,7 @@ impl OpenAiProviderBuilder {
self
}
pub fn custom_models(mut self, custom_models: Option<Vec<String>>) -> Self {
pub fn custom_models(mut self, custom_models: Option<Vec<ModelInfo>>) -> Self {
self.custom_models = custom_models;
self
}
@@ -447,6 +448,13 @@ impl OpenAiProvider {
}
}
fn declared_model(&self, model_name: &str) -> Option<&ModelInfo> {
self.custom_models
.as_ref()?
.iter()
.find(|m| m.name == model_name)
}
fn sanitize_request_for_compat(
&self,
mut payload: serde_json::Value,
@@ -709,8 +717,9 @@ impl Provider for OpenAiProvider {
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
if let Some(custom_models) = &self.custom_models {
let names: Vec<String> = custom_models.iter().map(|m| m.name.clone()).collect();
if self.dynamic_models == Some(false) {
return Ok(custom_models.clone());
return Ok(names);
}
match self.fetch_models_from_api().await {
Ok(models) => return Ok(models),
@@ -720,7 +729,7 @@ impl Provider for OpenAiProvider {
self.name,
e
);
return Ok(custom_models.clone());
return Ok(names);
}
Err(e) => return Err(e),
}
@@ -749,7 +758,11 @@ impl Provider for OpenAiProvider {
)
.await
} else {
let payload = create_request_with_options(
let declared_model = self.declared_model(&model_config.model_name);
let thinking_preservation_format =
declared_model.and_then(|m| m.thinking_preservation_format);
let mut payload = create_request_with_options(
model_config,
system,
messages,
@@ -757,9 +770,16 @@ impl Provider for OpenAiProvider {
&ImageFormat::OpenAi,
self.supports_streaming,
OpenAiFormatOptions {
preserve_thinking_context: self.preserve_thinking_context,
preserve_thinking_context: self.preserve_thinking_context
|| thinking_preservation_format.is_some(),
thinking_preservation_format,
},
)?;
if let Some(params) = declared_model.and_then(|m| m.request_params.as_ref()) {
apply_declared_request_params(&mut payload, params);
}
let payload = self.sanitize_request_for_compat(payload, model_config);
let mut log = start_log(model_config, &payload)?;
@@ -807,19 +827,31 @@ impl Provider for OpenAiProvider {
}
}
/// Merges a model's declared `request_params` into an already-built payload.
///
/// Reserved keys are skipped so a declaration cannot clobber the streaming setup.
fn apply_declared_request_params(
payload: &mut serde_json::Value,
params: &HashMap<String, serde_json::Value>,
) {
let Some(object) = payload.as_object_mut() else {
return;
};
for (key, value) in params {
if !is_reserved_request_param_key(key) {
object.insert(key.clone(), value.clone());
}
}
}
pub fn from_declarative_config(
config: DeclarativeProviderConfig,
tls_config: Option<TlsConfig>,
key_resolver: impl KeyResolver,
) -> Result<OpenAiProviderBuilder> {
let custom_models = if !config.models.is_empty() {
Some(
config
.models
.iter()
.map(|m| m.name.clone())
.collect::<Vec<String>>(),
)
Some(config.models.clone())
} else {
None
};
@@ -1395,4 +1427,89 @@ mod tests {
let r = derive_base_path("/api/voice");
assert_eq!(r, "api/voice/v1/chat/completions");
}
use crate::base::ThinkingPreservationFormat;
fn cerebras_config() -> DeclarativeProviderConfig {
crate::declarative::fixed_provider_configs()
.expect("bundled providers should load")
.into_iter()
.find(|config| config.name == "cerebras")
.expect("cerebras should be bundled")
}
fn cerebras_provider() -> OpenAiProvider {
struct StaticKeyResolver;
impl KeyResolver for StaticKeyResolver {
type Error = std::convert::Infallible;
fn resolve_key(&self, _key: &str) -> std::result::Result<String, Self::Error> {
Ok("test-key".to_string())
}
}
from_declarative_config(cerebras_config(), None, StaticKeyResolver)
.expect("cerebras config should build a provider")
.build()
}
#[test]
fn cerebras_models_declare_thinking_preservation_and_reasoning_format() {
let provider = cerebras_provider();
for (model, expected_format) in [
("gpt-oss-120b", ThinkingPreservationFormat::ContentPrepend),
("zai-glm-4.7", ThinkingPreservationFormat::ContentXml),
("gemma-4-31b", ThinkingPreservationFormat::ContentPrepend),
] {
let declared = provider
.declared_model(model)
.unwrap_or_else(|| panic!("{model} should be declared"));
assert_eq!(declared.thinking_preservation_format, Some(expected_format));
let reasoning_format = declared
.request_params
.as_ref()
.and_then(|params| params.get("reasoning_format"));
assert_eq!(
reasoning_format,
Some(&json!("parsed")),
"{model} must request parsed reasoning"
);
}
assert!(provider.declared_model("not-a-cerebras-model").is_none());
}
#[test]
fn cerebras_preserves_thinking_by_default() {
assert!(cerebras_config().preserves_thinking);
}
#[test]
fn apply_declared_request_params_skips_reserved_keys() {
let mut payload = json!({
"model": "zai-glm-4.7",
"stream": true,
"stream_options": {"include_usage": true},
"messages": [{"role": "user", "content": "hi"}]
});
let params = HashMap::from([
("reasoning_format".to_string(), json!("parsed")),
("model".to_string(), json!("hijacked")),
("stream".to_string(), json!(false)),
("stream_options".to_string(), json!(null)),
("messages".to_string(), json!([])),
]);
apply_declared_request_params(&mut payload, &params);
assert_eq!(payload["reasoning_format"], json!("parsed"));
assert_eq!(payload["model"], json!("zai-glm-4.7"));
assert_eq!(payload["stream"], json!(true));
assert_eq!(payload["stream_options"], json!({"include_usage": true}));
assert_eq!(payload["messages"].as_array().unwrap().len(), 1);
}
}
@@ -72,6 +72,7 @@ impl OpenAiCompatibleProvider {
for_streaming,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
)
.map_err(|e| ProviderError::RequestFailed(format!("Failed to create request: {}", e)))
@@ -549,6 +549,8 @@ mod tests {
currency: None,
supports_cache_control: None,
reasoning: false,
thinking_preservation_format: None,
request_params: None,
}],
headers: None,
timeout_seconds: None,
@@ -232,14 +232,9 @@ impl ProviderRegistry {
.models
.iter()
.map(|m| ModelInfo {
name: m.name.clone(),
resolved_model: None,
context_limit: m.context_limit,
input_token_cost: m.input_token_cost,
output_token_cost: m.output_token_cost,
currency: m.currency.clone(),
supports_cache_control: Some(m.supports_cache_control.unwrap_or(false)),
reasoning: m.reasoning,
..m.clone()
})
.collect();
+2
View File
@@ -1650,6 +1650,7 @@ mod tests {
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
let has_reasoning_on_tool_call = spec.iter().any(|m| {
@@ -2042,6 +2043,7 @@ mod tests {
&ImageFormat::OpenAi,
OpenAiFormatOptions {
preserve_thinking_context: true,
..Default::default()
},
);
+32 -1
View File
@@ -13,6 +13,7 @@ activities:
- Test extension discovery and management
- Test load tool for knowledge injection and discovery
- Test delegate tool for task delegation (sync and async)
- Test multi-turn thinking preservation for providers that reject replayed reasoning_content
- Test error boundaries including nested delegation prevention
- Generate comprehensive test report
@@ -21,7 +22,7 @@ parameters:
input_type: string
requirement: optional
default: "all"
description: "Which test phases to run: all, basic, extensions, delegation, advanced"
description: "Which test phases to run: all, basic, extensions, delegation, reasoning, advanced"
- key: test_depth
input_type: string
@@ -357,6 +358,36 @@ prompt: |
Log results to: {{ workspace_dir }}/phase3b_vision.md
{% endif %}
{% if test_phases == "all" or "reasoning" in test_phases %}
## 🧠 PHASE 3C: Thinking Preservation Testing
**Prerequisites**: CEREBRAS_API_KEY must be set and the session must be running a Cerebras
model with thinking enabled. Skip this phase and record it as SKIPPED otherwise.
Cerebras rejects multi-turn requests that replay thinking in `messages[].reasoning_content`
with `400 wrong_api_format`. Models declare a `thinking_preservation_format` so their
thinking is replayed inline in `content` instead, and `request_params.reasoning_format`
asks Cerebras to return reasoning in a structured field.
### Multi-Turn Thinking Replay Test
1. Confirm the active model is one of `zai-glm-4.7` (content_xml), `gpt-oss-120b`
(content_prepend), or `gemma-4-31b` (content_prepend).
2. Ask a question that requires reasoning, e.g. "How many times does the letter r appear
in strawberry, raspberry, and blackberry combined? Reason it through."
3. Verify the response includes visible thinking and a final answer.
4. Ask a follow-up in the same session that depends on the first answer, e.g.
"Now subtract the number of r's in blackberry from that total."
5. Verify the second turn succeeds. A `400 wrong_api_format` here is a FAILURE — it means
thinking was replayed as `reasoning_content` instead of inline in `content`.
6. Run at least one more follow-up turn to confirm the session stays healthy as history grows.
### Thinking Format Regression Test
1. Repeat the multi-turn exchange for a `content_xml` model and a `content_prepend` model.
2. Verify neither turn errors and that earlier reasoning is still reflected in later answers.
Log results to: {{ workspace_dir }}/phase3c_reasoning.md
{% endif %}
{% if test_phases == "all" or "advanced" in test_phases %}
## 🔬 PHASE 4: Advanced Testing