agent: Preserve thinking blocks between requests (#29055)

Looks like the required backend component of this was deployed.

https://github.com/zed-industries/monorepo/actions/runs/14541199197

Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>
Co-authored-by: Agus Zubiaga <hi@aguz.me>
Co-authored-by: Richard Feldman <oss@rtfeldman.com>
Co-authored-by: Nathan Sobo <nathan@zed.dev>

crates/agent/src/active_thread.rs

@@ -133,18 +133,23 @@ impl RenderedMessage {
     }
 
     fn push_segment(&mut self, segment: &MessageSegment, cx: &mut App) {
-        let rendered_segment = match segment {
-            MessageSegment::Thinking(text) => RenderedMessageSegment::Thinking {
-                content: parse_markdown(text.into(), self.language_registry.clone(), cx),
-                scroll_handle: ScrollHandle::default(),
-            },
-            MessageSegment::Text(text) => RenderedMessageSegment::Text(parse_markdown(
-                text.into(),
-                self.language_registry.clone(),
-                cx,
-            )),
+        match segment {
+            MessageSegment::Thinking { text, .. } => {
+                self.segments.push(RenderedMessageSegment::Thinking {
+                    content: parse_markdown(text.into(), self.language_registry.clone(), cx),
+                    scroll_handle: ScrollHandle::default(),
+                })
+            }
+            MessageSegment::Text(text) => {
+                self.segments
+                    .push(RenderedMessageSegment::Text(parse_markdown(
+                        text.into(),
+                        self.language_registry.clone(),
+                        cx,
+                    )))
+            }
+            MessageSegment::RedactedThinking(_) => {}
         };
-        self.segments.push(rendered_segment);
     }
 }
 

crates/agent/src/thread.rs

@@ -113,12 +113,21 @@ impl Message {
         self.segments.iter().all(|segment| segment.should_display())
     }
 
-    pub fn push_thinking(&mut self, text: &str) {
-        if let Some(MessageSegment::Thinking(segment)) = self.segments.last_mut() {
+    pub fn push_thinking(&mut self, text: &str, signature: Option<String>) {
+        if let Some(MessageSegment::Thinking {
+            text: segment,
+            signature: current_signature,
+        }) = self.segments.last_mut()
+        {
+            if let Some(signature) = signature {
+                *current_signature = Some(signature);
+            }
             segment.push_str(text);
         } else {
-            self.segments
-                .push(MessageSegment::Thinking(text.to_string()));
+            self.segments.push(MessageSegment::Thinking {
+                text: text.to_string(),
+                signature,
+            });
         }
     }
 
@@ -140,11 +149,12 @@ impl Message {
         for segment in &self.segments {
             match segment {
                 MessageSegment::Text(text) => result.push_str(text),
-                MessageSegment::Thinking(text) => {
-                    result.push_str("<think>");
+                MessageSegment::Thinking { text, .. } => {
+                    result.push_str("<think>\n");
                     result.push_str(text);
-                    result.push_str("</think>");
+                    result.push_str("\n</think>");
                 }
+                MessageSegment::RedactedThinking(_) => {}
             }
         }
 
@@ -155,24 +165,22 @@ impl Message {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum MessageSegment {
     Text(String),
-    Thinking(String),
+    Thinking {
+        text: String,
+        signature: Option<String>,
+    },
+    RedactedThinking(Vec<u8>),
 }
 
 impl MessageSegment {
-    pub fn text_mut(&mut self) -> &mut String {
-        match self {
-            Self::Text(text) => text,
-            Self::Thinking(text) => text,
-        }
-    }
-
     pub fn should_display(&self) -> bool {
         // We add USING_TOOL_MARKER when making a request that includes tool uses
         // without non-whitespace text around them, and this can cause the model
         // to mimic the pattern, so we consider those segments not displayable.
         match self {
             Self::Text(text) => text.is_empty() || text.trim() == USING_TOOL_MARKER,
-            Self::Thinking(text) => text.is_empty() || text.trim() == USING_TOOL_MARKER,
+            Self::Thinking { text, .. } => text.is_empty() || text.trim() == USING_TOOL_MARKER,
+            Self::RedactedThinking(_) => false,
         }
     }
 }
@@ -408,8 +416,11 @@ impl Thread {
                         .into_iter()
                         .map(|segment| match segment {
                             SerializedMessageSegment::Text { text } => MessageSegment::Text(text),
-                            SerializedMessageSegment::Thinking { text } => {
-                                MessageSegment::Thinking(text)
+                            SerializedMessageSegment::Thinking { text, signature } => {
+                                MessageSegment::Thinking { text, signature }
+                            }
+                            SerializedMessageSegment::RedactedThinking { data } => {
+                                MessageSegment::RedactedThinking(data)
                             }
                         })
                         .collect(),
@@ -862,9 +873,10 @@ impl Thread {
             for segment in &message.segments {
                 match segment {
                     MessageSegment::Text(content) => text.push_str(content),
-                    MessageSegment::Thinking(content) => {
+                    MessageSegment::Thinking { text: content, .. } => {
                         text.push_str(&format!("<think>{}</think>", content))
                     }
+                    MessageSegment::RedactedThinking(_) => {}
                 }
             }
             text.push('\n');
@@ -894,8 +906,16 @@ impl Thread {
                                 MessageSegment::Text(text) => {
                                     SerializedMessageSegment::Text { text: text.clone() }
                                 }
-                                MessageSegment::Thinking(text) => {
-                                    SerializedMessageSegment::Thinking { text: text.clone() }
+                                MessageSegment::Thinking { text, signature } => {
+                                    SerializedMessageSegment::Thinking {
+                                        text: text.clone(),
+                                        signature: signature.clone(),
+                                    }
+                                }
+                                MessageSegment::RedactedThinking(data) => {
+                                    SerializedMessageSegment::RedactedThinking {
+                                        data: data.clone(),
+                                    }
                                 }
                             })
                             .collect(),
@@ -1038,10 +1058,35 @@ impl Thread {
                 }
             }
 
-            if !message.segments.is_empty() {
+            if !message.context.is_empty() {
                 request_message
                     .content
-                    .push(MessageContent::Text(message.to_string()));
+                    .push(MessageContent::Text(message.context.to_string()));
+            }
+
+            for segment in &message.segments {
+                match segment {
+                    MessageSegment::Text(text) => {
+                        if !text.is_empty() {
+                            request_message
+                                .content
+                                .push(MessageContent::Text(text.into()));
+                        }
+                    }
+                    MessageSegment::Thinking { text, signature } => {
+                        if !text.is_empty() {
+                            request_message.content.push(MessageContent::Thinking {
+                                text: text.into(),
+                                signature: signature.clone(),
+                            });
+                        }
+                    }
+                    MessageSegment::RedactedThinking(data) => {
+                        request_message
+                            .content
+                            .push(MessageContent::RedactedThinking(data.clone()));
+                    }
+                };
             }
 
             match request_kind {
@@ -1187,10 +1232,13 @@ impl Thread {
                                     };
                                 }
                             }
-                            LanguageModelCompletionEvent::Thinking(chunk) => {
+                            LanguageModelCompletionEvent::Thinking {
+                                text: chunk,
+                                signature,
+                            } => {
                                 if let Some(last_message) = thread.messages.last_mut() {
                                     if last_message.role == Role::Assistant {
-                                        last_message.push_thinking(&chunk);
+                                        last_message.push_thinking(&chunk, signature);
                                         cx.emit(ThreadEvent::StreamedAssistantThinking(
                                             last_message.id,
                                             chunk,
@@ -1203,7 +1251,10 @@ impl Thread {
                                         // will result in duplicating the text of the chunk in the rendered Markdown.
                                         thread.insert_message(
                                             Role::Assistant,
-                                            vec![MessageSegment::Thinking(chunk.to_string())],
+                                            vec![MessageSegment::Thinking {
+                                                text: chunk.to_string(),
+                                                signature,
+                                            }],
                                             cx,
                                         );
                                     };
@@ -1893,9 +1944,10 @@ impl Thread {
             for segment in &message.segments {
                 match segment {
                     MessageSegment::Text(text) => writeln!(markdown, "{}\n", text)?,
-                    MessageSegment::Thinking(text) => {
-                        writeln!(markdown, "<think>{}</think>\n", text)?
+                    MessageSegment::Thinking { text, .. } => {
+                        writeln!(markdown, "<think>\n{}\n</think>\n", text)?
                     }
+                    MessageSegment::RedactedThinking(_) => {}
                 }
             }
 

crates/agent/src/thread_store.rs

@@ -660,9 +660,18 @@ pub struct SerializedMessage {
 #[serde(tag = "type")]
 pub enum SerializedMessageSegment {
     #[serde(rename = "text")]
-    Text { text: String },
+    Text {
+        text: String,
+    },
     #[serde(rename = "thinking")]
-    Thinking { text: String },
+    Thinking {
+        text: String,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        signature: Option<String>,
+    },
+    RedactedThinking {
+        data: Vec<u8>,
+    },
 }
 
 #[derive(Debug, Serialize, Deserialize)]