Add Message Metadata for Visibility Control (#4538)

crates/goose-server/src/openapi.rs

@@ -14,8 +14,9 @@ use rmcp::model::{
 use utoipa::{OpenApi, ToSchema};
 
 use goose::conversation::message::{
-    ContextLengthExceeded, FrontendToolRequest, Message, MessageContent, RedactedThinkingContent,
-    SummarizationRequested, ThinkingContent, ToolConfirmationRequest, ToolRequest, ToolResponse,
+    ContextLengthExceeded, FrontendToolRequest, Message, MessageContent, MessageMetadata,
+    RedactedThinkingContent, SummarizationRequested, ThinkingContent, ToolConfirmationRequest,
+    ToolRequest, ToolResponse,
 };
 use utoipa::openapi::schema::{
     AdditionalProperties, AnyOfBuilder, ArrayBuilder, ObjectBuilder, OneOfBuilder, Schema,
@@ -421,6 +422,7 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
         super::routes::session::SessionHistoryResponse,
         Message,
         MessageContent,
+        MessageMetadata,
         ContentSchema,
         EmbeddedResourceSchema,
         ImageContentSchema,

crates/goose-server/src/routes/context.rs

@@ -62,7 +62,12 @@ async fn manage_context(
     }
 
     Ok(Json(ContextManageResponse {
-        messages: processed_messages.messages().clone(),
+        messages: processed_messages
+            .messages()
+            .iter()
+            .filter(|m| m.is_user_visible())
+            .cloned()
+            .collect(),
         token_counts,
     }))
 }

crates/goose-server/src/routes/reply.rs

@@ -309,7 +309,11 @@ async fn reply_handler(
                             }
 
                             all_messages.push(message.clone());
-                            stream_event(MessageEvent::Message { message }, &tx, &cancel_token).await;
+
+                            // Only send message to client if it's user_visible
+                            if message.is_user_visible() {
+                                stream_event(MessageEvent::Message { message }, &tx, &cancel_token).await;
+                            }
                         }
                         Ok(Some(Ok(AgentEvent::HistoryReplaced(new_messages)))) => {
                             // Replace the message history with the compacted messages

crates/goose-server/src/routes/session.rs

@@ -124,10 +124,18 @@ async fn get_session_history(
         }
     };
 
+    // Filter messages to only include user_visible ones
+    let user_visible_messages: Vec<Message> = messages
+        .messages()
+        .iter()
+        .filter(|m| m.is_user_visible())
+        .cloned()
+        .collect();
+
     Ok(Json(SessionHistoryResponse {
         session_id,
         metadata,
-        messages: messages.messages().clone(),
+        messages: user_visible_messages,
     }))
 }
 

crates/goose/src/agents/context.rs

@@ -1,6 +1,6 @@
 use anyhow::Ok;
 
-use crate::conversation::message::Message;
+use crate::conversation::message::{Message, MessageMetadata};
 use crate::conversation::Conversation;
 use crate::token_counter::create_async_token_counter;
 
