feat: Improved support for Anthropic's extended thinking

src/client/bedrock.rs

@@ -529,7 +529,7 @@ fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
         bail!("Invalid response data: {data}");
     }
 
-    let output = ChatCompletionsOutput { text, tool_calls };
+    let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
     Ok(output)
 }
 

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<bool> {
         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("<think>\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</think>\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<ChatCompletionsOu
     let mut text = String::new();
     let mut reasoning = None;
     let mut tool_calls = vec![];
+    let mut thinking = vec![];
     if let Some(list) = data["content"].as_array() {
         for item in list {
             match item["type"].as_str() {
                 Some("thinking") => {
                     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<ChatCompletionsOu
     let output = ChatCompletionsOutput {
         text: text.to_string(),
         tool_calls,
+        thinking,
     };
     Ok(output)
 }

src/client/cohere.rs

@@ -244,6 +244,6 @@ fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
     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)
 }

src/client/common.rs

@@ -295,6 +295,7 @@ pub struct ChatCompletionsData {
 pub struct ChatCompletionsOutput {
     pub text: String,
     pub tool_calls: Vec<ToolCall>,
+    pub thinking: Vec<ThinkingBlock>,
 }
 
 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()));

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<ToolResult>,

src/client/openai.rs

@@ -533,7 +533,7 @@ pub fn openai_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
     } else {
         text.to_string()
     };
-    let output = ChatCompletionsOutput { text, tool_calls };
+    let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
     Ok(output)
 }
 
@@ -696,7 +696,7 @@ pub fn openai_extract_responses(data: &Value) -> Result<ChatCompletionsOutput> {
     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(

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<ToolCall>,
+    thinking: Vec<ThinkingBlock>,
     last_tool_calls: Vec<ToolCall>,
     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<ToolCall>) {
+    pub fn take(self) -> (String, Vec<ToolCall>, Vec<ThinkingBlock>) {
         let Self {
-            buffer, tool_calls, ..
+            buffer,
+            tool_calls,
+            thinking,
+            ..
         } = self;
-        (buffer, tool_calls)
+        (buffer, tool_calls, thinking)
     }
 }
 

src/client/vertexai.rs

@@ -322,7 +322,7 @@ fn gemini_extract_chat_completions_text(data: &Value) -> Result<ChatCompletionsO
             bail!("Invalid response data: {data}");
         }
     }
-    let output = ChatCompletionsOutput { text, tool_calls };
+    let output = ChatCompletionsOutput { text, tool_calls, ..Default::default() };
     Ok(output)
 }
 

src/function/mod.rs

@@ -5,6 +5,7 @@ pub(crate) mod todo;
 pub(crate) mod user_interaction;
 
 use crate::{
+    client::ThinkingBlock,
     config::{Agent, RequestContext},
     graph,
     utils::*,
@@ -205,6 +206,8 @@ pub struct ToolResult {
     pub output: Value,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub text: Option<String>,
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub thinking: Vec<ThinkingBlock>,
 }
 
 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());
+    }
 }