From 5ae92b6799c6ade6c69104735d85581eb5536138 Mon Sep 17 00:00:00 2001 From: Douwe Osinga Date: Fri, 17 Oct 2025 10:15:39 -0400 Subject: [PATCH] WIP --- crates/goose/src/agents/reply_parts.rs | 5 +-- crates/goose/src/conversation/message.rs | 22 +++++++----- crates/goose/src/conversation/mod.rs | 6 +--- .../goose/src/providers/formats/anthropic.rs | 25 ++++++------- crates/goose/src/session/session_manager.rs | 36 +++++++++++++------ ui/desktop/openapi.json | 5 ++- ui/desktop/src/api/types.gen.ts | 5 +-- 7 files changed, 57 insertions(+), 47 deletions(-) diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index c4dd379b68..75f71c2fb2 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -224,10 +224,7 @@ impl Agent { 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); - } + filtered_message = filtered_message.with_id(response.id.clone()); // Categorize tool requests let mut frontend_requests = Vec::new(); diff --git a/crates/goose/src/conversation/message.rs b/crates/goose/src/conversation/message.rs index b47d05774a..00814c36b2 100644 --- a/crates/goose/src/conversation/message.rs +++ b/crates/goose/src/conversation/message.rs @@ -1,4 +1,6 @@ +use crate::conversation::tool_result_serde; use crate::mcp_utils::ToolResult; +use crate::utils::sanitize_unicode_tags; use chrono::Utc; use rmcp::model::{ AnnotateAble, CallToolRequestParam, Content, ImageContent, JsonObject, PromptMessage, @@ -9,9 +11,7 @@ use serde::{Deserialize, Deserializer, Serialize}; use std::collections::HashSet; use std::fmt; use utoipa::ToSchema; - -use crate::conversation::tool_result_serde; -use crate::utils::sanitize_unicode_tags; +use uuid::Uuid; #[derive(ToSchema)] pub enum ToolCallResult { @@ -450,10 +450,9 @@ impl MessageMetadata { } #[derive(ToSchema, Clone, PartialEq, Serialize, Deserialize, Debug)] -/// A message to or from an LLM #[serde(rename_all = "camelCase")] pub struct Message { - pub id: Option, + pub id: String, pub role: Role, pub created: i64, #[serde(deserialize_with = "deserialize_sanitized_content")] @@ -462,15 +461,20 @@ pub struct Message { } impl Message { + pub fn msg_id() -> String { + format!("msg_{}", Uuid::new_v4()) + } + pub fn new(role: Role, created: i64, content: Vec) -> Self { Message { - id: None, + id: Message::msg_id(), role, created, content, metadata: MessageMetadata::default(), } } + pub fn debug(&self) -> String { format!("{:?}", self) } @@ -478,7 +482,7 @@ impl Message { /// Create a new user message with the current timestamp pub fn user() -> Self { Message { - id: None, + id: Message::msg_id(), role: Role::User, created: Utc::now().timestamp(), content: Vec::new(), @@ -489,7 +493,7 @@ impl Message { /// Create a new assistant message with the current timestamp pub fn assistant() -> Self { Message { - id: None, + id: Message::msg_id(), role: Role::Assistant, created: Utc::now().timestamp(), content: Vec::new(), @@ -498,7 +502,7 @@ impl Message { } pub fn with_id>(mut self, id: S) -> Self { - self.id = Some(id.into()); + self.id = id.into(); self } diff --git a/crates/goose/src/conversation/mod.rs b/crates/goose/src/conversation/mod.rs index b0ae237faa..35b59e2f34 100644 --- a/crates/goose/src/conversation/mod.rs +++ b/crates/goose/src/conversation/mod.rs @@ -42,11 +42,7 @@ impl Conversation { } pub fn push(&mut self, message: Message) { - if let Some(last) = self - .0 - .last_mut() - .filter(|m| m.id.is_some() && m.id == message.id) - { + if let Some(last) = self.0.last_mut().filter(|m| m.id == message.id) { match (last.content.last_mut(), message.content.last()) { (Some(MessageContent::Text(ref mut last)), Some(MessageContent::Text(new))) if message.content.len() == 1 => diff --git a/crates/goose/src/providers/formats/anthropic.rs b/crates/goose/src/providers/formats/anthropic.rs index 6142d2188c..32751a8bf6 100644 --- a/crates/goose/src/providers/formats/anthropic.rs +++ b/crates/goose/src/providers/formats/anthropic.rs @@ -449,17 +449,16 @@ pub fn create_request( Ok(payload) } -/// Process streaming response from Anthropic's API pub fn response_to_streaming_message( mut stream: S, ) -> impl futures::Stream< - Item = anyhow::Result<( + Item = Result<( Option, Option, )>, > + 'static where - S: futures::Stream> + Unpin + Send + 'static, + S: futures::Stream> + Unpin + Send + 'static, { use async_stream::try_stream; use futures::StreamExt; @@ -483,19 +482,16 @@ where while let Some(line_result) = stream.next().await { let line = line_result?; - // Skip empty lines and non-data lines if line.trim().is_empty() || !line.starts_with("data: ") { continue; } let data_part = line.strip_prefix("data: ").unwrap_or(&line); - // Handle end of stream if data_part.trim() == "[DONE]" { break; } - // Parse the JSON event let event: StreamingEvent = match serde_json::from_str(data_part) { Ok(event) => event, Err(e) => { @@ -555,7 +551,9 @@ where chrono::Utc::now().timestamp(), vec![MessageContent::text(text)], ); - message.id = message_id.clone(); + if let Some(msg_id) = message_id.clone() { + message = message.with_id(msg_id); + } yield (Some(message), None); } } else if delta.get("type") == Some(&json!("input_json_delta")) { @@ -593,7 +591,9 @@ where chrono::Utc::now().timestamp(), vec![MessageContent::tool_request(tool_id, Err(error))], ); - message.id = message_id.clone(); + if let Some(msg_id) = message_id.clone() { + message = message.with_id(msg_id); + } yield (Some(message), None); continue; } @@ -607,7 +607,9 @@ where chrono::Utc::now().timestamp(), vec![MessageContent::tool_request(tool_id, Ok(tool_call))], ); - message.id = message_id.clone(); + if let Some(msg_id) = message_id.clone() { + message = message.with_id(msg_id); + } yield (Some(message), None); } } @@ -678,9 +680,8 @@ where } } - // Yield final usage information if available - if let Some(usage) = final_usage { - yield (None, Some(usage)); + if let Some(usage) = &final_usage { + yield (None, Some(usage.clone())); } else { tracing::debug!("🔍 Anthropic no final usage to yield"); } diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 0002d4199a..7dd799f14c 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -18,7 +18,7 @@ use tokio::sync::OnceCell; use tracing::{info, warn}; use utoipa::ToSchema; -const CURRENT_SCHEMA_VERSION: i32 = 3; +const CURRENT_SCHEMA_VERSION: i32 = 4; static SESSION_STORAGE: OnceCell> = OnceCell::const_new(); @@ -620,6 +620,18 @@ impl SessionStorage { .execute(&self.pool) .await?; } + 4 => { + sqlx::query( + r#" + ALTER TABLE messages ADD COLUMN msg_id TEXT + "#, + ) + .execute(&self.pool) + .await?; + sqlx::query("CREATE INDEX idx_messages_msg_id ON messages(msg_id)") + .execute(&self.pool) + .await?; + } _ => { anyhow::bail!("Unknown migration version: {}", version); } @@ -784,15 +796,15 @@ impl SessionStorage { } async fn get_conversation(&self, session_id: &str) -> Result { - let rows = sqlx::query_as::<_, (String, String, i64, Option)>( - "SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? ORDER BY timestamp", + let rows = sqlx::query_as::<_, (String, String, i64, Option, Option)>( + "SELECT role, content_json, created_timestamp, metadata_json, msg_id FROM messages WHERE session_id = ? ORDER BY timestamp", ) .bind(session_id) .fetch_all(&self.pool) .await?; let mut messages = Vec::new(); - for (role_str, content_json, created_timestamp, metadata_json) in rows { + for (role_str, content_json, created_timestamp, metadata_json, msg_id) in rows { let role = match role_str.as_str() { "user" => Role::User, "assistant" => Role::Assistant, @@ -806,6 +818,9 @@ impl SessionStorage { let mut message = Message::new(role, created_timestamp, content); message.metadata = metadata; + if let Some(msg_id) = msg_id { + message = message.with_id(msg_id); + } messages.push(message); } @@ -817,8 +832,8 @@ impl SessionStorage { sqlx::query( r#" - INSERT INTO messages (session_id, role, content_json, created_timestamp, metadata_json) - VALUES (?, ?, ?, ?, ?) + INSERT INTO messages (session_id, role, content_json, created_timestamp, metadata_json, msg_id) + VALUES (?, ?, ?, ?, ?, ?) "#, ) .bind(session_id) @@ -826,6 +841,7 @@ impl SessionStorage { .bind(serde_json::to_string(&message.content)?) .bind(message.created) .bind(metadata_json) + .bind(&message.id) .execute(&self.pool) .await?; @@ -999,7 +1015,7 @@ mod tests { .add_message( &session.id, &Message { - id: None, + id: Message::msg_id(), role: Role::User, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text("hello world")], @@ -1013,7 +1029,7 @@ mod tests { .add_message( &session.id, &Message { - id: None, + id: Message::msg_id(), role: Role::Assistant, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text("sup world?")], @@ -1102,7 +1118,7 @@ mod tests { .add_message( &original.id, &Message { - id: None, + id: Message::msg_id(), role: Role::User, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text(USER_MESSAGE)], @@ -1116,7 +1132,7 @@ mod tests { .add_message( &original.id, &Message { - id: None, + id: Message::msg_id(), role: Role::Assistant, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text(ASSISTANT_MESSAGE)], diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index 2f52cbbf2c..956f743c23 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -3128,8 +3128,8 @@ }, "Message": { "type": "object", - "description": "A message to or from an LLM", "required": [ + "id", "role", "created", "content", @@ -3147,8 +3147,7 @@ "format": "int64" }, "id": { - "type": "string", - "nullable": true + "type": "string" }, "metadata": { "$ref": "#/components/schemas/MessageMetadata" diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index c82e302f46..64575e2f06 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -374,13 +374,10 @@ export type LoadedProvider = { is_editable: boolean; }; -/** - * A message to or from an LLM - */ export type Message = { content: Array; created: number; - id?: string | null; + id: string; metadata: MessageMetadata; role: Role; };