@@ -64,17 +64,8 @@ impl Agent {
         let provider = self.provider().await?;
         let summary_result = summarize_messages(provider.clone(), messages).await?;
 
-        let (mut new_messages, mut new_token_counts, summarization_usage) = match summary_result {
-            Some((summary_message, provider_usage)) => {
-                // For token counting purposes, we use the output tokens (the actual summary content)
-                // since that's what will be in the context going forward
-                let total_tokens = provider_usage.usage.output_tokens.unwrap_or(0) as usize;
-                (
-                    vec![summary_message],
-                    vec![total_tokens],
-                    Some(provider_usage),
-                )
-            }
+        let (summary_message, summarization_usage) = match summary_result {
+            Some((summary_message, provider_usage)) => (summary_message, Some(provider_usage)),
             None => {
                 // No summary was generated (empty input)
                 tracing::warn!("Summarization failed. Returning empty messages.");
@@ -82,30 +73,55 @@ impl Agent {
             }
         };
 
-        // Add an assistant message to the summarized messages to ensure the assistant's response is included in the context.
-        if new_messages.len() == 1 {
-            let compaction_marker = Message::assistant()
-                .with_summarization_requested("Conversation compacted and summarized");
-            let compaction_marker_tokens: usize = 8;
+        // Create the final message list with updated visibility metadata:
+        // 1. Original messages become user_visible but not agent_visible
+        // 2. Summary message becomes agent_visible but not user_visible
+        // 3. Assistant messages to continue the conversation remain both user_visible and agent_visible
 
-            // Insert the marker before the summary message
-            new_messages.insert(0, compaction_marker);
-            new_token_counts.insert(0, compaction_marker_tokens);
+        let mut final_messages = Vec::new();
+        let mut final_token_counts = Vec::new();
 
-            // Add an assistant message to continue the conversation
-            let assistant_message = Message::assistant().with_text("
-                The previous message contains a summary that was prepared because a context limit was reached.
-                Do not mention that you read a summary or that conversation summarization occurred
-                Just continue the conversation naturally based on the summarized context
-            ");
-            let assistant_message_tokens: usize = 41;
-            new_messages.push(assistant_message);
-            new_token_counts.push(assistant_message_tokens);
+        // Add all original messages with updated visibility (preserve user_visible, set agent_visible=false)
+        for msg in messages.iter().cloned() {
+            let updated_metadata = msg.metadata.with_agent_invisible();
+            let updated_msg = msg.with_metadata(updated_metadata);
+            final_messages.push(updated_msg);
+            // Token count doesn't matter for agent_visible=false messages, but we'll use 0
+            final_token_counts.push(0);
         }
 
+        // Add the compaction marker (user_visible=true, agent_visible=false)
+        let compaction_marker = Message::assistant()
+            .with_summarization_requested("Conversation compacted and summarized")
+            .with_metadata(MessageMetadata::user_only());
+        let compaction_marker_tokens: usize = 0; // Not counted since agent_visible=false
+        final_messages.push(compaction_marker);
+        final_token_counts.push(compaction_marker_tokens);
+
+        // Add the summary message (agent_visible=true, user_visible=false)
+        let summary_msg = summary_message.with_metadata(MessageMetadata::agent_only());
+        // For token counting purposes, we use the output tokens (the actual summary content)
+        // since that's what will be in the context going forward
+        let summary_tokens = summarization_usage
+            .as_ref()
+            .and_then(|usage| usage.usage.output_tokens)
+            .unwrap_or(0) as usize;
+        final_messages.push(summary_msg);
+        final_token_counts.push(summary_tokens);
+
+        // Add an assistant message to continue the conversation (agent_visible=true, user_visible=false)
+        let assistant_message = Message::assistant().with_text("
+            The previous message contains a summary that was prepared because a context limit was reached.
+            Do not mention that you read a summary or that conversation summarization occurred
+            Just continue the conversation naturally based on the summarized context
+        ").with_metadata(MessageMetadata::agent_only());
+        let assistant_message_tokens: usize = 0; // Not counted since it's for agent context only
+        final_messages.push(assistant_message);
+        final_token_counts.push(assistant_message_tokens);
+
         Ok((
-            Conversation::new_unvalidated(new_messages),
-            new_token_counts,
+            Conversation::new_unvalidated(final_messages),
+            final_token_counts,
             summarization_usage,
         ))
     }

crates/goose/src/agents/reply_parts.rs

@@ -245,12 +245,13 @@ impl Agent {
             }
         }
 
-        let filtered_message = Message {
-            id: response.id.clone(),
-            role: response.role.clone(),
-            created: response.created,
-            content: filtered_content,
-        };
+        let mut filtered_message =
+            Message::new(response.role.clone(), response.created, filtered_content);
+
+        // Preserve the ID if it exists
+        if let Some(id) = response.id.clone() {
+            filtered_message = filtered_message.with_id(id);
+        }
 
         // Categorize tool requests
         let mut frontend_requests = Vec::new();

crates/goose/src/context_mgmt/auto_compact.rs

@@ -207,8 +207,35 @@ pub async fn check_and_compact_messages(
         check_result.usage_ratio * 100.0
     );
 
-    // Use perform_compaction to do the actual work
-    perform_compaction(agent, messages).await
+    // Check if the most recent message is a user message
+    let (messages_to_compact, preserved_user_message) = if let Some(last_message) = messages.last()
+    {
+        if matches!(last_message.role, rmcp::model::Role::User) {
+            // Remove the last user message before auto-compaction
+            (&messages[..messages.len() - 1], Some(last_message.clone()))
+        } else {
+            (messages, None)
+        }
+    } else {
+        (messages, None)
+    };
+
+    // Perform the compaction on messages excluding the preserved user message
+    // The summarize_context method already handles the visibility properly
+    let (mut summary_messages, _, summarization_usage) =
+        agent.summarize_context(messages_to_compact).await?;
+
+    // Add back the preserved user message if it exists
+    // (keeps default visibility: both true)
+    if let Some(user_message) = preserved_user_message {
+        summary_messages.push(user_message);
+    }
+
+    Ok(AutoCompactResult {
+        compacted: true,
+        messages: summary_messages,
+        summarization_usage,
+    })
 }
 
 #[cfg(test)]
@@ -455,8 +482,9 @@ mod tests {
             );
         }
 
-        // Should have fewer messages (summarized)
-        assert!(result.messages.len() <= messages.len());
+        // After visibility implementation, we keep all messages plus summary
+        // Original messages become user_visible only, summary becomes agent_visible only
+        assert!(result.messages.len() > messages.len());
     }
 
     #[tokio::test]
@@ -641,8 +669,9 @@ mod tests {
         // Verify the compacted messages are returned
         assert!(!result.messages.is_empty());
 
-        // Should have fewer messages after compaction
-        assert!(result.messages.len() <= messages.len());
+        // After visibility implementation, we keep all messages plus summary
+        // Original messages become user_visible only, summary becomes agent_visible only
+        assert!(result.messages.len() > messages.len());
     }
 
     #[tokio::test]

crates/goose/src/conversation/message.rs

@@ -373,6 +373,89 @@ impl From<PromptMessage> for Message {
     }
 }
 
