From f5ad3384b82733e9d239383b1acc8bda599a28f7 Mon Sep 17 00:00:00 2001 From: Rabi Mishra Date: Thu, 26 Mar 2026 22:22:42 +0530 Subject: [PATCH] fix(acp): separate acp sessions from user sessions (#7857) Signed-off-by: Rabi Mishra Co-authored-by: Douwe Osinga --- crates/goose-acp/src/server.rs | 6 +- crates/goose-server/src/routes/session.rs | 13 +- .../agents/platform_extensions/chatrecall.rs | 9 + .../goose/src/session/chat_history_search.rs | 18 ++ crates/goose/src/session/session_manager.rs | 252 +++++++++++++++--- ui/desktop/openapi.json | 3 +- ui/desktop/src/api/types.gen.ts | 2 +- 7 files changed, 251 insertions(+), 52 deletions(-) diff --git a/crates/goose-acp/src/server.rs b/crates/goose-acp/src/server.rs index 730c1fc0..281e678f 100644 --- a/crates/goose-acp/src/server.rs +++ b/crates/goose-acp/src/server.rs @@ -848,7 +848,7 @@ impl GooseAcpAgent { .create_session( args.cwd.clone(), "ACP Session".to_string(), - SessionType::User, + SessionType::Acp, self.goose_mode, ) .await @@ -1246,7 +1246,7 @@ impl GooseAcpAgent { async fn on_list_sessions(&self) -> Result { let sessions = self .session_manager - .list_sessions() + .list_sessions_by_types(&[SessionType::Acp]) .await .map_err(|e| sacp::Error::internal_error().data(e.to_string()))?; let session_infos: Vec = sessions @@ -1422,7 +1422,7 @@ impl GooseAcpAgent { ) -> Result { let session = self .session_manager - .import_session(&req.data) + .import_session(&req.data, Some(SessionType::Acp)) .await .map_err(|e| sacp::Error::internal_error().data(e.to_string()))?; let session_json = serde_json::to_value(&session) diff --git a/crates/goose-server/src/routes/session.rs b/crates/goose-server/src/routes/session.rs index d66ed390..d83194e4 100644 --- a/crates/goose-server/src/routes/session.rs +++ b/crates/goose-server/src/routes/session.rs @@ -11,7 +11,7 @@ use axum::{ }; use goose::agents::ExtensionConfig; use goose::recipe::Recipe; -use goose::session::session_manager::SessionInsights; +use goose::session::session_manager::{SessionInsights, SessionType}; use goose::session::{EnabledExtensionsState, Session}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -357,7 +357,7 @@ async fn import_session( ) -> Result, StatusCode> { let session = state .session_manager() - .import_session(&request.json) + .import_session(&request.json, Some(SessionType::User)) .await .map_err(|_| StatusCode::BAD_REQUEST)?; @@ -583,7 +583,14 @@ async fn search_sessions( let search_results = state .session_manager() - .search_chat_history(query, Some(limit), after_date, before_date, None) + .search_chat_history( + query, + Some(limit), + after_date, + before_date, + None, + vec![SessionType::User, SessionType::Scheduled], + ) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; diff --git a/crates/goose/src/agents/platform_extensions/chatrecall.rs b/crates/goose/src/agents/platform_extensions/chatrecall.rs index aa135417..13a6a75e 100644 --- a/crates/goose/src/agents/platform_extensions/chatrecall.rs +++ b/crates/goose/src/agents/platform_extensions/chatrecall.rs @@ -1,6 +1,7 @@ use crate::agents::extension::PlatformExtensionContext; use crate::agents::mcp_client::{Error, McpClientTrait}; use crate::agents::tool_execution::ToolCallContext; +use crate::session::session_manager::SessionType; use anyhow::Result; use async_trait::async_trait; use indoc::indoc; @@ -58,6 +59,13 @@ impl ChatRecallClient { Ok(Self { info, context }) } + fn search_session_types(&self) -> Vec { + match self.context.session.as_ref().map(|s| s.session_type) { + Some(SessionType::Acp) => vec![SessionType::Acp], + _ => vec![SessionType::User, SessionType::Scheduled], + } + } + #[allow(clippy::too_many_lines)] async fn handle_chatrecall( &self, @@ -177,6 +185,7 @@ impl ChatRecallClient { after_date, before_date, exclude_session_id, + self.search_session_types(), ) .await { diff --git a/crates/goose/src/session/chat_history_search.rs b/crates/goose/src/session/chat_history_search.rs index f3436701..c06abb84 100644 --- a/crates/goose/src/session/chat_history_search.rs +++ b/crates/goose/src/session/chat_history_search.rs @@ -1,4 +1,5 @@ use crate::conversation::message::MessageContent; +use crate::session::session_manager::SessionType; use anyhow::Result; use chrono::{DateTime, Utc}; use serde::Serialize; @@ -52,6 +53,7 @@ pub struct ChatHistorySearch<'a> { after_date: Option>, before_date: Option>, exclude_session_id: Option, + session_types: Vec, } impl<'a> ChatHistorySearch<'a> { @@ -62,6 +64,7 @@ impl<'a> ChatHistorySearch<'a> { after_date: Option>, before_date: Option>, exclude_session_id: Option, + session_types: Vec, ) -> Self { Self { pool, @@ -70,6 +73,7 @@ impl<'a> ChatHistorySearch<'a> { after_date, before_date, exclude_session_id, + session_types, } } @@ -102,6 +106,10 @@ impl<'a> ChatHistorySearch<'a> { query_builder = query_builder.bind(exclude_id); } + for t in &self.session_types { + query_builder = query_builder.bind(t.to_string()); + } + if let Some(after) = self.after_date { query_builder = query_builder.bind(after); } @@ -159,6 +167,16 @@ impl<'a> ChatHistorySearch<'a> { sql.push_str(" AND s.id != ?"); } + if !self.session_types.is_empty() { + let placeholders: String = self + .session_types + .iter() + .map(|_| "?") + .collect::>() + .join(", "); + sql.push_str(&format!(" AND s.session_type IN ({})", placeholders)); + } + if self.after_date.is_some() { sql.push_str(" AND m.timestamp >= ?"); } diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index da447d7c..5f95bc97 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -19,12 +19,25 @@ use std::sync::{Arc, LazyLock}; use tracing::{info, warn}; use utoipa::ToSchema; -pub const CURRENT_SCHEMA_VERSION: i32 = 8; +pub const CURRENT_SCHEMA_VERSION: i32 = 9; pub const SESSIONS_FOLDER: &str = "sessions"; pub const DB_NAME: &str = "sessions.db"; -#[derive(Debug, Clone, Copy, Serialize, Deserialize, ToSchema, PartialEq, Eq, Default)] +#[derive( + Debug, + Clone, + Copy, + Serialize, + Deserialize, + ToSchema, + PartialEq, + Eq, + Default, + strum::Display, + strum::EnumString, +)] #[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] pub enum SessionType { #[default] User, @@ -33,35 +46,7 @@ pub enum SessionType { Hidden, Terminal, Gateway, -} - -impl std::fmt::Display for SessionType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - SessionType::User => write!(f, "user"), - SessionType::SubAgent => write!(f, "sub_agent"), - SessionType::Hidden => write!(f, "hidden"), - SessionType::Scheduled => write!(f, "scheduled"), - SessionType::Terminal => write!(f, "terminal"), - SessionType::Gateway => write!(f, "gateway"), - } - } -} - -impl std::str::FromStr for SessionType { - type Err = anyhow::Error; - - fn from_str(s: &str) -> Result { - match s { - "user" => Ok(SessionType::User), - "sub_agent" => Ok(SessionType::SubAgent), - "hidden" => Ok(SessionType::Hidden), - "scheduled" => Ok(SessionType::Scheduled), - "terminal" => Ok(SessionType::Terminal), - "gateway" => Ok(SessionType::Gateway), - _ => Err(anyhow::anyhow!("Invalid session type: {}", s)), - } - } + Acp, } static SESSION_STORAGE: LazyLock> = @@ -323,15 +308,23 @@ impl SessionManager { } pub async fn get_insights(&self) -> Result { - self.storage.get_insights().await + self.storage + .get_insights(&[SessionType::User, SessionType::Scheduled]) + .await } pub async fn export_session(&self, id: &str) -> Result { self.storage.export_session(id).await } - pub async fn import_session(&self, json: &str) -> Result { - self.storage.import_session(self, json).await + pub async fn import_session( + &self, + json: &str, + session_type_override: Option, + ) -> Result { + self.storage + .import_session(self, json, session_type_override) + .await } pub async fn copy_session(&self, session_id: &str, new_name: String) -> Result { @@ -376,9 +369,17 @@ impl SessionManager { after_date: Option>, before_date: Option>, exclude_session_id: Option, + session_types: Vec, ) -> Result { self.storage - .search_chat_history(query, limit, after_date, before_date, exclude_session_id) + .search_chat_history( + query, + limit, + after_date, + before_date, + exclude_session_id, + session_types, + ) .await } @@ -919,6 +920,19 @@ impl SessionStorage { .execute(&mut **tx) .await?; } + 9 => { + sqlx::query( + r#" + UPDATE sessions + SET session_type = 'acp' + WHERE session_type = 'user' + AND name = 'ACP Session' + AND user_set_name = FALSE + "#, + ) + .execute(&mut **tx) + .await?; + } _ => { anyhow::bail!("Unknown migration version: {}", version); } @@ -1316,17 +1330,32 @@ impl SessionStorage { Ok(()) } - async fn get_insights(&self) -> Result { - let pool = self.pool().await?; - let row = sqlx::query_as::<_, (i64, Option)>( + async fn get_insights(&self, types: &[SessionType]) -> Result { + if types.is_empty() { + return Ok(SessionInsights { + total_sessions: 0, + total_tokens: 0, + }); + } + + let placeholders: String = types.iter().map(|_| "?").collect::>().join(", "); + let query = format!( r#" SELECT COUNT(*) as total_sessions, COALESCE(SUM(COALESCE(accumulated_total_tokens, total_tokens, 0)), 0) as total_tokens FROM sessions + WHERE session_type IN ({}) "#, - ) - .fetch_one(pool) - .await?; + placeholders + ); + + let pool = self.pool().await?; + let mut q = sqlx::query_as::<_, (i64, Option)>(&query); + for t in types { + q = q.bind(t.to_string()); + } + + let row = q.fetch_one(pool).await?; Ok(SessionInsights { total_sessions: row.0 as usize, @@ -1343,6 +1372,7 @@ impl SessionStorage { &self, session_manager: &SessionManager, json: &str, + session_type_override: Option, ) -> Result { let import: Session = serde_json::from_str(json)?; @@ -1350,7 +1380,7 @@ impl SessionStorage { .create_session( import.working_dir.clone(), import.name.clone(), - import.session_type, + session_type_override.unwrap_or(import.session_type), import.goose_mode, ) .await?; @@ -1445,6 +1475,7 @@ impl SessionStorage { after_date: Option>, before_date: Option>, exclude_session_id: Option, + session_types: Vec, ) -> Result { use crate::session::chat_history_search::ChatHistorySearch; @@ -1456,6 +1487,7 @@ impl SessionStorage { after_date, before_date, exclude_session_id, + session_types, ) .execute() .await @@ -1753,7 +1785,7 @@ mod tests { .unwrap(); let exported = sm.export_session(&original.id).await.unwrap(); - let imported = sm.import_session(&exported).await.unwrap(); + let imported = sm.import_session(&exported, None).await.unwrap(); assert_ne!(imported.id, original.id); assert_eq!(imported.name, DESCRIPTION); @@ -1770,6 +1802,69 @@ mod tests { assert_eq!(conversation.messages()[1].role, Role::Assistant); } + #[tokio::test] + async fn test_list_sessions_filters_by_type() { + let temp_dir = TempDir::new().unwrap(); + let sm = SessionManager::new(temp_dir.path().to_path_buf()); + + let user_session = sm + .create_session( + PathBuf::from("/tmp/test"), + "User session".to_string(), + SessionType::User, + GooseMode::default(), + ) + .await + .unwrap(); + + sm.add_message( + &user_session.id, + &Message { + id: None, + role: Role::User, + created: chrono::Utc::now().timestamp_millis(), + content: vec![MessageContent::text("hello world")], + metadata: Default::default(), + }, + ) + .await + .unwrap(); + + let acp_session = sm + .create_session( + PathBuf::from("/tmp/test"), + "ACP session".to_string(), + SessionType::Acp, + GooseMode::default(), + ) + .await + .unwrap(); + + sm.add_message( + &acp_session.id, + &Message { + id: None, + role: Role::User, + created: chrono::Utc::now().timestamp_millis(), + content: vec![MessageContent::text("hello acp")], + metadata: Default::default(), + }, + ) + .await + .unwrap(); + + let default_sessions = sm.list_sessions().await.unwrap(); + assert_eq!(default_sessions.len(), 1); + assert_eq!(default_sessions[0].name, "User session"); + + let acp_sessions = sm + .list_sessions_by_types(&[SessionType::Acp]) + .await + .unwrap(); + assert_eq!(acp_sessions.len(), 1); + assert_eq!(acp_sessions[0].name, "ACP session"); + } + #[tokio::test] async fn test_import_session_with_description_field() { const OLD_FORMAT_JSON: &str = r#"{ @@ -1786,7 +1881,7 @@ mod tests { let temp_dir = TempDir::new().unwrap(); let sm = SessionManager::new(temp_dir.path().to_path_buf()); - let imported = sm.import_session(OLD_FORMAT_JSON).await.unwrap(); + let imported = sm.import_session(OLD_FORMAT_JSON, None).await.unwrap(); assert_eq!(imported.name, "Old format session"); assert!(imported.user_set_name); @@ -1865,4 +1960,73 @@ mod tests { let reloaded = sm.get_session(&session.id, false).await.unwrap(); assert_eq!(reloaded.goose_mode, GooseMode::default()); } + + #[tokio::test] + async fn test_acp_session_migration() { + let temp_dir = TempDir::new().unwrap(); + let db_path = temp_dir.path().join(SESSIONS_FOLDER).join(DB_NAME); + + if let Some(parent) = db_path.parent() { + std::fs::create_dir_all(parent).unwrap(); + } + + let pool = SqlitePoolOptions::new() + .connect_with( + SqliteConnectOptions::new() + .filename(&db_path) + .create_if_missing(true), + ) + .await + .unwrap(); + + SessionStorage::create_schema(&pool).await.unwrap(); + + // Demote the schema back to v8 to simulate a database + // that has never seen migration 9. + sqlx::query("UPDATE schema_version SET version = 8") + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO sessions (id, name, user_set_name, session_type, working_dir, extension_data, goose_mode) + VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .bind("user_id") + .bind("User Session") + .bind(false) + .bind("user") + .bind("/tmp") + .bind("{}") + .bind("auto") + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO sessions (id, name, user_set_name, session_type, working_dir, extension_data, goose_mode) + VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .bind("acp_id") + .bind("ACP Session") + .bind(false) + .bind("user") + .bind("/tmp") + .bind("{}") + .bind("auto") + .execute(&pool) + .await + .unwrap(); + + pool.close().await; + + let sm = SessionManager::new(temp_dir.path().to_path_buf()); + sm.storage().pool().await.unwrap(); // Triggers migration + + let user_session = sm.storage().get_session("user_id", false).await.unwrap(); + assert_eq!(user_session.session_type, SessionType::User); + + let acp_session = sm.storage().get_session("acp_id", false).await.unwrap(); + assert_eq!(acp_session.session_type, SessionType::Acp); + } } diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index af47a77d..f16e6597 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -7904,7 +7904,8 @@ "sub_agent", "hidden", "terminal", - "gateway" + "gateway", + "acp" ] }, "SessionsQuery": { diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 10d7b59b..eaffc44e 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -1276,7 +1276,7 @@ export type SessionReplyResponse = { request_id: string; }; -export type SessionType = 'user' | 'scheduled' | 'sub_agent' | 'hidden' | 'terminal' | 'gateway'; +export type SessionType = 'user' | 'scheduled' | 'sub_agent' | 'hidden' | 'terminal' | 'gateway' | 'acp'; export type SessionsQuery = { limit: number;