diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 7f3735743..b30e9b52a 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -133,6 +133,13 @@ pub type AcpProviderFactory = Arc< + Sync, >; +const ACP_VISIBLE_SESSION_TYPES: [SessionType; 3] = + [SessionType::User, SessionType::Scheduled, SessionType::Acp]; + +fn is_acp_visible_session_type(session_type: &SessionType) -> bool { + ACP_VISIBLE_SESSION_TYPES.contains(session_type) +} + /// Convenience conversions from any `Display` error into an `agent_client_protocol::Error`. /// /// Replaces the repetitive `.internal_err()` diff --git a/crates/goose/src/acp/server/list_sessions.rs b/crates/goose/src/acp/server/list_sessions.rs index 8ba208553..7a620f900 100644 --- a/crates/goose/src/acp/server/list_sessions.rs +++ b/crates/goose/src/acp/server/list_sessions.rs @@ -1,4 +1,7 @@ -use super::{build_session_info, meta_string, GooseAcpAgent, ResultExt}; +use super::{ + build_session_info, is_acp_visible_session_type, meta_string, GooseAcpAgent, ResultExt, + ACP_VISIBLE_SESSION_TYPES, +}; use crate::session::session_manager::{ SessionListCursor, SessionListFilters, SessionListPageQuery, SessionType, }; @@ -10,8 +13,6 @@ use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; const SESSION_LIST_PAGE_SIZE: usize = 50; -const ACP_SESSION_LIST_TYPES: [SessionType; 3] = - [SessionType::User, SessionType::Scheduled, SessionType::Acp]; #[derive(Debug, Serialize, Deserialize)] struct SessionListCursorToken { @@ -48,10 +49,10 @@ fn session_types_from_meta( meta: Option<&Meta>, ) -> Result, agent_client_protocol::Error> { let Some(value) = meta.and_then(|meta| meta.get("types")) else { - return Ok(ACP_SESSION_LIST_TYPES.to_vec()); + return Ok(ACP_VISIBLE_SESSION_TYPES.to_vec()); }; if value.is_null() { - return Ok(ACP_SESSION_LIST_TYPES.to_vec()); + return Ok(ACP_VISIBLE_SESSION_TYPES.to_vec()); } let session_types = @@ -60,11 +61,11 @@ fn session_types_from_meta( .data("types must be an array of session type strings") })?; if session_types.is_empty() { - Ok(ACP_SESSION_LIST_TYPES.to_vec()) + Ok(ACP_VISIBLE_SESSION_TYPES.to_vec()) } else { if session_types .iter() - .any(|session_type| !ACP_SESSION_LIST_TYPES.contains(session_type)) + .any(|session_type| !is_acp_visible_session_type(session_type)) { return Err(agent_client_protocol::Error::invalid_params() .data("types may only include user, scheduled, or acp")); diff --git a/crates/goose/src/acp/server/manage_sessions.rs b/crates/goose/src/acp/server/manage_sessions.rs index 24d89d370..cb1e712ca 100644 --- a/crates/goose/src/acp/server/manage_sessions.rs +++ b/crates/goose/src/acp/server/manage_sessions.rs @@ -229,12 +229,23 @@ impl GooseAcpAgent { &self, req: UpdateSessionProjectRequest, ) -> Result { - self.session_manager - .update(&req.session_id) - .project_id(req.project_id) - .apply() + let session_id = &req.session_id; + let session_not_found = || { + agent_client_protocol::Error::resource_not_found(Some(session_id.to_string())) + .data(format!("Session not found: {session_id}")) + }; + let updated = self + .session_manager + .update_project_for_session_types( + session_id, + req.project_id, + &ACP_VISIBLE_SESSION_TYPES, + ) .await .internal_err()?; + if !updated { + return Err(session_not_found()); + } Ok(EmptyResponse {}) } diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 5adf1f7d1..f7bc9f98c 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -436,6 +436,17 @@ impl SessionManager { SessionUpdateBuilder::new(self, id.to_string()) } + pub(crate) async fn update_project_for_session_types( + &self, + id: &str, + project_id: Option, + session_types: &[SessionType], + ) -> Result { + self.storage + .update_project_for_session_types(id, project_id, session_types) + .await + } + async fn apply_update_inner(&self, builder: SessionUpdateBuilder<'_>) -> Result<()> { self.storage.apply_update(builder).await } @@ -1798,6 +1809,34 @@ impl SessionStorage { Ok(()) } + async fn update_project_for_session_types( + &self, + id: &str, + project_id: Option, + session_types: &[SessionType], + ) -> Result { + if session_types.is_empty() { + return Ok(false); + } + + let placeholders = session_types + .iter() + .map(|_| "?") + .collect::>() + .join(", "); + let query = format!( + "UPDATE sessions SET project_id = ?, updated_at = datetime('now') \ + WHERE id = ? AND session_type IN ({placeholders})" + ); + let pool = self.pool().await?; + let mut query = sqlx::query(AssertSqlSafe(query)).bind(project_id).bind(id); + for session_type in session_types { + query = query.bind(session_type.to_string()); + } + + Ok(query.execute(pool).await?.rows_affected() == 1) + } + async fn get_conversation(&self, session_id: &str) -> Result { let pool = self.pool().await?; let rows = sqlx::query_as::<_, (String, String, i64, Option, Option)>( diff --git a/crates/goose/tests/acp_server_test.rs b/crates/goose/tests/acp_server_test.rs index 7501216e4..085242fdc 100644 --- a/crates/goose/tests/acp_server_test.rs +++ b/crates/goose/tests/acp_server_test.rs @@ -31,7 +31,9 @@ use common_tests::{ }; use goose::config::GooseMode; use goose::conversation::message::{Message, MessageMetadata}; -use goose::custom_requests::{GetSessionInfoRequest, GetSessionInfoResponse}; +use goose::custom_requests::{ + GetSessionInfoRequest, GetSessionInfoResponse, UpdateSessionProjectRequest, +}; use goose::recipe::{Recipe, Settings}; use goose::recipe_deeplink; use goose::session::{SessionManager, SessionType}; @@ -514,6 +516,176 @@ fn test_get_session_info() { }); } +#[test] +fn test_update_session_project_rejects_hidden_session_types() { + run_test(async { + let data_root = tempfile::tempdir().unwrap(); + let working_dir = tempfile::tempdir().unwrap(); + let session_manager = SessionManager::new(data_root.path().to_path_buf()); + let mut hidden_sessions = Vec::new(); + + for session_type in [ + SessionType::Gateway, + SessionType::SubAgent, + SessionType::Hidden, + SessionType::Terminal, + ] { + hidden_sessions.push( + session_manager + .create_session( + working_dir.path().to_path_buf(), + format!("{session_type} session"), + session_type, + GooseMode::default(), + ) + .await + .unwrap(), + ); + } + + let conn = new_connection(data_root.path()).await; + for session in hidden_sessions { + let error = conn + .cx() + .send_request(UpdateSessionProjectRequest { + session_id: session.id.clone(), + project_id: Some("untrusted-project".to_string()), + }) + .block_task() + .await + .unwrap_err(); + assert_eq!(error.code, ErrorCode::ResourceNotFound); + + let stored = session_manager + .get_session(&session.id, false) + .await + .unwrap(); + assert_eq!(stored.project_id, None); + } + }); +} + +#[test] +fn test_update_session_project_rejects_unknown_persisted_session_type() { + run_test(async { + let data_root = tempfile::tempdir().unwrap(); + let working_dir = tempfile::tempdir().unwrap(); + let session_manager = SessionManager::new(data_root.path().to_path_buf()); + let session = session_manager + .create_session( + working_dir.path().to_path_buf(), + "Future session".to_string(), + SessionType::User, + GooseMode::default(), + ) + .await + .unwrap(); + let db_path = data_root + .path() + .join(goose::session::session_manager::SESSIONS_FOLDER) + .join(goose::session::session_manager::DB_NAME); + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .connect_with(sqlx::sqlite::SqliteConnectOptions::new().filename(db_path)) + .await + .unwrap(); + sqlx::query("UPDATE sessions SET session_type = 'future_type' WHERE id = ?") + .bind(&session.id) + .execute(&pool) + .await + .unwrap(); + + assert_eq!( + session_manager + .get_session(&session.id, false) + .await + .unwrap() + .session_type, + SessionType::User + ); + + let conn = new_connection(data_root.path()).await; + let error = conn + .cx() + .send_request(UpdateSessionProjectRequest { + session_id: session.id.clone(), + project_id: Some("untrusted-project".to_string()), + }) + .block_task() + .await + .unwrap_err(); + assert_eq!(error.code, ErrorCode::ResourceNotFound); + + let project_id = + sqlx::query_scalar::<_, Option>("SELECT project_id FROM sessions WHERE id = ?") + .bind(&session.id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(project_id, None); + }); +} + +#[test] +fn test_update_session_project_allows_visible_session_types() { + run_test(async { + let data_root = tempfile::tempdir().unwrap(); + let working_dir = tempfile::tempdir().unwrap(); + let session_manager = SessionManager::new(data_root.path().to_path_buf()); + let mut visible_sessions = Vec::new(); + + for session_type in [SessionType::User, SessionType::Scheduled, SessionType::Acp] { + visible_sessions.push( + session_manager + .create_session( + working_dir.path().to_path_buf(), + format!("{session_type} session"), + session_type, + GooseMode::default(), + ) + .await + .unwrap(), + ); + } + + let conn = new_connection(data_root.path()).await; + for session in visible_sessions { + conn.cx() + .send_request(UpdateSessionProjectRequest { + session_id: session.id.clone(), + project_id: Some("project".to_string()), + }) + .block_task() + .await + .unwrap(); + assert_eq!( + session_manager + .get_session(&session.id, false) + .await + .unwrap() + .project_id, + Some("project".to_string()) + ); + + conn.cx() + .send_request(UpdateSessionProjectRequest { + session_id: session.id.clone(), + project_id: None, + }) + .block_task() + .await + .unwrap(); + assert_eq!( + session_manager + .get_session(&session.id, false) + .await + .unwrap() + .project_id, + None + ); + } + }); +} + #[test] fn test_session_name_update_notification() { run_test(async { run_session_name_update_notification::().await });