+#[derive(ToSchema, Clone, Copy, PartialEq, Serialize, Deserialize)]
+/// Metadata for message visibility
+#[serde(rename_all = "camelCase")]
+pub struct MessageMetadata {
+    /// Whether the message should be visible to the user in the UI
+    #[serde(default = "default_true")]
+    pub user_visible: bool,
+    /// Whether the message should be included in the agent's context window
+    #[serde(default = "default_true")]
+    pub agent_visible: bool,
+}
+
+impl Default for MessageMetadata {
+    fn default() -> Self {
+        MessageMetadata {
+            user_visible: true,
+            agent_visible: true,
+        }
+    }
+}
+
+impl MessageMetadata {
+    /// Create metadata for messages visible only to the agent
+    pub fn agent_only() -> Self {
+        MessageMetadata {
+            user_visible: false,
+            agent_visible: true,
+        }
+    }
+
+    /// Create metadata for messages visible only to the user
+    pub fn user_only() -> Self {
+        MessageMetadata {
+            user_visible: true,
+            agent_visible: false,
+        }
+    }
+
+    /// Create metadata for messages visible to neither user nor agent (archived)
+    pub fn invisible() -> Self {
+        MessageMetadata {
+            user_visible: false,
+            agent_visible: false,
+        }
+    }
+
+    /// Return a copy with agent_visible set to false
+    pub fn with_agent_invisible(self) -> Self {
+        Self {
+            agent_visible: false,
+            ..self
+        }
+    }
+
+    /// Return a copy with user_visible set to false
+    pub fn with_user_invisible(self) -> Self {
+        Self {
+            user_visible: false,
+            ..self
+        }
+    }
+
+    /// Return a copy with agent_visible set to true
+    pub fn with_agent_visible(self) -> Self {
+        Self {
+            agent_visible: true,
+            ..self
+        }
+    }
+
+    /// Return a copy with user_visible set to true
+    pub fn with_user_visible(self) -> Self {
+        Self {
+            user_visible: true,
+            ..self
+        }
+    }
+}
+
+fn default_true() -> bool {
+    true
+}
+
 #[derive(ToSchema, Clone, PartialEq, Serialize, Deserialize)]
 /// A message to or from an LLM
 #[serde(rename_all = "camelCase")]
@@ -383,6 +466,8 @@ pub struct Message {
     pub created: i64,
     #[serde(deserialize_with = "deserialize_sanitized_content")]
     pub content: Vec<MessageContent>,
+    #[serde(default)]
+    pub metadata: MessageMetadata,
 }
 
 impl fmt::Debug for Message {
@@ -409,6 +494,7 @@ impl Message {
             role,
             created,
             content,
+            metadata: MessageMetadata::default(),
         }
     }
     pub fn debug(&self) -> String {
@@ -422,6 +508,7 @@ impl Message {
             role: Role::User,
             created: Utc::now().timestamp(),
             content: Vec::new(),
+            metadata: MessageMetadata::default(),
         }
     }
 
