fix(providers): Buffer thinking across stream chunks for DeepSeek/Kimi tool calls (#9704)
Signed-off-by: thepetk <thepetk@gmail.com>
This commit is contained in:
Theofanis Petkos
@@ -182,12 +182,37 @@ pub fn format_messages_with_options(
|
||||
) -> Vec<Value> {
|
||||
let mut messages_spec = Vec::new();
|
||||
let mut pending_assistant_reasoning = String::new();
|
||||
// Reasoning to propagate across consecutive tool-call messages in the same turn.
|
||||
// DeepSeek/Kimi require reasoning_content on every assistant tool-call message.
|
||||
let mut tool_call_turn_reasoning = String::new();
|
||||
let mut saw_tool_response = false;
|
||||
|
||||
for message in messages {
|
||||
if options.preserve_thinking_context && message.role != Role::Assistant {
|
||||
pending_assistant_reasoning.clear();
|
||||
}
|
||||
|
||||
if options.preserve_thinking_context && message.role == Role::User {
|
||||
if message
|
||||
.content
|
||||
.iter()
|
||||
.any(|c| matches!(c, MessageContent::ToolResponse(_)))
|
||||
{
|
||||
saw_tool_response = true;
|
||||
} else {
|
||||
tool_call_turn_reasoning.clear();
|
||||
saw_tool_response = false;
|
||||
}
|
||||
}
|
||||
|
||||
// A new assistant message after tool results creates a new turn.
|
||||
// Prevents reasoning from the previous turn leaking into the new one.
|
||||
if options.preserve_thinking_context && message.role == Role::Assistant && saw_tool_response
|
||||
{
|
||||
tool_call_turn_reasoning.clear();
|
||||
saw_tool_response = false;
|
||||
}
|
||||
|
||||
let mut converted = json!({
|
||||
"role": message.role
|
||||
});
|
||||
@@ -408,6 +433,25 @@ pub fn format_messages_with_options(
|
||||
merge_reasoning_text(&pending_assistant_reasoning, &reasoning_text);
|
||||
pending_assistant_reasoning.clear();
|
||||
}
|
||||
|
||||
let has_tool_calls = converted
|
||||
.get("tool_calls")
|
||||
.and_then(|tc| tc.as_array())
|
||||
.is_some_and(|a| !a.is_empty());
|
||||
|
||||
if has_tool_calls {
|
||||
if reasoning_text.is_empty() {
|
||||
reasoning_text = tool_call_turn_reasoning.clone();
|
||||
} else {
|
||||
tool_call_turn_reasoning = reasoning_text.clone();
|
||||
}
|
||||
} else {
|
||||
// Carry reasoning forward even through non-tool assistant messages
|
||||
// (e.g., a visible text chunk that's is sent before a tool-call chunk
|
||||
// in the same streaming turn). An empty reasoning_text is equivalent
|
||||
// to clear.
|
||||
tool_call_turn_reasoning = reasoning_text.clone();
|
||||
}
|
||||
}
|
||||
|
||||
// Include reasoning_content only when non-empty. Kimi rejects empty
|
||||
@@ -3135,6 +3179,180 @@ data: [DONE]"#;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_messages_carries_reasoning_through_text_only_chunks() -> anyhow::Result<()> {
|
||||
// Scenario B from the streaming bug: thinking arrives first, then multiple
|
||||
// text-only assistant messages, then a tool call with thinking re-attached
|
||||
// by agent.rs (via the earlier-chunk lookback).
|
||||
// Text-only messages set tool_call_turn_reasoning="" (line 453 else-branch),
|
||||
// but the TC's own Thinking content must repopulate it.
|
||||
let messages = vec![
|
||||
Message::assistant().with_content(MessageContent::thinking("reason", "")),
|
||||
Message::assistant().with_text("partial answer"),
|
||||
Message::assistant().with_text("more text"),
|
||||
// agent.rs attaches the earlier thinking to the TC message
|
||||
Message::assistant()
|
||||
.with_content(MessageContent::thinking("reason", ""))
|
||||
.with_tool_request(
|
||||
"tool1",
|
||||
Ok(CallToolRequestParams::new("test_tool").with_arguments(object!({}))),
|
||||
),
|
||||
];
|
||||
|
||||
let spec = format_messages_with_options(
|
||||
&messages,
|
||||
&ImageFormat::OpenAi,
|
||||
OpenAiFormatOptions {
|
||||
preserve_thinking_context: true,
|
||||
},
|
||||
);
|
||||
|
||||
let tool_call_msgs: Vec<_> = spec
|
||||
.iter()
|
||||
.filter(|m| {
|
||||
m.get("tool_calls")
|
||||
.and_then(|tc| tc.as_array())
|
||||
.is_some_and(|a| !a.is_empty())
|
||||
})
|
||||
.collect();
|
||||
|
||||
assert_eq!(tool_call_msgs.len(), 1);
|
||||
assert_eq!(
|
||||
tool_call_msgs[0]["reasoning_content"], "reason",
|
||||
"reasoning_content must survive text-only chunks between thinking and tool call"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_messages_carries_reasoning_to_all_split_tool_calls() -> anyhow::Result<()> {
|
||||
// Simulates DeepSeek/Kimi streaming: a thinking-only chunk arrives first,
|
||||
// then the agent splits two tool calls into separate messages, each with
|
||||
// the same reasoning attached (as agent.rs does via response_thinking).
|
||||
// The formatter must keep reasoning_content on both so that
|
||||
// merge_split_tool_call_messages can reunite them into one assistant message.
|
||||
let tool_result1 = Message::user().with_tool_response(
|
||||
"tool1",
|
||||
Ok(rmcp::model::CallToolResult::success(vec![
|
||||
rmcp::model::Content::text("result1"),
|
||||
])),
|
||||
);
|
||||
let messages = vec![
|
||||
// Standalone thinking message (created by agent.rs alongside request_msgs)
|
||||
Message::assistant().with_content(MessageContent::thinking("reasoning", "")),
|
||||
// Each request_msg has thinking explicitly attached (agent.rs behaviour)
|
||||
Message::assistant()
|
||||
.with_content(MessageContent::thinking("reasoning", ""))
|
||||
.with_tool_request(
|
||||
"tool1",
|
||||
Ok(CallToolRequestParams::new("tool_a").with_arguments(object!({}))),
|
||||
),
|
||||
tool_result1,
|
||||
Message::assistant()
|
||||
.with_content(MessageContent::thinking("reasoning", ""))
|
||||
.with_tool_request(
|
||||
"tool2",
|
||||
Ok(CallToolRequestParams::new("tool_b").with_arguments(object!({}))),
|
||||
),
|
||||
];
|
||||
|
||||
let spec = format_messages_with_options(
|
||||
&messages,
|
||||
&ImageFormat::OpenAi,
|
||||
OpenAiFormatOptions {
|
||||
preserve_thinking_context: true,
|
||||
},
|
||||
);
|
||||
|
||||
// After merge: one assistant message with both tool calls
|
||||
let assistant_msgs: Vec<_> = spec
|
||||
.iter()
|
||||
.filter(|m| m.get("role") == Some(&json!("assistant")))
|
||||
.collect();
|
||||
assert_eq!(assistant_msgs.len(), 1);
|
||||
assert_eq!(assistant_msgs[0]["reasoning_content"], "reasoning");
|
||||
let tool_calls = assistant_msgs[0]["tool_calls"].as_array().unwrap();
|
||||
assert_eq!(tool_calls.len(), 2);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sequential_tool_calls_not_merged() -> anyhow::Result<()> {
|
||||
// Verifies that two tool calls from *different* turns are never merged,
|
||||
// even when the second call carries no fresh reasoning (the previous
|
||||
// turn's reasoning must not leak into it).
|
||||
let tool_result1 = Message::user().with_tool_response(
|
||||
"tool1",
|
||||
Ok(rmcp::model::CallToolResult::success(vec![
|
||||
rmcp::model::Content::text("result1"),
|
||||
])),
|
||||
);
|
||||
let messages = vec![
|
||||
// Turn 1: thinking then tool call
|
||||
Message::assistant().with_content(MessageContent::thinking("turn1_reasoning", "")),
|
||||
Message::assistant().with_tool_request(
|
||||
"tool1",
|
||||
Ok(CallToolRequestParams::new("tool_a").with_arguments(object!({}))),
|
||||
),
|
||||
tool_result1,
|
||||
// Turn 2: new tool call, no fresh thinking
|
||||
Message::assistant().with_tool_request(
|
||||
"tool2",
|
||||
Ok(CallToolRequestParams::new("tool_b").with_arguments(object!({}))),
|
||||
),
|
||||
];
|
||||
|
||||
let spec = format_messages_with_options(
|
||||
&messages,
|
||||
&ImageFormat::OpenAi,
|
||||
OpenAiFormatOptions {
|
||||
preserve_thinking_context: true,
|
||||
},
|
||||
);
|
||||
|
||||
let assistant_msgs: Vec<_> = spec
|
||||
.iter()
|
||||
.filter(|m| m.get("role") == Some(&json!("assistant")))
|
||||
.collect();
|
||||
|
||||
// Must remain two separate assistant messages — not merged across turns.
|
||||
assert_eq!(
|
||||
assistant_msgs.len(),
|
||||
2,
|
||||
"sequential tool calls must not be merged"
|
||||
);
|
||||
|
||||
// Turn 1 carries reasoning; turn 2 must not inherit it.
|
||||
assert_eq!(assistant_msgs[0]["reasoning_content"], "turn1_reasoning");
|
||||
assert!(
|
||||
assistant_msgs[1].get("reasoning_content").is_none()
|
||||
|| assistant_msgs[1]["reasoning_content"].is_null(),
|
||||
"turn 2 must not inherit stale reasoning from turn 1"
|
||||
);
|
||||
|
||||
// The tool result must appear between the two assistant messages.
|
||||
let tool_idx = spec
|
||||
.iter()
|
||||
.position(|m| m.get("role") == Some(&json!("tool")))
|
||||
.expect("tool result must be present");
|
||||
let asst1_idx = spec
|
||||
.iter()
|
||||
.position(|m| m.get("role") == Some(&json!("assistant")))
|
||||
.unwrap();
|
||||
let asst2_idx = spec
|
||||
.iter()
|
||||
.rposition(|m| m.get("role") == Some(&json!("assistant")))
|
||||
.unwrap();
|
||||
assert!(
|
||||
asst1_idx < tool_idx && tool_idx < asst2_idx,
|
||||
"tool result must sit between the two assistant messages"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test_case(
|
||||
"data: {\"error\":{\"message\":\"Internal server error\",\"type\":\"server_error\",\"code\":500}}\ndata: [DONE]",
|
||||
"Internal server error";
|
||||
|
||||
@@ -2239,27 +2239,54 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
// Preserve thinking/reasoning content from the original response
|
||||
// Gemini (and other thinking models) require thinking to be echoed back
|
||||
// Kimi/DeepSeek require reasoning_content on assistant tool call messages
|
||||
let thinking_content: Vec<MessageContent> = response.content.iter()
|
||||
.filter(|c| matches!(c, MessageContent::Thinking(_)))
|
||||
// Preserve thinking/reasoning content from the original response.
|
||||
// Gemini (and other thinking models) require thinking to be echoed back.
|
||||
// Kimi/DeepSeek require reasoning_content on assistant tool call messages.
|
||||
let direct_thinking: Vec<MessageContent> = response
|
||||
.content
|
||||
.iter()
|
||||
.filter(|c| {
|
||||
matches!(
|
||||
c,
|
||||
MessageContent::Thinking(_)
|
||||
| MessageContent::RedactedThinking(_)
|
||||
)
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
if !thinking_content.is_empty() {
|
||||
if !direct_thinking.is_empty() {
|
||||
let thinking_msg = Message::new(
|
||||
response.role.clone(),
|
||||
response.created,
|
||||
thinking_content,
|
||||
).with_id(format!("msg_{}", Uuid::new_v4()));
|
||||
direct_thinking.clone(),
|
||||
)
|
||||
.with_id(format!("msg_{}", Uuid::new_v4()));
|
||||
messages_to_add.push(thinking_msg);
|
||||
}
|
||||
|
||||
// Collect reasoning content to attach to tool request messages
|
||||
let reasoning_content: Vec<MessageContent> = response.content.iter()
|
||||
.filter(|c| matches!(c, MessageContent::Thinking(_)))
|
||||
.cloned()
|
||||
.collect();
|
||||
// When thinking arrived in an earlier stream chunk (stored as a
|
||||
// thinking-only message) and this chunk has only tool calls,
|
||||
// reuse that thinking so each split request_msg carries it.
|
||||
let response_thinking = if direct_thinking.is_empty() {
|
||||
messages_to_add
|
||||
.messages()
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|m| {
|
||||
m.role == response.role
|
||||
&& !m.content.is_empty()
|
||||
&& m.content.iter().all(|c| {
|
||||
matches!(
|
||||
c,
|
||||
MessageContent::Thinking(_)
|
||||
| MessageContent::RedactedThinking(_)
|
||||
)
|
||||
})
|
||||
})
|
||||
.map(|m| m.content.clone())
|
||||
.unwrap_or_default()
|
||||
} else {
|
||||
direct_thinking
|
||||
};
|
||||
|
||||
for request in frontend_requests.iter().chain(remaining_requests.iter()) {
|
||||
if let Err(err) = &request.tool_call {
|
||||
@@ -2274,10 +2301,8 @@ impl Agent {
|
||||
let mut request_msg = Message::assistant()
|
||||
.with_id(format!("msg_{}", Uuid::new_v4()));
|
||||
|
||||
// Providers like Kimi require reasoning_content on all assistant
|
||||
// messages with tool_calls when thinking mode is enabled.
|
||||
for rc in &reasoning_content {
|
||||
request_msg = request_msg.with_content(rc.clone());
|
||||
for thinking in &response_thinking {
|
||||
request_msg = request_msg.with_content(thinking.clone());
|
||||
}
|
||||
|
||||
request_msg = request_msg
|
||||
|
||||
@@ -1122,6 +1122,566 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod thinking_preservation_tests {
|
||||
use super::*;
|
||||
use async_trait::async_trait;
|
||||
use goose::agents::{AgentConfig, SessionConfig};
|
||||
use goose::config::permission::PermissionManager;
|
||||
use goose::config::GooseMode;
|
||||
use goose::conversation::message::{Message, MessageContent};
|
||||
use goose::providers::base::{MessageStream, Provider, ProviderDef, ProviderMetadata};
|
||||
use goose::session::session_manager::SessionType;
|
||||
use goose::session::SessionManager;
|
||||
use goose_providers::conversation::token_usage::{ProviderUsage, Usage};
|
||||
use goose_providers::errors::ProviderError;
|
||||
use goose_providers::model::ModelConfig;
|
||||
use rmcp::model::{CallToolRequestParams, Tool};
|
||||
use rmcp::object;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
/// Simulates DeepSeek/Kimi streaming: reasoning_content arrives in an early
|
||||
/// chunk, the tool call arrives in a later chunk with no reasoning_content.
|
||||
struct ThinkingStreamProvider {
|
||||
call_count: AtomicUsize,
|
||||
name: &'static str,
|
||||
}
|
||||
|
||||
impl ThinkingStreamProvider {
|
||||
fn new(name: &'static str) -> Self {
|
||||
Self {
|
||||
call_count: AtomicUsize::new(0),
|
||||
name,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl goose::providers::base::ProviderDescriptor for ThinkingStreamProvider {
|
||||
fn metadata() -> ProviderMetadata {
|
||||
ProviderMetadata {
|
||||
name: "thinking-stream-mock".to_string(),
|
||||
display_name: "Thinking Stream Mock".to_string(),
|
||||
description: "Mock for thinking preservation tests".to_string(),
|
||||
default_model: "mock-model".to_string(),
|
||||
known_models: vec![],
|
||||
model_doc_link: "".to_string(),
|
||||
config_keys: vec![],
|
||||
setup_steps: vec![],
|
||||
model_selection_hint: None,
|
||||
fast_model: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderDef for ThinkingStreamProvider {
|
||||
type Provider = Self;
|
||||
|
||||
fn from_env(
|
||||
_extensions: Vec<goose::config::ExtensionConfig>,
|
||||
_tls_config: Option<goose::providers::api_client::TlsConfig>,
|
||||
) -> futures::future::BoxFuture<'static, anyhow::Result<Self>> {
|
||||
unimplemented!()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for ThinkingStreamProvider {
|
||||
async fn stream(
|
||||
&self,
|
||||
_model_config: &ModelConfig,
|
||||
_session_id: &str,
|
||||
_system_prompt: &str,
|
||||
_messages: &[Message],
|
||||
_tools: &[Tool],
|
||||
) -> Result<MessageStream, ProviderError> {
|
||||
let call = self.call_count.fetch_add(1, Ordering::SeqCst);
|
||||
let usage = ProviderUsage::new(
|
||||
"mock-model".to_string(),
|
||||
Usage::new(Some(10), Some(20), Some(30)),
|
||||
);
|
||||
match call {
|
||||
0 => {
|
||||
// Chunk 1: reasoning_content only (no tool call)
|
||||
let thinking =
|
||||
Message::assistant().with_thinking("I should call test_tool", "sig_0");
|
||||
// Chunk 2: tool call only (no reasoning_content) — the bug scenario
|
||||
let tool_call = CallToolRequestParams::new("test_tool")
|
||||
.with_arguments(object!({"param": "value"}));
|
||||
let tool_msg =
|
||||
Message::assistant().with_tool_request("call_1", Ok(tool_call));
|
||||
let stream = futures::stream::iter(vec![
|
||||
Ok((Some(thinking), None)),
|
||||
Ok((Some(tool_msg), Some(usage))),
|
||||
]);
|
||||
Ok(Box::pin(stream))
|
||||
}
|
||||
_ => {
|
||||
let msg = Message::assistant().with_text("Done.");
|
||||
Ok(Box::pin(futures::stream::once(async move {
|
||||
Ok((Some(msg), Some(usage)))
|
||||
})))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_name(&self) -> &str {
|
||||
self.name
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_and_collect(provider_name: &'static str) -> Result<Vec<Message>> {
|
||||
let temp_dir = tempfile::tempdir()?;
|
||||
let session_manager = Arc::new(SessionManager::new(temp_dir.path().to_path_buf()));
|
||||
let config = AgentConfig::new(
|
||||
session_manager.clone(),
|
||||
PermissionManager::instance(),
|
||||
None,
|
||||
GooseMode::Auto,
|
||||
true,
|
||||
GoosePlatform::GooseCli,
|
||||
);
|
||||
let agent = Agent::with_config(config);
|
||||
let provider = Arc::new(ThinkingStreamProvider::new(provider_name));
|
||||
|
||||
let session = session_manager
|
||||
.create_session(
|
||||
PathBuf::default(),
|
||||
format!("{provider_name}-thinking-test"),
|
||||
SessionType::Hidden,
|
||||
GooseMode::default(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let session_id = session.id.clone();
|
||||
agent
|
||||
.update_provider(provider, ModelConfig::new("mock-model"), &session_id)
|
||||
.await?;
|
||||
|
||||
let session_config = SessionConfig {
|
||||
id: session_id.clone(),
|
||||
schedule_id: None,
|
||||
max_turns: Some(2),
|
||||
retry_config: None,
|
||||
};
|
||||
|
||||
let reply_stream = agent
|
||||
.reply(
|
||||
Message::user().with_text("Use the test tool"),
|
||||
session_config,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
tokio::pin!(reply_stream);
|
||||
|
||||
while let Some(event) = reply_stream.next().await {
|
||||
event?;
|
||||
}
|
||||
|
||||
let reloaded = session_manager.get_session(&session_id, true).await?;
|
||||
Ok(reloaded
|
||||
.conversation
|
||||
.expect("should have conversation")
|
||||
.messages()
|
||||
.to_vec())
|
||||
}
|
||||
|
||||
fn assert_formatter_adds_reasoning_to_tool_calls(messages: &[Message], provider: &str) {
|
||||
use goose_providers::formats::openai::{
|
||||
format_messages_with_options, OpenAiFormatOptions,
|
||||
};
|
||||
use goose_providers::images::ImageFormat;
|
||||
|
||||
assert!(
|
||||
messages.iter().any(|m| m
|
||||
.content
|
||||
.iter()
|
||||
.any(|c| matches!(c, MessageContent::Thinking(_)))),
|
||||
"{provider}: conversation must contain at least one Thinking message"
|
||||
);
|
||||
assert!(
|
||||
messages.iter().any(|m| m
|
||||
.content
|
||||
.iter()
|
||||
.any(|c| matches!(c, MessageContent::ToolRequest(_)))),
|
||||
"{provider}: conversation must contain at least one tool-call message"
|
||||
);
|
||||
|
||||
let spec = format_messages_with_options(
|
||||
messages,
|
||||
&ImageFormat::OpenAi,
|
||||
OpenAiFormatOptions {
|
||||
preserve_thinking_context: true,
|
||||
},
|
||||
);
|
||||
let has_reasoning_on_tool_call = spec.iter().any(|m| {
|
||||
m.get("tool_calls")
|
||||
.and_then(|tc| tc.as_array())
|
||||
.is_some_and(|a| !a.is_empty())
|
||||
&& m.get("reasoning_content").is_some()
|
||||
});
|
||||
assert!(
|
||||
has_reasoning_on_tool_call,
|
||||
"{provider}: formatter must produce reasoning_content on assistant tool-call \
|
||||
messages — {provider} returns HTTP 400 when it is absent on the next turn"
|
||||
);
|
||||
}
|
||||
|
||||
/// DeepSeek streams reasoning_content before the tool-call chunk. The formatter
|
||||
/// must attach it to the tool-call message so the next turn is accepted.
|
||||
#[tokio::test]
|
||||
async fn test_deepseek_thinking_preserved_in_tool_call_message() -> Result<()> {
|
||||
let messages = run_and_collect("deepseek-mock").await?;
|
||||
assert_formatter_adds_reasoning_to_tool_calls(&messages, "DeepSeek");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Kimi has the same streaming behaviour as DeepSeek.
|
||||
#[tokio::test]
|
||||
async fn test_kimi_thinking_preserved_in_tool_call_message() -> Result<()> {
|
||||
let messages = run_and_collect("kimi-mock").await?;
|
||||
assert_formatter_adds_reasoning_to_tool_calls(&messages, "Kimi");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Simulates a provider that emits reasoning and a tool call in the same
|
||||
/// streamed message (no prior thinking-only chunk).
|
||||
struct CombinedThinkingToolProvider {
|
||||
call_count: AtomicUsize,
|
||||
}
|
||||
|
||||
impl CombinedThinkingToolProvider {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
call_count: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl goose::providers::base::ProviderDescriptor for CombinedThinkingToolProvider {
|
||||
fn metadata() -> ProviderMetadata {
|
||||
ProviderMetadata {
|
||||
name: "combined-thinking-tool-mock".to_string(),
|
||||
display_name: "Combined Thinking+Tool Mock".to_string(),
|
||||
description: "Mock for combined thinking+tool call in one chunk".to_string(),
|
||||
default_model: "mock-model".to_string(),
|
||||
known_models: vec![],
|
||||
model_doc_link: "".to_string(),
|
||||
config_keys: vec![],
|
||||
setup_steps: vec![],
|
||||
model_selection_hint: None,
|
||||
fast_model: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderDef for CombinedThinkingToolProvider {
|
||||
type Provider = Self;
|
||||
|
||||
fn from_env(
|
||||
_extensions: Vec<goose::config::ExtensionConfig>,
|
||||
_tls_config: Option<goose::providers::api_client::TlsConfig>,
|
||||
) -> futures::future::BoxFuture<'static, anyhow::Result<Self>> {
|
||||
unimplemented!()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for CombinedThinkingToolProvider {
|
||||
async fn stream(
|
||||
&self,
|
||||
_model_config: &ModelConfig,
|
||||
_session_id: &str,
|
||||
_system_prompt: &str,
|
||||
_messages: &[Message],
|
||||
_tools: &[Tool],
|
||||
) -> Result<MessageStream, ProviderError> {
|
||||
let call = self.call_count.fetch_add(1, Ordering::SeqCst);
|
||||
let usage = ProviderUsage::new(
|
||||
"mock-model".to_string(),
|
||||
Usage::new(Some(10), Some(20), Some(30)),
|
||||
);
|
||||
match call {
|
||||
0 => {
|
||||
// Single chunk: reasoning_content AND tool call together
|
||||
let tool_call = CallToolRequestParams::new("test_tool")
|
||||
.with_arguments(object!({"param": "value"}));
|
||||
let combined = Message::assistant()
|
||||
.with_thinking("I should call test_tool", "sig_0")
|
||||
.with_tool_request("call_1", Ok(tool_call));
|
||||
Ok(Box::pin(futures::stream::once(async move {
|
||||
Ok((Some(combined), Some(usage)))
|
||||
})))
|
||||
}
|
||||
_ => {
|
||||
let msg = Message::assistant().with_text("Done.");
|
||||
Ok(Box::pin(futures::stream::once(async move {
|
||||
Ok((Some(msg), Some(usage)))
|
||||
})))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_name(&self) -> &str {
|
||||
"combined-thinking-tool-mock"
|
||||
}
|
||||
}
|
||||
|
||||
/// When reasoning arrives in the same chunk as the tool call (no prior
|
||||
/// thinking-only message), the agent must attach it to the persisted
|
||||
/// request_msg so the formatter can emit reasoning_content on the next turn.
|
||||
#[tokio::test]
|
||||
async fn test_reasoning_preserved_when_combined_with_tool_call() -> Result<()> {
|
||||
let temp_dir = tempfile::tempdir()?;
|
||||
let session_manager = Arc::new(SessionManager::new(temp_dir.path().to_path_buf()));
|
||||
let config = AgentConfig::new(
|
||||
session_manager.clone(),
|
||||
PermissionManager::instance(),
|
||||
None,
|
||||
GooseMode::Auto,
|
||||
true,
|
||||
GoosePlatform::GooseCli,
|
||||
);
|
||||
let agent = Agent::with_config(config);
|
||||
let provider = Arc::new(CombinedThinkingToolProvider::new());
|
||||
|
||||
let session = session_manager
|
||||
.create_session(
|
||||
PathBuf::default(),
|
||||
"combined-thinking-tool-test".to_string(),
|
||||
SessionType::Hidden,
|
||||
GooseMode::default(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let session_id = session.id.clone();
|
||||
agent
|
||||
.update_provider(provider, ModelConfig::new("mock-model"), &session_id)
|
||||
.await?;
|
||||
|
||||
let session_config = SessionConfig {
|
||||
id: session_id.clone(),
|
||||
schedule_id: None,
|
||||
max_turns: Some(2),
|
||||
retry_config: None,
|
||||
};
|
||||
|
||||
let reply_stream = agent
|
||||
.reply(
|
||||
Message::user().with_text("Use the test tool"),
|
||||
session_config,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
tokio::pin!(reply_stream);
|
||||
while let Some(event) = reply_stream.next().await {
|
||||
match event {
|
||||
Ok(_) => {}
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
let reloaded = session_manager.get_session(&session_id, true).await?;
|
||||
let messages = reloaded
|
||||
.conversation
|
||||
.expect("should have conversation")
|
||||
.messages()
|
||||
.to_vec();
|
||||
|
||||
assert_formatter_adds_reasoning_to_tool_calls(&messages, "combined-thinking-tool");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Simulates the DeepSeek/Kimi multi-tool-call case: thinking arrives as a
|
||||
/// separate stream chunk, then both tool calls arrive together in a second
|
||||
/// chunk with no thinking. Before the fix, the second tool-call message
|
||||
/// (asst(TC2)) received no reasoning_content because lines 210-213 in
|
||||
/// format_messages_with_options cleared tool_call_turn_reasoning after the
|
||||
/// first tool result.
|
||||
struct MultiToolThinkingProvider {
|
||||
call_count: AtomicUsize,
|
||||
}
|
||||
|
||||
impl MultiToolThinkingProvider {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
call_count: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl goose::providers::base::ProviderDescriptor for MultiToolThinkingProvider {
|
||||
fn metadata() -> ProviderMetadata {
|
||||
ProviderMetadata {
|
||||
name: "multi-tool-thinking-mock".to_string(),
|
||||
display_name: "Multi Tool Thinking Mock".to_string(),
|
||||
description: "Mock for multi-tool thinking preservation".to_string(),
|
||||
default_model: "mock-model".to_string(),
|
||||
known_models: vec![],
|
||||
model_doc_link: "".to_string(),
|
||||
config_keys: vec![],
|
||||
setup_steps: vec![],
|
||||
model_selection_hint: None,
|
||||
fast_model: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderDef for MultiToolThinkingProvider {
|
||||
type Provider = Self;
|
||||
|
||||
fn from_env(
|
||||
_extensions: Vec<goose::config::ExtensionConfig>,
|
||||
_tls_config: Option<goose::providers::api_client::TlsConfig>,
|
||||
) -> futures::future::BoxFuture<'static, anyhow::Result<Self>> {
|
||||
unimplemented!()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for MultiToolThinkingProvider {
|
||||
async fn stream(
|
||||
&self,
|
||||
_model_config: &ModelConfig,
|
||||
_session_id: &str,
|
||||
_system_prompt: &str,
|
||||
_messages: &[Message],
|
||||
_tools: &[Tool],
|
||||
) -> Result<MessageStream, ProviderError> {
|
||||
let call = self.call_count.fetch_add(1, Ordering::SeqCst);
|
||||
let usage = ProviderUsage::new(
|
||||
"mock-model".to_string(),
|
||||
Usage::new(Some(10), Some(20), Some(30)),
|
||||
);
|
||||
match call {
|
||||
0 => {
|
||||
// Chunk 1: reasoning only (no tool calls)
|
||||
let thinking =
|
||||
Message::assistant().with_thinking("multi-tool reasoning", "sig_0");
|
||||
// Chunk 2: two tool calls, no reasoning — the multi-tool bug scenario
|
||||
let tc1 = CallToolRequestParams::new("tool_a")
|
||||
.with_arguments(object!({"p": "1"}));
|
||||
let tc2 = CallToolRequestParams::new("tool_b")
|
||||
.with_arguments(object!({"p": "2"}));
|
||||
let tool_msg = Message::assistant()
|
||||
.with_tool_request("call_1", Ok(tc1))
|
||||
.with_tool_request("call_2", Ok(tc2));
|
||||
let stream = futures::stream::iter(vec![
|
||||
Ok((Some(thinking), None)),
|
||||
Ok((Some(tool_msg), Some(usage))),
|
||||
]);
|
||||
Ok(Box::pin(stream))
|
||||
}
|
||||
_ => {
|
||||
let msg = Message::assistant().with_text("Done.");
|
||||
Ok(Box::pin(futures::stream::once(async move {
|
||||
Ok((Some(msg), Some(usage)))
|
||||
})))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_name(&self) -> &str {
|
||||
"multi-tool-thinking-mock"
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_reasoning_preserved_on_all_tool_calls_when_thinking_in_separate_chunk(
|
||||
) -> Result<()> {
|
||||
use goose_providers::formats::openai::{
|
||||
format_messages_with_options, OpenAiFormatOptions,
|
||||
};
|
||||
use goose_providers::images::ImageFormat;
|
||||
|
||||
let temp_dir = tempfile::tempdir()?;
|
||||
let session_manager = Arc::new(SessionManager::new(temp_dir.path().to_path_buf()));
|
||||
let config = AgentConfig::new(
|
||||
session_manager.clone(),
|
||||
PermissionManager::instance(),
|
||||
None,
|
||||
GooseMode::Auto,
|
||||
true,
|
||||
GoosePlatform::GooseCli,
|
||||
);
|
||||
let agent = Agent::with_config(config);
|
||||
let provider = Arc::new(MultiToolThinkingProvider::new());
|
||||
|
||||
let session = session_manager
|
||||
.create_session(
|
||||
PathBuf::default(),
|
||||
"multi-tool-thinking-test".to_string(),
|
||||
SessionType::Hidden,
|
||||
GooseMode::default(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let session_id = session.id.clone();
|
||||
agent
|
||||
.update_provider(provider, ModelConfig::new("mock-model"), &session_id)
|
||||
.await?;
|
||||
|
||||
let session_config = SessionConfig {
|
||||
id: session_id.clone(),
|
||||
schedule_id: None,
|
||||
max_turns: Some(2),
|
||||
retry_config: None,
|
||||
};
|
||||
|
||||
let reply_stream = agent
|
||||
.reply(
|
||||
Message::user().with_text("Use both tools"),
|
||||
session_config,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
tokio::pin!(reply_stream);
|
||||
while let Some(event) = reply_stream.next().await {
|
||||
event?;
|
||||
}
|
||||
|
||||
let reloaded = session_manager.get_session(&session_id, true).await?;
|
||||
let messages = reloaded
|
||||
.conversation
|
||||
.expect("should have conversation")
|
||||
.messages()
|
||||
.to_vec();
|
||||
|
||||
let spec = format_messages_with_options(
|
||||
&messages,
|
||||
&ImageFormat::OpenAi,
|
||||
OpenAiFormatOptions {
|
||||
preserve_thinking_context: true,
|
||||
},
|
||||
);
|
||||
|
||||
// Both tool calls must end up in one merged assistant message with reasoning_content.
|
||||
let assistant_msgs: Vec<_> = spec
|
||||
.iter()
|
||||
.filter(|m| m.get("role") == Some(&serde_json::json!("assistant")))
|
||||
.filter(|m| {
|
||||
m.get("tool_calls")
|
||||
.and_then(|tc| tc.as_array())
|
||||
.is_some_and(|a| !a.is_empty())
|
||||
})
|
||||
.collect();
|
||||
|
||||
assert_eq!(
|
||||
assistant_msgs.len(),
|
||||
1,
|
||||
"both tool calls must be merged into one assistant message"
|
||||
);
|
||||
assert_eq!(
|
||||
assistant_msgs[0]["reasoning_content"], "multi-tool reasoning",
|
||||
"merged message must carry reasoning_content"
|
||||
);
|
||||
let tool_calls = assistant_msgs[0]["tool_calls"].as_array().unwrap();
|
||||
assert_eq!(tool_calls.len(), 2, "both tool calls must be present");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod goal_checking_tests {
|
||||
use super::*;
|
||||
|
||||
Reference in New Issue
Block a user