diff --git a/crates/goose-cli/src/commands/session.rs b/crates/goose-cli/src/commands/session.rs index 7c45eaf93..31e8715dd 100644 --- a/crates/goose-cli/src/commands/session.rs +++ b/crates/goose-cli/src/commands/session.rs @@ -155,6 +155,18 @@ fn session_activity_at(session: &Session) -> chrono::DateTime { session.last_message_at.unwrap_or(session.updated_at) } +fn session_list_limit_if_safe_to_push_down( + ascending: bool, + working_dir: Option<&Path>, + limit: Option, +) -> Option { + if !ascending && working_dir.is_none() { + limit + } else { + None + } +} + pub async fn handle_session_list( format: String, ascending: bool, @@ -162,7 +174,13 @@ pub async fn handle_session_list( limit: Option, ) -> Result<()> { let session_manager = SessionManager::instance(); - let mut sessions = session_manager.list_sessions().await?; + let mut sessions = if let Some(limit) = + session_list_limit_if_safe_to_push_down(ascending, working_dir.as_deref(), limit) + { + session_manager.list_sessions_with_limit(limit).await? + } else { + session_manager.list_sessions().await? + }; if let Some(ref pat) = working_dir { let pat_lower = pat.to_string_lossy().to_lowercase(); @@ -223,6 +241,31 @@ pub async fn handle_session_list( Ok(()) } +#[cfg(test)] +mod session_list_tests { + use super::*; + + #[test] + fn pushes_down_only_safe_session_list_limits() { + assert_eq!( + session_list_limit_if_safe_to_push_down(false, None, Some(50)), + Some(50) + ); + assert_eq!( + session_list_limit_if_safe_to_push_down(true, None, Some(50)), + None + ); + assert_eq!( + session_list_limit_if_safe_to_push_down(false, Some(Path::new("/tmp")), Some(50)), + None + ); + assert_eq!( + session_list_limit_if_safe_to_push_down(false, None, None), + None + ); + } +} + pub async fn handle_session_export( session_id: String, output_path: Option, diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index f7bc9f98c..7e1b0b616 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -28,6 +28,7 @@ pub const CURRENT_SCHEMA_VERSION: i32 = 16; pub const SESSIONS_FOLDER: &str = "sessions"; pub const DB_NAME: &str = "sessions.db"; const MILLISECOND_TIMESTAMP_THRESHOLD: i64 = 10_000_000_000; +const SESSION_COUNT_BATCH_SIZE: usize = 900; #[derive( Debug, @@ -463,6 +464,19 @@ impl SessionManager { self.storage.list_sessions().await } + pub async fn list_sessions_with_limit(&self, limit: usize) -> Result> { + self.storage + .list_sessions_matching(SessionListQuery { + filters: SessionListFilters { + types: Some(&[SessionType::User, SessionType::Scheduled]), + ..Default::default() + }, + limit: Some(limit), + ..Default::default() + }) + .await + } + pub async fn list_sessions_by_types(&self, types: &[SessionType]) -> Result> { self.storage.list_sessions_by_types(Some(types)).await } @@ -1975,6 +1989,7 @@ impl SessionStorage { return Ok(Vec::new()); } + let has_limit = query.limit.is_some(); let keywords = keyword_terms(filters.keyword); let mut where_clauses = Vec::new(); let mut having_clauses = Vec::new(); @@ -2015,6 +2030,14 @@ impl SessionStorage { let order_by = "ORDER BY sort_timestamp DESC, s.id DESC"; let limit_clause = if query.limit.is_some() { "LIMIT ?" } else { "" }; + let message_count_sql = if has_limit { + "0".to_string() + } else { + format!( + "COUNT(m.id) FILTER (WHERE {})", + user_visible_message_sql("m.metadata_json") + ) + }; let sql = format!( r#" SELECT s.id, s.working_dir, s.name, s.description, s.user_set_name, s.session_type, s.created_at, s.updated_at, s.extension_data, @@ -2026,7 +2049,7 @@ impl SessionStorage { s.schedule_id, s.recipe_json, s.user_recipe_values_json, s.provider_name, s.model_config_json, s.goose_mode, s.archived_at, s.project_id, s.parent_session_id, - COUNT(m.id) FILTER (WHERE {}) as message_count, + {} as message_count, MAX({}) as last_message_timestamp, {} as sort_timestamp FROM sessions s @@ -2037,7 +2060,7 @@ impl SessionStorage { {} {} "#, - user_visible_message_sql("m.metadata_json"), + message_count_sql, normalized_message_timestamp, sort_timestamp_sql, message_join, @@ -2070,7 +2093,50 @@ impl SessionStorage { } let pool = self.pool().await?; - q.fetch_all(pool).await.map_err(Into::into) + if has_limit { + let mut tx = pool.begin().await?; + let mut sessions = q.fetch_all(&mut *tx).await?; + Self::populate_visible_message_counts(&mut tx, &mut sessions).await?; + tx.commit().await?; + Ok(sessions) + } else { + q.fetch_all(pool).await.map_err(Into::into) + } + } + + async fn populate_visible_message_counts( + tx: &mut sqlx::Transaction<'_, Sqlite>, + sessions: &mut [Session], + ) -> Result<()> { + let mut counts = HashMap::with_capacity(sessions.len()); + + for chunk in sessions.chunks(SESSION_COUNT_BATCH_SIZE) { + let placeholders = chunk.iter().map(|_| "?").collect::>().join(", "); + let sql = format!( + r#" + SELECT m.session_id, + COUNT(m.id) FILTER (WHERE {}) as message_count + FROM messages m + WHERE m.session_id IN ({}) + GROUP BY m.session_id + "#, + user_visible_message_sql("m.metadata_json"), + placeholders + ); + let mut q = sqlx::query_as::<_, (String, i64)>(AssertSqlSafe(sql)); + for session in chunk { + q = q.bind(&session.id); + } + for (session_id, message_count) in q.fetch_all(&mut **tx).await? { + counts.insert(session_id, message_count as usize); + } + } + + for session in sessions { + session.message_count = counts.get(&session.id).copied().unwrap_or_default(); + } + + Ok(()) } async fn list_sessions_by_types(&self, types: Option<&[SessionType]>) -> Result> { @@ -3112,6 +3178,10 @@ mod tests { let listed = sm.list_sessions().await.unwrap(); let listed_session = listed.iter().find(|s| s.id == session.id).unwrap(); assert_eq!(listed_session.message_count, 2); + + let limited = sm.list_sessions_with_limit(1).await.unwrap(); + assert_eq!(limited[0].id, session.id); + assert_eq!(limited[0].message_count, 2); } #[tokio::test] @@ -3568,6 +3638,114 @@ mod tests { assert_session_list_page(&sm, cursor.as_ref(), None, 2, &expected_ids[4..5], false).await; } + #[tokio::test] + async fn test_list_sessions_with_limit_counts_visible_and_legacy_messages() { + let temp_dir = TempDir::new().unwrap(); + let sm = SessionManager::new(temp_dir.path().to_path_buf()); + let older = create_session_for_list(&sm, "/tmp/session-list", false).await; + add_message_at(&sm, &older, "older", "2026-01-01T00:00:00Z").await; + + let selected = create_session_for_list(&sm, "/tmp/session-list", false).await; + sm.add_message( + &selected, + &Message::user().with_id("legacy").with_text("legacy"), + ) + .await + .unwrap(); + set_message_timestamp(&sm, &selected, "legacy", "2026-01-02T00:00:00Z").await; + sqlx::query( + "UPDATE messages SET metadata_json = '{}' WHERE session_id = ? AND message_id = ?", + ) + .bind(&selected) + .bind("legacy") + .execute(sm.storage().pool().await.unwrap()) + .await + .unwrap(); + + sm.add_message( + &selected, + &Message::user() + .with_id("hidden") + .with_text("hidden") + .with_visibility(false, true), + ) + .await + .unwrap(); + set_message_timestamp(&sm, &selected, "hidden", "2026-01-03T00:00:00Z").await; + + sm.add_message( + &selected, + &Message::assistant().with_id("visible").with_text("visible"), + ) + .await + .unwrap(); + set_message_timestamp(&sm, &selected, "visible", "2026-01-04T00:00:00Z").await; + + let sessions = sm.list_sessions_with_limit(1).await.unwrap(); + + assert_eq!(sessions.len(), 1); + assert_eq!(sessions[0].id, selected); + assert_eq!(sessions[0].message_count, 2); + assert_eq!( + sessions[0].last_message_at, + Some( + chrono::DateTime::parse_from_rfc3339("2026-01-04T00:00:00Z") + .unwrap() + .with_timezone(&Utc) + ) + ); + } + + #[tokio::test] + async fn test_session_list_paged_includes_hidden_only_session_with_zero_count() { + let temp_dir = TempDir::new().unwrap(); + let sm = SessionManager::new(temp_dir.path().to_path_buf()); + let hidden_only = create_session_for_list(&sm, "/tmp/session-list", false).await; + sm.add_message( + &hidden_only, + &Message::user() + .with_text("hidden") + .with_visibility(false, true), + ) + .await + .unwrap(); + create_session_for_list(&sm, "/tmp/session-list", false).await; + + let types = [SessionType::User]; + let page = sm + .list_sessions_paged(SessionListPageQuery { + filters: SessionListFilters { + types: Some(&types), + only_sessions_with_messages: true, + ..Default::default() + }, + cursor: None, + page_size: 10, + include_last_message_snippet: false, + }) + .await + .unwrap(); + + assert_eq!(page.sessions.len(), 1); + assert_eq!(page.sessions[0].id, hidden_only); + assert_eq!(page.sessions[0].message_count, 0); + assert!(page.sessions[0].last_message_at.is_some()); + } + + #[tokio::test] + async fn test_list_sessions_with_limit_returns_empty_session_with_zero_count() { + let temp_dir = TempDir::new().unwrap(); + let sm = SessionManager::new(temp_dir.path().to_path_buf()); + let empty = create_session_for_list(&sm, "/tmp/session-list", false).await; + + let sessions = sm.list_sessions_with_limit(1).await.unwrap(); + + assert_eq!(sessions.len(), 1); + assert_eq!(sessions[0].id, empty); + assert_eq!(sessions[0].message_count, 0); + assert_eq!(sessions[0].last_message_at, None); + } + #[tokio::test] async fn test_session_list_paged_sorts_by_last_message_at() { let temp_dir = TempDir::new().unwrap();