@@ -432,6 +519,7 @@ impl Message {
             role: Role::Assistant,
             created: Utc::now().timestamp(),
             content: Vec::new(),
+            metadata: MessageMetadata::default(),
         }
     }
 
@@ -597,11 +685,48 @@ impl Message {
     pub fn with_summarization_requested<S: Into<String>>(self, msg: S) -> Self {
         self.with_content(MessageContent::summarization_requested(msg))
     }
+
+    /// Set the visibility metadata for the message
+    pub fn with_visibility(mut self, user_visible: bool, agent_visible: bool) -> Self {
+        self.metadata.user_visible = user_visible;
+        self.metadata.agent_visible = agent_visible;
+        self
+    }
+
+    /// Set the entire metadata for the message
+    pub fn with_metadata(mut self, metadata: MessageMetadata) -> Self {
+        self.metadata = metadata;
+        self
+    }
+
+    /// Mark the message as only visible to the user (not the agent)
+    pub fn user_only(mut self) -> Self {
+        self.metadata.user_visible = true;
+        self.metadata.agent_visible = false;
+        self
+    }
+
+    /// Mark the message as only visible to the agent (not the user)
+    pub fn agent_only(mut self) -> Self {
+        self.metadata.user_visible = false;
+        self.metadata.agent_visible = true;
+        self
+    }
+
+    /// Check if the message is visible to the user
+    pub fn is_user_visible(&self) -> bool {
+        self.metadata.user_visible
+    }
+
+    /// Check if the message is visible to the agent
+    pub fn is_agent_visible(&self) -> bool {
+        self.metadata.agent_visible
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::conversation::message::{Message, MessageContent};
+    use crate::conversation::message::{Message, MessageContent, MessageMetadata};
     use crate::conversation::*;
     use mcp_core::ToolCall;
     use rmcp::model::{
@@ -953,4 +1078,148 @@ mod tests {
 
         assert_eq!(message.as_concat_text(), "Hello world 世界 🌍");
     }
+
+    #[test]
+    fn test_message_metadata_defaults() {
+        let message = Message::user().with_text("Test");
+
+        // By default, messages should be both user and agent visible
+        assert!(message.is_user_visible());
+        assert!(message.is_agent_visible());
+    }
+
+    #[test]
+    fn test_message_visibility_methods() {
+        // Test user_only
+        let user_only_msg = Message::user().with_text("User only").user_only();
+        assert!(user_only_msg.is_user_visible());
+        assert!(!user_only_msg.is_agent_visible());
+
+        // Test agent_only
+        let agent_only_msg = Message::assistant().with_text("Agent only").agent_only();
+        assert!(!agent_only_msg.is_user_visible());
+        assert!(agent_only_msg.is_agent_visible());
+
+        // Test with_visibility
+        let custom_msg = Message::user()
+            .with_text("Custom visibility")
+            .with_visibility(false, true);
+        assert!(!custom_msg.is_user_visible());
+        assert!(custom_msg.is_agent_visible());
+    }
+
+    #[test]
+    fn test_message_metadata_serialization() {
+        let message = Message::user()
+            .with_text("Test message")
+            .with_visibility(false, true);
+
+        let json_str = serde_json::to_string(&message).unwrap();
+        let value: Value = serde_json::from_str(&json_str).unwrap();
+
+        assert_eq!(value["metadata"]["userVisible"], false);
+        assert_eq!(value["metadata"]["agentVisible"], true);
+    }
+
+    #[test]
+    fn test_message_metadata_deserialization() {
+        // Test with explicit metadata
+        let json_with_metadata = r#"{
+            "role": "user",
+            "created": 1640995200,
+            "content": [{
+                "type": "text",
+                "text": "Test"
+            }],
+            "metadata": {
+                "userVisible": false,
+                "agentVisible": true
+            }
+        }"#;
+
+        let message: Message = serde_json::from_str(json_with_metadata).unwrap();
+        assert!(!message.is_user_visible());
+        assert!(message.is_agent_visible());
+
+        // Test without metadata (should use defaults)
+        let json_without_metadata = r#"{
+            "role": "user",
+            "created": 1640995200,
+            "content": [{
+                "type": "text",
+                "text": "Test"
+            }]
+        }"#;
+
+        let message: Message = serde_json::from_str(json_without_metadata).unwrap();
+        assert!(message.is_user_visible());
+        assert!(message.is_agent_visible());
+    }
+
+    #[test]
+    fn test_message_metadata_static_methods() {
+        // Test MessageMetadata::agent_only()
+        let agent_only_metadata = MessageMetadata::agent_only();
+        assert!(!agent_only_metadata.user_visible);
+        assert!(agent_only_metadata.agent_visible);
+
+        // Test MessageMetadata::user_only()
+        let user_only_metadata = MessageMetadata::user_only();
+        assert!(user_only_metadata.user_visible);
+        assert!(!user_only_metadata.agent_visible);
+
+        // Test MessageMetadata::invisible()
+        let invisible_metadata = MessageMetadata::invisible();
+        assert!(!invisible_metadata.user_visible);
+        assert!(!invisible_metadata.agent_visible);
+
+        // Test using them with messages
+        let agent_msg = Message::assistant()
+            .with_text("Agent only message")
+            .with_metadata(MessageMetadata::agent_only());
+        assert!(!agent_msg.is_user_visible());
+        assert!(agent_msg.is_agent_visible());
+
+        let user_msg = Message::user()
+            .with_text("User only message")
+            .with_metadata(MessageMetadata::user_only());
+        assert!(user_msg.is_user_visible());
+        assert!(!user_msg.is_agent_visible());
+
+        let invisible_msg = Message::user()
+            .with_text("Invisible message")
+            .with_metadata(MessageMetadata::invisible());
+        assert!(!invisible_msg.is_user_visible());
+        assert!(!invisible_msg.is_agent_visible());
+    }
+
+    #[test]
+    fn test_message_metadata_builder_methods() {
+        // Test with_agent_invisible
+        let metadata = MessageMetadata::default().with_agent_invisible();
+        assert!(metadata.user_visible);
+        assert!(!metadata.agent_visible);
+
+        // Test with_user_invisible
+        let metadata = MessageMetadata::default().with_user_invisible();
+        assert!(!metadata.user_visible);
+        assert!(metadata.agent_visible);
+
+        // Test with_agent_visible
+        let metadata = MessageMetadata::invisible().with_agent_visible();
+        assert!(!metadata.user_visible);
+        assert!(metadata.agent_visible);
+
+        // Test with_user_visible
+        let metadata = MessageMetadata::invisible().with_user_visible();
+        assert!(metadata.user_visible);
+        assert!(!metadata.agent_visible);
+
+        // Test chaining
+        let metadata = MessageMetadata::invisible()
+            .with_user_visible()
+            .with_agent_visible();
+        assert!(metadata.user_visible);
+        assert!(metadata.agent_visible);
+    }
 }

