From b8f82bc9d680fa9bb411d3aaa34be035b89efd17 Mon Sep 17 00:00:00 2001 From: Vincent Huang Date: Tue, 11 Nov 2025 22:08:06 -0800 Subject: [PATCH 1/2] Added transaction commits to multi sql functions in session_manager --- crates/goose/src/session/session_manager.rs | 122 +++++++++++++++++++- 1 file changed, 117 insertions(+), 5 deletions(-) diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 8dfc81e232cd..db4d6b97bd9c 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -970,6 +970,8 @@ impl SessionStorage { } async fn add_message(&self, session_id: &str, message: &Message) -> Result<()> { + let mut tx = self.pool.begin().await?; + let metadata_json = serde_json::to_string(&message.metadata)?; sqlx::query( @@ -983,14 +985,15 @@ impl SessionStorage { .bind(serde_json::to_string(&message.content)?) .bind(message.created) .bind(metadata_json) - .execute(&self.pool) + .execute(&mut *tx) .await?; sqlx::query("UPDATE sessions SET updated_at = datetime('now') WHERE id = ?") .bind(session_id) - .execute(&self.pool) + .execute(&mut *tx) .await?; + tx.commit().await?; Ok(()) } @@ -1049,10 +1052,12 @@ impl SessionStorage { } async fn delete_session(&self, session_id: &str) -> Result<()> { + let mut tx = self.pool.begin().await?; + let exists = sqlx::query_scalar::<_, bool>("SELECT EXISTS(SELECT 1 FROM sessions WHERE id = ?)") .bind(session_id) - .fetch_one(&self.pool) + .fetch_one(&mut *tx) .await?; if !exists { @@ -1061,14 +1066,15 @@ impl SessionStorage { sqlx::query("DELETE FROM messages WHERE session_id = ?") .bind(session_id) - .execute(&self.pool) + .execute(&mut *tx) .await?; sqlx::query("DELETE FROM sessions WHERE id = ?") .bind(session_id) - .execute(&self.pool) + .execute(&mut *tx) .await?; + tx.commit().await?; Ok(()) } @@ -1358,4 +1364,110 @@ mod tests { assert!(imported.user_set_name); assert_eq!(imported.working_dir, PathBuf::from("/tmp/test")); } + + #[tokio::test] + async fn test_delete_session() { + let temp_dir = TempDir::new().unwrap(); + let db_path = temp_dir.path().join("test_delete.db"); + let storage = Arc::new(SessionStorage::create(&db_path).await.unwrap()); + + let session = storage + .create_session( + PathBuf::from("/tmp/test"), + "Test session".to_string(), + SessionType::User, + ) + .await + .unwrap(); + + storage + .add_message( + &session.id, + &Message { + id: None, + role: Role::User, + created: chrono::Utc::now().timestamp_millis(), + content: vec![MessageContent::text("test message")], + metadata: Default::default(), + }, + ) + .await + .unwrap(); + + let retrieved = storage.get_session(&session.id, true).await.unwrap(); + assert_eq!(retrieved.message_count, 1); + + storage.delete_session(&session.id).await.unwrap(); + + let result = storage.get_session(&session.id, false).await; + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Session not found")); + } + + #[tokio::test] + async fn test_add_message() { + const MESSAGE_COUNT: usize = 5; + + let temp_dir = TempDir::new().unwrap(); + let db_path = temp_dir.path().join("test_add_message.db"); + let storage = Arc::new(SessionStorage::create(&db_path).await.unwrap()); + + let session = storage + .create_session( + PathBuf::from("/tmp/test"), + "Test session".to_string(), + SessionType::User, + ) + .await + .unwrap(); + + // Add multiple messages with alternating roles + for i in 0..MESSAGE_COUNT { + let message = Message::new( + if i % 2 == 0 { + Role::User + } else { + Role::Assistant + }, + chrono::Utc::now().timestamp_millis(), + vec![MessageContent::text(format!("message {}", i))], + ); + + storage.add_message(&session.id, &message).await.unwrap(); + } + + // Add a message with custom metadata + let mut message_with_metadata = Message::new( + Role::User, + chrono::Utc::now().timestamp_millis(), + vec![MessageContent::text("message with custom metadata")], + ); + message_with_metadata.metadata.user_visible = false; + message_with_metadata.metadata.agent_visible = true; + storage + .add_message(&session.id, &message_with_metadata) + .await + .unwrap(); + + // Verify messages were added + let updated_session = storage.get_session(&session.id, true).await.unwrap(); + assert_eq!(updated_session.message_count, MESSAGE_COUNT + 1); + + // Verify messages and roles + let conversation = updated_session.conversation.unwrap(); + assert_eq!(conversation.messages().len(), MESSAGE_COUNT + 1); + assert_eq!(conversation.messages()[0].role, Role::User); + assert_eq!(conversation.messages()[1].role, Role::Assistant); + assert_eq!(conversation.messages()[2].role, Role::User); + assert_eq!(conversation.messages()[3].role, Role::Assistant); + assert_eq!(conversation.messages()[4].role, Role::User); + + // Verify custom metadata was preserved on the last message + let last_message = conversation.messages().last().unwrap(); + assert!(!last_message.metadata.user_visible); + assert!(last_message.metadata.agent_visible); + } } From 3622c979125f3ca877d3bbc8876ffc9b25e6c80c Mon Sep 17 00:00:00 2001 From: Vincent Huang Date: Wed, 12 Nov 2025 09:23:56 -0800 Subject: [PATCH 2/2] Removed tests Signed-off-by: Vincent Huang --- crates/goose/src/session/session_manager.rs | 106 -------------------- 1 file changed, 106 deletions(-) diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index db4d6b97bd9c..564911acbc1e 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -1364,110 +1364,4 @@ mod tests { assert!(imported.user_set_name); assert_eq!(imported.working_dir, PathBuf::from("/tmp/test")); } - - #[tokio::test] - async fn test_delete_session() { - let temp_dir = TempDir::new().unwrap(); - let db_path = temp_dir.path().join("test_delete.db"); - let storage = Arc::new(SessionStorage::create(&db_path).await.unwrap()); - - let session = storage - .create_session( - PathBuf::from("/tmp/test"), - "Test session".to_string(), - SessionType::User, - ) - .await - .unwrap(); - - storage - .add_message( - &session.id, - &Message { - id: None, - role: Role::User, - created: chrono::Utc::now().timestamp_millis(), - content: vec![MessageContent::text("test message")], - metadata: Default::default(), - }, - ) - .await - .unwrap(); - - let retrieved = storage.get_session(&session.id, true).await.unwrap(); - assert_eq!(retrieved.message_count, 1); - - storage.delete_session(&session.id).await.unwrap(); - - let result = storage.get_session(&session.id, false).await; - assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("Session not found")); - } - - #[tokio::test] - async fn test_add_message() { - const MESSAGE_COUNT: usize = 5; - - let temp_dir = TempDir::new().unwrap(); - let db_path = temp_dir.path().join("test_add_message.db"); - let storage = Arc::new(SessionStorage::create(&db_path).await.unwrap()); - - let session = storage - .create_session( - PathBuf::from("/tmp/test"), - "Test session".to_string(), - SessionType::User, - ) - .await - .unwrap(); - - // Add multiple messages with alternating roles - for i in 0..MESSAGE_COUNT { - let message = Message::new( - if i % 2 == 0 { - Role::User - } else { - Role::Assistant - }, - chrono::Utc::now().timestamp_millis(), - vec![MessageContent::text(format!("message {}", i))], - ); - - storage.add_message(&session.id, &message).await.unwrap(); - } - - // Add a message with custom metadata - let mut message_with_metadata = Message::new( - Role::User, - chrono::Utc::now().timestamp_millis(), - vec![MessageContent::text("message with custom metadata")], - ); - message_with_metadata.metadata.user_visible = false; - message_with_metadata.metadata.agent_visible = true; - storage - .add_message(&session.id, &message_with_metadata) - .await - .unwrap(); - - // Verify messages were added - let updated_session = storage.get_session(&session.id, true).await.unwrap(); - assert_eq!(updated_session.message_count, MESSAGE_COUNT + 1); - - // Verify messages and roles - let conversation = updated_session.conversation.unwrap(); - assert_eq!(conversation.messages().len(), MESSAGE_COUNT + 1); - assert_eq!(conversation.messages()[0].role, Role::User); - assert_eq!(conversation.messages()[1].role, Role::Assistant); - assert_eq!(conversation.messages()[2].role, Role::User); - assert_eq!(conversation.messages()[3].role, Role::Assistant); - assert_eq!(conversation.messages()[4].role, Role::User); - - // Verify custom metadata was preserved on the last message - let last_message = conversation.messages().last().unwrap(); - assert!(!last_message.metadata.user_visible); - assert!(last_message.metadata.agent_visible); - } }