From 8629c1ca151c0c3fcda8e2dccb77723845ee4233 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Fri, 17 Jul 2026 17:26:41 -0600 Subject: [PATCH] feat: Improved support for Anthropic's extended thinking --- src/client/bedrock.rs | 2 +- src/client/claude.rs | 60 +++++++++++++++++++++++++++++++++--------- src/client/cohere.rs | 2 +- src/client/common.rs | 14 +++++++--- src/client/message.rs | 11 ++++++++ src/client/openai.rs | 4 +-- src/client/stream.rs | 17 +++++++++--- src/client/vertexai.rs | 2 +- src/function/mod.rs | 33 +++++++++++++++++++++++ 9 files changed, 121 insertions(+), 24 deletions(-) diff --git a/src/client/bedrock.rs b/src/client/bedrock.rs index 73671a7..1b978de 100644 --- a/src/client/bedrock.rs +++ b/src/client/bedrock.rs @@ -529,7 +529,7 @@ fn extract_chat_completions(data: &Value) -> Result { bail!("Invalid response data: {data}"); } - let output = ChatCompletionsOutput { text, tool_calls }; + let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() }; Ok(output) } diff --git a/src/client/claude.rs b/src/client/claude.rs index 06fa603..f9854c8 100644 --- a/src/client/claude.rs +++ b/src/client/claude.rs @@ -168,12 +168,22 @@ pub async fn claude_chat_completions_streaming( let mut function_arguments = String::new(); let mut function_id = String::new(); let mut reasoning_state = 0; + let mut thinking_text = String::new(); + let mut thinking_signature = String::new(); let handle = |message: SseMessage| -> Result { let data: Value = serde_json::from_str(&message.data)?; debug!("stream-data: {data}"); if let Some(typ) = data["type"].as_str() { match typ { "content_block_start" => { + if let (Some("redacted_thinking"), Some(redacted_data)) = ( + data["content_block"]["type"].as_str(), + data["content_block"]["data"].as_str(), + ) { + handler.thinking_block(ThinkingBlock::RedactedThinking { + data: redacted_data.to_string(), + }); + } if let (Some("tool_use"), Some(name), Some(id)) = ( data["content_block"]["type"].as_str(), data["content_block"]["name"].as_str(), @@ -206,7 +216,10 @@ pub async fn claude_chat_completions_streaming( handler.text("\n")?; reasoning_state = 1; } + thinking_text.push_str(text); handler.text(text)?; + } else if let Some(signature) = data["delta"]["signature"].as_str() { + thinking_signature.push_str(signature); } else if let (true, Some(partial_json)) = ( !function_name.is_empty(), data["delta"]["partial_json"].as_str(), @@ -218,6 +231,10 @@ pub async fn claude_chat_completions_streaming( if reasoning_state == 1 { handler.text("\n\n\n")?; reasoning_state = 0; + handler.thinking_block(ThinkingBlock::Thinking { + thinking: std::mem::take(&mut thinking_text), + signature: std::mem::take(&mut thinking_signature), + }); } if !function_name.is_empty() { let arguments: Value = if function_arguments.is_empty() { @@ -313,18 +330,24 @@ pub fn claude_build_chat_completions_body( }) => { let mut assistant_parts = vec![]; let mut user_parts = vec![]; - if !text.is_empty() { - assistant_parts.push(json!({ - "type": "text", - "text": text, - })) - } - for tool_result in tool_results { - if let Some(round_text) = &tool_result.text { - assistant_parts.push(json!({ - "type": "text", - "text": round_text, - })) + for (index, tool_result) in tool_results.iter().enumerate() { + for block in &tool_result.thinking { + assistant_parts.push(json!(block)); + } + let round_text = if index == 0 && !text.is_empty() { + Some(text.as_str()) + } else { + tool_result.text.as_deref() + }; + if let Some(round_text) = round_text { + let round_text = strip_think_tag(round_text); + let round_text = round_text.trim(); + if !round_text.is_empty() { + assistant_parts.push(json!({ + "type": "text", + "text": round_text, + })) + } } assistant_parts.push(json!({ "type": "tool_use", @@ -409,12 +432,24 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result { if let Some(v) = item["thinking"].as_str() { reasoning = Some(v.to_string()); + thinking.push(ThinkingBlock::Thinking { + thinking: v.to_string(), + signature: item["signature"].as_str().unwrap_or_default().to_string(), + }); + } + } + Some("redacted_thinking") => { + if let Some(v) = item["data"].as_str() { + thinking.push(ThinkingBlock::RedactedThinking { + data: v.to_string(), + }); } } Some("text") => { @@ -453,6 +488,7 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result Result { if text.is_empty() && tool_calls.is_empty() { bail!("Invalid response data: {data}"); } - let output = ChatCompletionsOutput { text, tool_calls }; + let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() }; Ok(output) } diff --git a/src/client/common.rs b/src/client/common.rs index 1fb0435..0172e98 100644 --- a/src/client/common.rs +++ b/src/client/common.rs @@ -295,6 +295,7 @@ pub struct ChatCompletionsData { pub struct ChatCompletionsOutput { pub text: String, pub tool_calls: Vec, + pub thinking: Vec, } impl ChatCompletionsOutput { @@ -435,6 +436,7 @@ pub async fn call_chat_completions( let ChatCompletionsOutput { mut text, tool_calls, + thinking, .. } = ret; if !text.is_empty() { @@ -445,7 +447,10 @@ pub async fn call_chat_completions( ctx.app.config.print_markdown(&text)?; } } - let tool_results = eval_tool_calls(ctx, tool_calls).await?; + let mut tool_results = eval_tool_calls(ctx, tool_calls).await?; + if let Some(first) = tool_results.first_mut() { + first.thinking = thinking; + } tool_results .iter() .for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone())); @@ -479,13 +484,16 @@ pub async fn call_chat_completions_streaming( render_ret?; - let (text, tool_calls) = handler.take(); + let (text, tool_calls, thinking) = handler.take(); match send_ret { Ok(_) => { if !text.is_empty() && !text.ends_with('\n') { println!(); } - let tool_results = eval_tool_calls(ctx, tool_calls).await?; + let mut tool_results = eval_tool_calls(ctx, tool_calls).await?; + if let Some(first) = tool_results.first_mut() { + first.thinking = thinking; + } tool_results .iter() .for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone())); diff --git a/src/client/message.rs b/src/client/message.rs index 3e4074c..ddad394 100644 --- a/src/client/message.rs +++ b/src/client/message.rs @@ -188,6 +188,17 @@ pub struct ImageUrl { pub url: String, } +/// An extended-thinking block returned by Anthropic-protocol models. +/// Serialized to match the API wire format (`type: thinking` / `type: redacted_thinking`) +/// so blocks can be replayed verbatim, signature intact, in subsequent +/// tool-loop rounds as the API requires. +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ThinkingBlock { + Thinking { thinking: String, signature: String }, + RedactedThinking { data: String }, +} + #[derive(Debug, Clone, Deserialize, Serialize)] pub struct MessageContentToolCalls { pub tool_results: Vec, diff --git a/src/client/openai.rs b/src/client/openai.rs index e5c3dd8..6633be3 100644 --- a/src/client/openai.rs +++ b/src/client/openai.rs @@ -533,7 +533,7 @@ pub fn openai_extract_chat_completions(data: &Value) -> Result Result { if text.is_empty() && tool_calls.is_empty() { bail!("Invalid response data: {data}"); } - Ok(ChatCompletionsOutput { text, tool_calls }) + Ok(ChatCompletionsOutput { text, tool_calls, ..Default::default() }) } pub async fn openai_responses_streaming( diff --git a/src/client/stream.rs b/src/client/stream.rs index c4db9f5..e1e12b8 100644 --- a/src/client/stream.rs +++ b/src/client/stream.rs @@ -1,4 +1,4 @@ -use super::{ToolCall, catch_error}; +use super::{ThinkingBlock, ToolCall, catch_error}; use crate::utils::AbortSignal; use anyhow::{Context, Result, anyhow, bail}; @@ -13,6 +13,7 @@ pub struct SseHandler { abort_signal: AbortSignal, buffer: String, tool_calls: Vec, + thinking: Vec, last_tool_calls: Vec, max_call_repeats: usize, call_repeat_chain_len: usize, @@ -26,6 +27,7 @@ impl SseHandler { abort_signal, buffer: String::new(), tool_calls: Vec::new(), + thinking: Vec::new(), last_tool_calls: Vec::new(), max_call_repeats: 2, call_repeat_chain_len: 3, @@ -170,6 +172,10 @@ impl SseHandler { message } + pub fn thinking_block(&mut self, block: ThinkingBlock) { + self.thinking.push(block); + } + pub fn abort(&self) -> AbortSignal { self.abort_signal.clone() } @@ -179,11 +185,14 @@ impl SseHandler { &self.last_tool_calls } - pub fn take(self) -> (String, Vec) { + pub fn take(self) -> (String, Vec, Vec) { let Self { - buffer, tool_calls, .. + buffer, + tool_calls, + thinking, + .. } = self; - (buffer, tool_calls) + (buffer, tool_calls, thinking) } } diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs index ae0cc73..85d689b 100644 --- a/src/client/vertexai.rs +++ b/src/client/vertexai.rs @@ -322,7 +322,7 @@ fn gemini_extract_chat_completions_text(data: &Value) -> Result, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub thinking: Vec, } impl ToolResult { @@ -213,6 +216,7 @@ impl ToolResult { call, output, text: None, + thinking: vec![], } } } @@ -1914,4 +1918,33 @@ mod tests { assert_eq!(result.call.name, "my_tool"); assert_eq!(result.output, json!({"result": "ok"})); } + + #[test] + fn thinking_block_matches_anthropic_wire_format() { + let block = ThinkingBlock::Thinking { + thinking: "chain of thought".to_string(), + signature: "sig123".to_string(), + }; + assert_eq!( + serde_json::to_value(&block).unwrap(), + json!({"type": "thinking", "thinking": "chain of thought", "signature": "sig123"}) + ); + + let redacted = ThinkingBlock::RedactedThinking { + data: "opaque".to_string(), + }; + assert_eq!( + serde_json::to_value(&redacted).unwrap(), + json!({"type": "redacted_thinking", "data": "opaque"}) + ); + } + + #[test] + fn tool_result_deserializes_without_text_and_thinking() { + let yaml = "call:\n name: my_tool\n arguments: {}\noutput: ok\n"; + let result: ToolResult = serde_yaml::from_str(yaml).unwrap(); + assert_eq!(result.call.name, "my_tool"); + assert!(result.text.is_none()); + assert!(result.thinking.is_empty()); + } }