crates/goose/src/providers/base.rs

@@ -328,6 +328,7 @@ pub trait Provider: Send + Sync {
     ) -> Result<(Message, ProviderUsage), ProviderError>;
 
     // Default implementation: use the provider's configured model
+    // This method filters messages to only include agent_visible ones
     async fn complete(
         &self,
         system: &str,
@@ -335,11 +336,20 @@ pub trait Provider: Send + Sync {
         tools: &[Tool],
     ) -> Result<(Message, ProviderUsage), ProviderError> {
         let model_config = self.get_model_config();
-        self.complete_with_model(&model_config, system, messages, tools)
+
+        // Filter messages to only include agent_visible ones
+        let agent_visible_messages: Vec<Message> = messages
+            .iter()
+            .filter(|m| m.is_agent_visible())
+            .cloned()
+            .collect();
+
+        self.complete_with_model(&model_config, system, &agent_visible_messages, tools)
             .await
     }
 
     // Check if a fast model is configured, otherwise fall back to regular model
+    // This method filters messages to only include agent_visible ones
     async fn complete_fast(
         &self,
         system: &str,
@@ -349,8 +359,15 @@ pub trait Provider: Send + Sync {
         let model_config = self.get_model_config();
         let fast_config = model_config.use_fast_model();
 
+        // Filter messages to only include agent_visible ones
+        let agent_visible_messages: Vec<Message> = messages
+            .iter()
+            .filter(|m| m.is_agent_visible())
+            .cloned()
+            .collect();
+
         match self
-            .complete_with_model(&fast_config, system, messages, tools)
+            .complete_with_model(&fast_config, system, &agent_visible_messages, tools)
             .await
         {
             Ok(result) => Ok(result),
@@ -362,7 +379,7 @@ pub trait Provider: Send + Sync {
                         e,
                         model_config.model_name
                     );
-                    self.complete_with_model(&model_config, system, messages, tools)
+                    self.complete_with_model(&model_config, system, &agent_visible_messages, tools)
                         .await
                 } else {
                     Err(e)

crates/goose/src/providers/claude_code.rs

@@ -282,12 +282,11 @@ impl ClaudeCodeProvider {
 
         let message_content = vec![MessageContent::text(combined_text)];
 
-        let response_message = Message {
-            id: None,
-            role: Role::Assistant,
-            created: chrono::Utc::now().timestamp(),
-            content: message_content,
-        };
+        let response_message = Message::new(
+            Role::Assistant,
+            chrono::Utc::now().timestamp(),
+            message_content,
+        );
 
         Ok((response_message, usage))
     }
@@ -433,12 +432,11 @@ impl ClaudeCodeProvider {
             println!("================================");
         }
 
-        let message = Message {
-            id: None,
-            role: Role::Assistant,
-            created: chrono::Utc::now().timestamp(),
-            content: vec![MessageContent::text(description.clone())],
-        };
+        let message = Message::new(
+            Role::Assistant,
+            chrono::Utc::now().timestamp(),
+            vec![MessageContent::text(description.clone())],
+        );
 
         let usage = Usage::default();
 

crates/goose/src/providers/cursor_agent.rs

@@ -214,12 +214,11 @@ impl CursorAgentProvider {
                         };
 
                         let message_content = vec![MessageContent::text(text_content)];
-                        let response_message = Message {
-                            id: None,
-                            role: Role::Assistant,
-                            created: chrono::Utc::now().timestamp(),
-                            content: message_content,
-                        };
+                        let response_message = Message::new(
+                            Role::Assistant,
+                            chrono::Utc::now().timestamp(),
+                            message_content,
+                        );
 
                         let usage = Usage::default();
 
@@ -233,12 +232,11 @@ impl CursorAgentProvider {
         let response_text = lines.join("\n");
 
         let message_content = vec![MessageContent::text(response_text)];
-        let response_message = Message {
-            id: None,
-            role: Role::Assistant,
-            created: chrono::Utc::now().timestamp(),
-            content: message_content,
-        };
+        let response_message = Message::new(
+            Role::Assistant,
+            chrono::Utc::now().timestamp(),
+            message_content,
+        );
         let usage = Usage::default();
 
         Ok((response_message, usage))
@@ -366,12 +364,11 @@ impl CursorAgentProvider {
             println!("================================");
         }
 
-        let message = Message {
-            id: None,
-            role: Role::Assistant,
-            created: chrono::Utc::now().timestamp(),
-            content: vec![MessageContent::text(description.clone())],
-        };
+        let message = Message::new(
+            Role::Assistant,
+            chrono::Utc::now().timestamp(),
+            vec![MessageContent::text(description.clone())],
+        );
 
         let usage = Usage::default();
 

crates/goose/src/providers/formats/openai.rs

@@ -534,23 +534,35 @@ where
                     }
                 }
 
+                let mut msg = Message::new(
+                    Role::Assistant,
+                    chrono::Utc::now().timestamp(),
+                    contents,
+                );
+
+                // Add ID if present
+                if let Some(id) = chunk.id {
+                    msg = msg.with_id(id);
+                }
+
                 yield (
-                    Some(Message {
-                        id: chunk.id,
-                        role: Role::Assistant,
-                        created: chrono::Utc::now().timestamp(),
-                        content: contents,
-                    }),
+                    Some(msg),
                     usage,
                 )
             } else if let Some(text) = &chunk.choices[0].delta.content {
+                let mut msg = Message::new(
+                    Role::Assistant,
+                    chrono::Utc::now().timestamp(),
+                    vec![MessageContent::text(text)],
+                );
+
+                // Add ID if present
+                if let Some(id) = chunk.id {
+                    msg = msg.with_id(id);
+                }
+
                 yield (
-                    Some(Message {
-                        id: chunk.id,
-                        role: Role::Assistant,
-                        created: chrono::Utc::now().timestamp(),
-                        content: vec![MessageContent::text(text)],
-                    }),
+                    Some(msg),
                     if chunk.choices[0].finish_reason.is_some() {
                         usage
                     } else {