From 82ef9292df6f97c12ae6820622ccc156cf5583ba Mon Sep 17 00:00:00 2001 From: Hashem Khalifa Date: Wed, 12 Aug 2026 23:49:59 +0200 Subject: [PATCH 1/4] perf(lcm): bound retention scans --- .../src/runtime/lcm/doctor.rs | 299 +++++++++++++++--- 1 file changed, 252 insertions(+), 47 deletions(-) diff --git a/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs b/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs index 8c524b277..40470785d 100644 --- a/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs +++ b/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs @@ -14,6 +14,7 @@ use super::{ }; const MAX_SAMPLES: usize = 20; +const RETENTION_SCAN_MESSAGE_LIMIT: usize = 10_000; const RETENTION_OLD_DAYS: f64 = 30.0; const RETENTION_HEAVY_CHARS: i64 = 128 * 1024; const SQLITE_IN_BATCH_SIZE: usize = 500; @@ -46,15 +47,21 @@ struct RepairRequest<'a> { } pub async fn doctor(conn: &Connection, request: DoctorRequest<'_>) -> Result { - let diagnostics = gather_diagnostics( - conn, - request.storage_root, - request.provider, - request.session_id, - &request.clean_config, - &request.gc_config, - ) - .await?; + let diagnostics = if request.mode == "retention" { + json!({ + "retention": retention_candidates(conn, request.provider, request.session_id).await?, + }) + } else { + gather_diagnostics( + conn, + request.storage_root, + request.provider, + request.session_id, + &request.clean_config, + &request.gc_config, + ) + .await? + }; let repairs = plan_and_apply_repairs( conn, RepairRequest { @@ -70,7 +77,11 @@ pub async fn doctor(conn: &Connection, request: DoctorRequest<'_>) -> Result::new(); + let mut messages_analyzed = 0_usize; + let mut scan_truncated = false; while let Some(row) = rows.next().await? { - analyzed += 1; + if messages_analyzed == RETENTION_SCAN_MESSAGE_LIMIT { + scan_truncated = true; + break; + } let session_id: String = row.get(0)?; - let message_count: i64 = row.get(1)?; - let retained_chars: i64 = row.get(2)?; - let first_message_at: i64 = row.get(3)?; - let last_message_at: i64 = row.get(4)?; - let summary_node_count: i64 = row.get(5)?; + let retained_chars: i64 = row.get(1)?; + let timestamp: i64 = row.get(2)?; + let stats = sessions.entry(session_id).or_default(); + stats.message_count += 1; + stats.retained_chars += retained_chars; + stats.first_message_at = Some( + stats + .first_message_at + .map_or(timestamp, |first| first.min(timestamp)), + ); + stats.last_message_at = Some( + stats + .last_message_at + .map_or(timestamp, |last| last.max(timestamp)), + ); + messages_analyzed += 1; + } + + let sampled_session_ids = sessions.keys().cloned().collect::>(); + let summary_counts = retention_summary_counts(conn, provider, &sampled_session_ids).await?; + let mut sessions = sessions.into_iter().collect::>(); + sessions.sort_by(|(left_id, left), (right_id, right)| { + right + .retained_chars + .cmp(&left.retained_chars) + .then_with(|| left.last_message_at.cmp(&right.last_message_at)) + .then_with(|| left_id.cmp(right_id)) + }); + let sessions_analyzed = sessions.len(); + let mut candidates = Vec::new(); + for (session_id, stats) in sessions { + let first_message_at = stats.first_message_at.unwrap_or_default(); + let last_message_at = stats.last_message_at.unwrap_or_default(); + let summary_node_count = summary_counts.get(&session_id).copied().unwrap_or_default(); let age_days = if last_message_at > 0 { (now.saturating_sub(last_message_at) as f64) / 86_400.0 } else { 0.0 }; let old = age_days >= RETENTION_OLD_DAYS; - let heavy = retained_chars >= RETENTION_HEAVY_CHARS; + let heavy = stats.retained_chars >= RETENTION_HEAVY_CHARS; let session_only = summary_node_count == 0; if old || heavy || session_only { candidates.push(json!({ "session_id": session_id, - "message_count": message_count, - "retained_chars": retained_chars, + "message_count": stats.message_count, + "retained_chars": stats.retained_chars, "first_message_at": first_message_at, "last_message_at": last_message_at, "age_days": age_days, @@ -827,12 +855,54 @@ async fn retention_candidates( Ok(json!({ "read_only": true, - "sessions_analyzed": analyzed, + "sessions_analyzed": sessions_analyzed, + "messages_analyzed": messages_analyzed, + "scan_limit": RETENTION_SCAN_MESSAGE_LIMIT, + "scan_truncated": scan_truncated, "candidate_count": candidates.len(), "candidates": candidates, })) } +#[derive(Default)] +struct RetentionSessionStats { + message_count: i64, + retained_chars: i64, + first_message_at: Option, + last_message_at: Option, +} + +async fn retention_summary_counts( + conn: &Connection, + provider: &str, + session_ids: &[String], +) -> Result, LcmError> { + let mut counts = BTreeMap::new(); + for session_chunk in session_ids.chunks(SQLITE_IN_BATCH_SIZE) { + if session_chunk.is_empty() { + continue; + } + let placeholders = sql_placeholders(session_chunk.len()); + let mut values = vec![SqlValue::Text(provider.to_string())]; + values.extend(session_chunk.iter().cloned().map(SqlValue::Text)); + let mut rows = conn + .query( + &format!( + "SELECT session_id, COUNT(*) + FROM lcm_summary_nodes + WHERE provider = ? AND session_id IN ({placeholders}) + GROUP BY session_id" + ), + values, + ) + .await?; + while let Some(row) = rows.next().await? { + counts.insert(row.get(0)?, row.get(1)?); + } + } + Ok(counts) +} + #[derive(Default)] struct CleanupSessionCandidate { classes: BTreeSet<&'static str>, @@ -1393,6 +1463,141 @@ mod tests { .map_err(|err| format!("read raw message count for {session_id}: {err}")) } + #[tokio::test] + async fn retention_mode_skips_unrelated_diagnostics() -> Result<(), String> { + let temp = tempfile::tempdir().map_err(|err| format!("create tempdir: {err}"))?; + let db_path = temp.path().join("sessions.db"); + let db = libsql::Builder::new_local(&db_path) + .build() + .await + .map_err(|err| format!("build test database: {err}"))?; + let conn = db + .connect() + .map_err(|err| format!("connect to test database: {err}"))?; + conn.execute_batch( + "CREATE TABLE lcm_raw_messages ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL, + index_text TEXT NOT NULL, + timestamp INTEGER + ); + CREATE TABLE lcm_summary_nodes ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL + ); + INSERT INTO lcm_raw_messages (provider, session_id, index_text, timestamp) + VALUES ('codex', 'retention-only-session', 'retention body', 1);", + ) + .await + .map_err(|err| format!("create retention-only fixture: {err}"))?; + + let result = doctor( + &conn, + DoctorRequest { + storage_root: temp.path(), + db_path: &db_path, + provider: "codex", + session_id: None, + mode: "retention", + apply: false, + clean_config: LcmCleanConfig::default(), + gc_config: LcmGcConfig::default(), + }, + ) + .await + .map_err(|err| format!("run retention doctor: {err}"))?; + + assert_eq!(result["status"], "ok"); + assert_eq!(result["mode"], "retention"); + assert_eq!(result["dry_run"], false); + assert_eq!(result["diagnostics"]["retention"]["read_only"], true); + assert_eq!(result["diagnostics"]["retention"]["candidate_count"], 1); + assert_eq!( + result["diagnostics"].as_object().map(serde_json::Map::len), + Some(1), + "retention mode must not run or report unrelated diagnostics" + ); + assert_eq!(result["repairs"]["planned_actions"], json!([])); + + let mut rows = conn + .query("SELECT COUNT(*) FROM lcm_raw_messages", ()) + .await + .map_err(|err| format!("count retained messages: {err}"))?; + let row = rows + .next() + .await + .map_err(|err| format!("read retained message count: {err}"))? + .ok_or_else(|| "missing retained message count row".to_string())?; + let message_count: i64 = row + .get(0) + .map_err(|err| format!("decode retained message count: {err}"))?; + assert_eq!(message_count, 1); + Ok(()) + } + + #[tokio::test] + async fn retention_mode_bounds_raw_message_scan_and_reports_truncation() -> Result<(), String> { + let temp = tempfile::tempdir().map_err(|err| format!("create tempdir: {err}"))?; + let db_path = temp.path().join("sessions.db"); + let db = libsql::Builder::new_local(&db_path) + .build() + .await + .map_err(|err| format!("build test database: {err}"))?; + let conn = db + .connect() + .map_err(|err| format!("connect to test database: {err}"))?; + conn.execute_batch(&format!( + "CREATE TABLE lcm_raw_messages ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL, + index_text TEXT NOT NULL, + timestamp INTEGER + ); + CREATE TABLE lcm_summary_nodes ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL + ); + WITH RECURSIVE messages(ordinal) AS ( + VALUES(1) + UNION ALL + SELECT ordinal + 1 FROM messages + WHERE ordinal <= {RETENTION_SCAN_MESSAGE_LIMIT} + ) + INSERT INTO lcm_raw_messages (provider, session_id, index_text, timestamp) + SELECT 'codex', 'bounded-retention-session', 'x', ordinal + FROM messages;" + )) + .await + .map_err(|err| format!("create oversized retention fixture: {err}"))?; + + let result = doctor( + &conn, + DoctorRequest { + storage_root: temp.path(), + db_path: &db_path, + provider: "codex", + session_id: Some("bounded-retention-session"), + mode: "retention", + apply: false, + clean_config: LcmCleanConfig::default(), + gc_config: LcmGcConfig::default(), + }, + ) + .await + .map_err(|err| format!("run bounded retention doctor: {err}"))?; + + let retention = &result["diagnostics"]["retention"]; + assert_eq!(retention["messages_analyzed"], RETENTION_SCAN_MESSAGE_LIMIT); + assert_eq!(retention["scan_limit"], RETENTION_SCAN_MESSAGE_LIMIT); + assert_eq!(retention["scan_truncated"], true); + assert_eq!(retention["sessions_analyzed"], 1); + assert_eq!( + retention["candidates"][0]["message_count"], + RETENTION_SCAN_MESSAGE_LIMIT + ); + Ok(()) + } + #[tokio::test] async fn clean_apply_backup_callback_runs_under_immediate_transaction() -> Result<(), String> { let temp = tempfile::tempdir().map_err(|err| format!("create tempdir: {err}"))?; From 417be8fa136b59c8b28cd07f904ea03591b4592e Mon Sep 17 00:00:00 2001 From: Hashem Khalifa Date: Thu, 13 Aug 2026 00:07:39 +0200 Subject: [PATCH 2/4] fix(lcm): avoid partial retention candidates --- .../src/runtime/lcm/doctor.rs | 208 ++++++++++++++---- 1 file changed, 168 insertions(+), 40 deletions(-) diff --git a/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs b/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs index 40470785d..156e1dcb4 100644 --- a/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs +++ b/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs @@ -770,48 +770,62 @@ async fn retention_candidates( session_id: Option<&str>, ) -> Result { let now = current_timestamp(); - let mut rows = conn - .query( - "SELECT session_id, LENGTH(index_text), COALESCE(timestamp, 0) - FROM lcm_raw_messages - WHERE provider = ?1 AND (?2 IS NULL OR session_id = ?2) - LIMIT ?3", - params![ - provider, - util::opt_text(session_id), - (RETENTION_SCAN_MESSAGE_LIMIT + 1) as i64 - ], - ) - .await?; + let (sampled_session_ids, session_scan_truncated) = + retention_session_ids(conn, provider, session_id).await?; + let per_session_limit = if sampled_session_ids.is_empty() { + 0 + } else { + RETENTION_SCAN_MESSAGE_LIMIT / sampled_session_ids.len() + }; let mut sessions = BTreeMap::::new(); let mut messages_analyzed = 0_usize; - let mut scan_truncated = false; - while let Some(row) = rows.next().await? { - if messages_analyzed == RETENTION_SCAN_MESSAGE_LIMIT { - scan_truncated = true; - break; + let mut incomplete_sessions = Vec::new(); + for sampled_session_id in &sampled_session_ids { + let mut rows = conn + .query( + "SELECT LENGTH(index_text), COALESCE(timestamp, 0) + FROM lcm_raw_messages + WHERE provider = ?1 AND session_id = ?2 + ORDER BY store_id + LIMIT ?3", + params![ + provider, + sampled_session_id.as_str(), + (per_session_limit + 1) as i64 + ], + ) + .await?; + let mut stats = RetentionSessionStats::default(); + while let Some(row) = rows.next().await? { + if stats.message_count == per_session_limit as i64 { + incomplete_sessions.push(sampled_session_id.clone()); + break; + } + let retained_chars: i64 = row.get(0)?; + let timestamp: i64 = row.get(1)?; + stats.message_count += 1; + stats.retained_chars += retained_chars; + stats.first_message_at = Some( + stats + .first_message_at + .map_or(timestamp, |first| first.min(timestamp)), + ); + stats.last_message_at = Some( + stats + .last_message_at + .map_or(timestamp, |last| last.max(timestamp)), + ); + messages_analyzed += 1; } - let session_id: String = row.get(0)?; - let retained_chars: i64 = row.get(1)?; - let timestamp: i64 = row.get(2)?; - let stats = sessions.entry(session_id).or_default(); - stats.message_count += 1; - stats.retained_chars += retained_chars; - stats.first_message_at = Some( - stats - .first_message_at - .map_or(timestamp, |first| first.min(timestamp)), - ); - stats.last_message_at = Some( - stats - .last_message_at - .map_or(timestamp, |last| last.max(timestamp)), - ); - messages_analyzed += 1; + sessions.insert(sampled_session_id.clone(), stats); } - let sampled_session_ids = sessions.keys().cloned().collect::>(); - let summary_counts = retention_summary_counts(conn, provider, &sampled_session_ids).await?; + let complete_session_ids = sampled_session_ids + .iter() + .filter(|session_id| !incomplete_sessions.contains(session_id)) + .cloned() + .collect::>(); + let summary_counts = retention_summary_counts(conn, provider, &complete_session_ids).await?; let mut sessions = sessions.into_iter().collect::>(); sessions.sort_by(|(left_id, left), (right_id, right)| { right @@ -821,8 +835,12 @@ async fn retention_candidates( .then_with(|| left_id.cmp(right_id)) }); let sessions_analyzed = sessions.len(); + let complete_sessions_analyzed = complete_session_ids.len(); let mut candidates = Vec::new(); for (session_id, stats) in sessions { + if incomplete_sessions.contains(&session_id) { + continue; + } let first_message_at = stats.first_message_at.unwrap_or_default(); let last_message_at = stats.last_message_at.unwrap_or_default(); let summary_node_count = summary_counts.get(&session_id).copied().unwrap_or_default(); @@ -856,14 +874,58 @@ async fn retention_candidates( Ok(json!({ "read_only": true, "sessions_analyzed": sessions_analyzed, + "complete_sessions_analyzed": complete_sessions_analyzed, "messages_analyzed": messages_analyzed, "scan_limit": RETENTION_SCAN_MESSAGE_LIMIT, - "scan_truncated": scan_truncated, + "per_session_scan_limit": per_session_limit, + "scan_truncated": session_scan_truncated || !incomplete_sessions.is_empty(), + "session_scan_truncated": session_scan_truncated, + "incomplete_session_count": incomplete_sessions.len(), + "incomplete_sessions": incomplete_sessions, "candidate_count": candidates.len(), "candidates": candidates, })) } +async fn retention_session_ids( + conn: &Connection, + provider: &str, + session_id: Option<&str>, +) -> Result<(Vec, bool), LcmError> { + let (sql, values) = if let Some(session_id) = session_id { + ( + "SELECT session_id + FROM lcm_raw_messages + WHERE provider = ? AND session_id = ? + LIMIT 1", + vec![ + SqlValue::Text(provider.to_string()), + SqlValue::Text(session_id.to_string()), + ], + ) + } else { + ( + "SELECT DISTINCT session_id + FROM lcm_raw_messages + WHERE provider = ? + ORDER BY session_id + LIMIT ?", + vec![ + SqlValue::Text(provider.to_string()), + SqlValue::Integer((MAX_SAMPLES + 1) as i64), + ], + ) + }; + let mut rows = conn.query(sql, values).await?; + let mut session_ids = Vec::new(); + while let Some(row) = rows.next().await? { + session_ids.push(row.get(0)?); + } + let truncated = session_id.is_none() && session_ids.len() > MAX_SAMPLES; + session_ids.truncate(MAX_SAMPLES); + Ok((session_ids, truncated)) +} + #[derive(Default)] struct RetentionSessionStats { message_count: i64, @@ -1478,6 +1540,7 @@ mod tests { "CREATE TABLE lcm_raw_messages ( provider TEXT NOT NULL, session_id TEXT NOT NULL, + store_id INTEGER PRIMARY KEY AUTOINCREMENT, index_text TEXT NOT NULL, timestamp INTEGER ); @@ -1550,6 +1613,7 @@ mod tests { "CREATE TABLE lcm_raw_messages ( provider TEXT NOT NULL, session_id TEXT NOT NULL, + store_id INTEGER PRIMARY KEY AUTOINCREMENT, index_text TEXT NOT NULL, timestamp INTEGER ); @@ -1591,9 +1655,73 @@ mod tests { assert_eq!(retention["scan_limit"], RETENTION_SCAN_MESSAGE_LIMIT); assert_eq!(retention["scan_truncated"], true); assert_eq!(retention["sessions_analyzed"], 1); + assert_eq!(retention["candidate_count"], 0); + assert_eq!(retention["complete_sessions_analyzed"], 0); + assert_eq!(retention["incomplete_session_count"], 1); assert_eq!( - retention["candidates"][0]["message_count"], - RETENTION_SCAN_MESSAGE_LIMIT + retention["incomplete_sessions"], + json!(["bounded-retention-session"]) + ); + Ok(()) + } + + #[tokio::test] + async fn retention_mode_samples_multiple_sessions_deterministically() -> Result<(), String> { + let temp = tempfile::tempdir().map_err(|err| format!("create tempdir: {err}"))?; + let db_path = temp.path().join("sessions.db"); + let db = libsql::Builder::new_local(&db_path) + .build() + .await + .map_err(|err| format!("build test database: {err}"))?; + let conn = db + .connect() + .map_err(|err| format!("connect to test database: {err}"))?; + conn.execute_batch(&format!( + "CREATE TABLE lcm_raw_messages ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL, + store_id INTEGER PRIMARY KEY AUTOINCREMENT, + index_text TEXT NOT NULL, + timestamp INTEGER + ); + CREATE TABLE lcm_summary_nodes ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL + ); + WITH RECURSIVE messages(ordinal) AS ( + VALUES(1) + UNION ALL + SELECT ordinal + 1 FROM messages + WHERE ordinal < {RETENTION_SCAN_MESSAGE_LIMIT} + ) + INSERT INTO lcm_raw_messages (provider, session_id, index_text, timestamp) + SELECT 'codex', 'a-large-session', 'x', ordinal + FROM messages; + INSERT INTO lcm_raw_messages (provider, session_id, index_text, timestamp) + VALUES ('codex', 'z-small-session', 'complete', 1);" + )) + .await + .map_err(|err| format!("create multi-session retention fixture: {err}"))?; + + let first = retention_candidates(&conn, "codex", None) + .await + .map_err(|err| format!("first retention scan: {err}"))?; + let second = retention_candidates(&conn, "codex", None) + .await + .map_err(|err| format!("second retention scan: {err}"))?; + + assert_eq!(first, second, "bounded sampling must be deterministic"); + assert_eq!(first["sessions_analyzed"], 2); + assert_eq!(first["complete_sessions_analyzed"], 1); + assert_eq!(first["incomplete_session_count"], 1); + assert_eq!(first["incomplete_sessions"], json!(["a-large-session"])); + assert!( + first["candidates"] + .as_array() + .unwrap() + .iter() + .any(|candidate| candidate["session_id"] == "z-small-session"), + "a large first session must not starve a later complete session" ); Ok(()) } From 2964f667d748de091b4c992a0ab04820546ade82 Mon Sep 17 00:00:00 2001 From: Hashem Khalifa Date: Thu, 13 Aug 2026 00:34:00 +0200 Subject: [PATCH 3/4] fix(lcm): bound retention discovery --- .../src/runtime/lcm/doctor.rs | 214 +++++++++++++----- 1 file changed, 154 insertions(+), 60 deletions(-) diff --git a/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs b/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs index 156e1dcb4..6abf7c245 100644 --- a/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs +++ b/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs @@ -770,8 +770,8 @@ async fn retention_candidates( session_id: Option<&str>, ) -> Result { let now = current_timestamp(); - let (sampled_session_ids, session_scan_truncated) = - retention_session_ids(conn, provider, session_id).await?; + let discovery = retention_session_ids(conn, provider, session_id).await?; + let sampled_session_ids = discovery.session_ids; let per_session_limit = if sampled_session_ids.is_empty() { 0 } else { @@ -817,15 +817,18 @@ async fn retention_candidates( ); messages_analyzed += 1; } - sessions.insert(sampled_session_id.clone(), stats); + if stats.message_count > 0 { + sessions.insert(sampled_session_id.clone(), stats); + } } - let complete_session_ids = sampled_session_ids - .iter() + let complete_session_ids = sessions + .keys() .filter(|session_id| !incomplete_sessions.contains(session_id)) .cloned() .collect::>(); - let summary_counts = retention_summary_counts(conn, provider, &complete_session_ids).await?; + let sessions_with_summaries = + retention_sessions_with_summaries(conn, provider, &complete_session_ids).await?; let mut sessions = sessions.into_iter().collect::>(); sessions.sort_by(|(left_id, left), (right_id, right)| { right @@ -843,7 +846,7 @@ async fn retention_candidates( } let first_message_at = stats.first_message_at.unwrap_or_default(); let last_message_at = stats.last_message_at.unwrap_or_default(); - let summary_node_count = summary_counts.get(&session_id).copied().unwrap_or_default(); + let has_summary = sessions_with_summaries.contains(&session_id); let age_days = if last_message_at > 0 { (now.saturating_sub(last_message_at) as f64) / 86_400.0 } else { @@ -851,7 +854,7 @@ async fn retention_candidates( }; let old = age_days >= RETENTION_OLD_DAYS; let heavy = stats.retained_chars >= RETENTION_HEAVY_CHARS; - let session_only = summary_node_count == 0; + let session_only = !has_summary; if old || heavy || session_only { candidates.push(json!({ "session_id": session_id, @@ -878,8 +881,10 @@ async fn retention_candidates( "messages_analyzed": messages_analyzed, "scan_limit": RETENTION_SCAN_MESSAGE_LIMIT, "per_session_scan_limit": per_session_limit, - "scan_truncated": session_scan_truncated || !incomplete_sessions.is_empty(), - "session_scan_truncated": session_scan_truncated, + "scan_truncated": discovery.truncated || !incomplete_sessions.is_empty(), + "session_scan_truncated": discovery.truncated, + "session_discovery_rows_scanned": discovery.rows_scanned, + "session_discovery_scan_limit": if session_id.is_some() { 0 } else { MAX_SAMPLES + 1 }, "incomplete_session_count": incomplete_sessions.len(), "incomplete_sessions": incomplete_sessions, "candidate_count": candidates.len(), @@ -887,43 +892,48 @@ async fn retention_candidates( })) } +struct RetentionSessionDiscovery { + session_ids: Vec, + truncated: bool, + rows_scanned: usize, +} + async fn retention_session_ids( conn: &Connection, provider: &str, session_id: Option<&str>, -) -> Result<(Vec, bool), LcmError> { - let (sql, values) = if let Some(session_id) = session_id { - ( - "SELECT session_id - FROM lcm_raw_messages - WHERE provider = ? AND session_id = ? - LIMIT 1", - vec![ - SqlValue::Text(provider.to_string()), - SqlValue::Text(session_id.to_string()), - ], - ) - } else { - ( - "SELECT DISTINCT session_id - FROM lcm_raw_messages - WHERE provider = ? - ORDER BY session_id - LIMIT ?", - vec![ - SqlValue::Text(provider.to_string()), - SqlValue::Integer((MAX_SAMPLES + 1) as i64), - ], - ) - }; +) -> Result { + if let Some(session_id) = session_id { + return Ok(RetentionSessionDiscovery { + session_ids: vec![session_id.to_string()], + truncated: false, + rows_scanned: 0, + }); + } + // Session discovery must stay on the narrow provider/session registry; + // reading distinct ids from lcm_raw_messages can visit its entire payload-heavy store. + let sql = "SELECT session_id + FROM sessions + WHERE provider = ? + ORDER BY session_id + LIMIT ?"; + let values = vec![ + SqlValue::Text(provider.to_string()), + SqlValue::Integer((MAX_SAMPLES + 1) as i64), + ]; let mut rows = conn.query(sql, values).await?; let mut session_ids = Vec::new(); while let Some(row) = rows.next().await? { session_ids.push(row.get(0)?); } - let truncated = session_id.is_none() && session_ids.len() > MAX_SAMPLES; + let rows_scanned = session_ids.len(); + let truncated = rows_scanned > MAX_SAMPLES; session_ids.truncate(MAX_SAMPLES); - Ok((session_ids, truncated)) + Ok(RetentionSessionDiscovery { + session_ids, + truncated, + rows_scanned, + }) } #[derive(Default)] @@ -934,35 +944,27 @@ struct RetentionSessionStats { last_message_at: Option, } -async fn retention_summary_counts( +async fn retention_sessions_with_summaries( conn: &Connection, provider: &str, session_ids: &[String], -) -> Result, LcmError> { - let mut counts = BTreeMap::new(); - for session_chunk in session_ids.chunks(SQLITE_IN_BATCH_SIZE) { - if session_chunk.is_empty() { - continue; - } - let placeholders = sql_placeholders(session_chunk.len()); - let mut values = vec![SqlValue::Text(provider.to_string())]; - values.extend(session_chunk.iter().cloned().map(SqlValue::Text)); +) -> Result, LcmError> { + let mut sessions_with_summaries = BTreeSet::new(); + for session_id in session_ids { let mut rows = conn .query( - &format!( - "SELECT session_id, COUNT(*) - FROM lcm_summary_nodes - WHERE provider = ? AND session_id IN ({placeholders}) - GROUP BY session_id" - ), - values, + "SELECT 1 + FROM lcm_summary_nodes + WHERE provider = ?1 AND session_id = ?2 + LIMIT 1", + params![provider, session_id.as_str()], ) .await?; - while let Some(row) = rows.next().await? { - counts.insert(row.get(0)?, row.get(1)?); + if rows.next().await?.is_some() { + sessions_with_summaries.insert(session_id.clone()); } } - Ok(counts) + Ok(sessions_with_summaries) } #[derive(Default)] @@ -1537,7 +1539,12 @@ mod tests { .connect() .map_err(|err| format!("connect to test database: {err}"))?; conn.execute_batch( - "CREATE TABLE lcm_raw_messages ( + "CREATE TABLE sessions ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL, + PRIMARY KEY(provider, session_id) + ); + CREATE TABLE lcm_raw_messages ( provider TEXT NOT NULL, session_id TEXT NOT NULL, store_id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -1548,6 +1555,8 @@ mod tests { provider TEXT NOT NULL, session_id TEXT NOT NULL ); + INSERT INTO sessions (provider, session_id) + VALUES ('codex', 'retention-only-session'); INSERT INTO lcm_raw_messages (provider, session_id, index_text, timestamp) VALUES ('codex', 'retention-only-session', 'retention body', 1);", ) @@ -1610,7 +1619,12 @@ mod tests { .connect() .map_err(|err| format!("connect to test database: {err}"))?; conn.execute_batch(&format!( - "CREATE TABLE lcm_raw_messages ( + "CREATE TABLE sessions ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL, + PRIMARY KEY(provider, session_id) + ); + CREATE TABLE lcm_raw_messages ( provider TEXT NOT NULL, session_id TEXT NOT NULL, store_id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -1621,6 +1635,8 @@ mod tests { provider TEXT NOT NULL, session_id TEXT NOT NULL ); + INSERT INTO sessions (provider, session_id) + VALUES ('codex', 'bounded-retention-session'); WITH RECURSIVE messages(ordinal) AS ( VALUES(1) UNION ALL @@ -1677,7 +1693,12 @@ mod tests { .connect() .map_err(|err| format!("connect to test database: {err}"))?; conn.execute_batch(&format!( - "CREATE TABLE lcm_raw_messages ( + "CREATE TABLE sessions ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL, + PRIMARY KEY(provider, session_id) + ); + CREATE TABLE lcm_raw_messages ( provider TEXT NOT NULL, session_id TEXT NOT NULL, store_id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -1688,6 +1709,8 @@ mod tests { provider TEXT NOT NULL, session_id TEXT NOT NULL ); + INSERT INTO sessions (provider, session_id) + VALUES ('codex', 'a-large-session'), ('codex', 'z-small-session'); WITH RECURSIVE messages(ordinal) AS ( VALUES(1) UNION ALL @@ -1726,6 +1749,77 @@ mod tests { Ok(()) } + #[tokio::test] + async fn retention_mode_bounds_session_discovery_before_large_raw_sessions() + -> Result<(), String> { + let temp = tempfile::tempdir().map_err(|err| format!("create tempdir: {err}"))?; + let db_path = temp.path().join("sessions.db"); + let db = libsql::Builder::new_local(&db_path) + .build() + .await + .map_err(|err| format!("build test database: {err}"))?; + let conn = db + .connect() + .map_err(|err| format!("connect to test database: {err}"))?; + conn.execute_batch(&format!( + "CREATE TABLE sessions ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL, + PRIMARY KEY(provider, session_id) + ); + CREATE TABLE lcm_raw_messages ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL, + store_id INTEGER PRIMARY KEY AUTOINCREMENT, + index_text TEXT NOT NULL, + timestamp INTEGER + ); + CREATE INDEX idx_lcm_raw_session_order + ON lcm_raw_messages(provider, session_id, store_id); + CREATE TABLE lcm_summary_nodes ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL + ); + WITH RECURSIVE session_numbers(ordinal) AS ( + VALUES(1) + UNION ALL + SELECT ordinal + 1 FROM session_numbers WHERE ordinal < 21 + ) + INSERT INTO sessions (provider, session_id) + SELECT 'codex', printf('session-%02d', ordinal) FROM session_numbers; + WITH RECURSIVE messages(ordinal) AS ( + VALUES(1) + UNION ALL + SELECT ordinal + 1 FROM messages + WHERE ordinal <= {RETENTION_SCAN_MESSAGE_LIMIT} + ) + INSERT INTO lcm_raw_messages (provider, session_id, index_text, timestamp) + SELECT 'codex', 'session-01', 'x', ordinal FROM messages; + INSERT INTO lcm_raw_messages (provider, session_id, index_text, timestamp) + SELECT 'codex', 'session-02', index_text, timestamp + FROM lcm_raw_messages WHERE session_id = 'session-01'; + INSERT INTO lcm_raw_messages (provider, session_id, index_text, timestamp) + SELECT 'codex', 'session-03', index_text, timestamp + FROM lcm_raw_messages WHERE session_id = 'session-01';" + )) + .await + .map_err(|err| format!("create bounded-discovery fixture: {err}"))?; + + let retention = retention_candidates(&conn, "codex", None) + .await + .map_err(|err| format!("run bounded retention scan: {err}"))?; + + assert_eq!(retention["session_discovery_rows_scanned"], MAX_SAMPLES + 1); + assert_eq!(retention["session_discovery_scan_limit"], MAX_SAMPLES + 1); + assert_eq!(retention["session_scan_truncated"], true); + assert!( + retention["messages_analyzed"].as_u64().unwrap_or_default() + <= RETENTION_SCAN_MESSAGE_LIMIT as u64, + "raw-message analysis must retain its total hard bound" + ); + Ok(()) + } + #[tokio::test] async fn clean_apply_backup_callback_runs_under_immediate_transaction() -> Result<(), String> { let temp = tempfile::tempdir().map_err(|err| format!("create tempdir: {err}"))?; From d8e99975258540991fe6aa2f2e29e963a8f70ae7 Mon Sep 17 00:00:00 2001 From: Hashem Khalifa Date: Thu, 13 Aug 2026 00:55:58 +0200 Subject: [PATCH 4/4] fix(lcm): preserve full diagnostics --- .../src/runtime/lcm/doctor.rs | 197 +++++++++++++++++- 1 file changed, 192 insertions(+), 5 deletions(-) diff --git a/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs b/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs index 6abf7c245..909bad9f3 100644 --- a/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs +++ b/crates/tracedecay-sessions/src/runtime/lcm/doctor.rs @@ -49,7 +49,7 @@ struct RepairRequest<'a> { pub async fn doctor(conn: &Connection, request: DoctorRequest<'_>) -> Result { let diagnostics = if request.mode == "retention" { json!({ - "retention": retention_candidates(conn, request.provider, request.session_id).await?, + "retention": retention_candidates_bounded(conn, request.provider, request.session_id).await?, }) } else { gather_diagnostics( @@ -768,6 +768,86 @@ async fn retention_candidates( conn: &Connection, provider: &str, session_id: Option<&str>, +) -> Result { + let now = current_timestamp(); + let mut rows = conn + .query( + "SELECT raw.session_id, + raw.message_count, + raw.retained_chars, + raw.first_message_at, + raw.last_message_at, + COALESCE(summary_counts.summary_node_count, 0) AS summary_node_count + FROM ( + SELECT session_id, + COUNT(*) AS message_count, + COALESCE(SUM(LENGTH(index_text)), 0) AS retained_chars, + MIN(COALESCE(timestamp, 0)) AS first_message_at, + MAX(COALESCE(timestamp, 0)) AS last_message_at + FROM lcm_raw_messages + WHERE provider = ?1 AND (?2 IS NULL OR session_id = ?2) + GROUP BY session_id + ) raw + LEFT JOIN ( + SELECT session_id, COUNT(*) AS summary_node_count + FROM lcm_summary_nodes + WHERE provider = ?1 AND (?2 IS NULL OR session_id = ?2) + GROUP BY session_id + ) summary_counts ON summary_counts.session_id = raw.session_id + ORDER BY raw.retained_chars DESC, raw.last_message_at ASC + LIMIT 100", + params![provider, util::opt_text(session_id)], + ) + .await?; + let mut candidates = Vec::new(); + let mut analyzed = 0; + while let Some(row) = rows.next().await? { + analyzed += 1; + let session_id: String = row.get(0)?; + let message_count: i64 = row.get(1)?; + let retained_chars: i64 = row.get(2)?; + let first_message_at: i64 = row.get(3)?; + let last_message_at: i64 = row.get(4)?; + let summary_node_count: i64 = row.get(5)?; + let age_days = if last_message_at > 0 { + (now.saturating_sub(last_message_at) as f64) / 86_400.0 + } else { + 0.0 + }; + let old = age_days >= RETENTION_OLD_DAYS; + let heavy = retained_chars >= RETENTION_HEAVY_CHARS; + let session_only = summary_node_count == 0; + if old || heavy || session_only { + candidates.push(json!({ + "session_id": session_id, + "message_count": message_count, + "retained_chars": retained_chars, + "first_message_at": first_message_at, + "last_message_at": last_message_at, + "age_days": age_days, + "old": old, + "heavy": heavy, + "session_only": session_only, + "protected": false, + })); + } + if candidates.len() >= MAX_SAMPLES { + break; + } + } + + Ok(json!({ + "read_only": true, + "sessions_analyzed": analyzed, + "candidate_count": candidates.len(), + "candidates": candidates, + })) +} + +async fn retention_candidates_bounded( + conn: &Connection, + provider: &str, + session_id: Option<&str>, ) -> Result { let now = current_timestamp(); let discovery = retention_session_ids(conn, provider, session_id).await?; @@ -1607,6 +1687,87 @@ mod tests { Ok(()) } + #[tokio::test] + async fn diagnose_retention_includes_candidate_beyond_bounded_session_sample() + -> Result<(), String> { + let temp = tempfile::tempdir().map_err(|err| format!("create tempdir: {err}"))?; + let db_path = temp.path().join("sessions.db"); + let db = libsql::Builder::new_local(&db_path) + .build() + .await + .map_err(|err| format!("build test database: {err}"))?; + let conn = db + .connect() + .map_err(|err| format!("connect to test database: {err}"))?; + conn.execute_batch( + "CREATE TABLE sessions ( + provider TEXT NOT NULL, + session_id TEXT NOT NULL, + project_key TEXT NOT NULL, + project_path TEXT NOT NULL, + title TEXT, + started_at INTEGER, + PRIMARY KEY(provider, session_id) + ); + CREATE TABLE session_messages ( + provider TEXT NOT NULL, + message_id TEXT NOT NULL, + session_id TEXT NOT NULL, + role TEXT NOT NULL, + timestamp INTEGER, + ordinal INTEGER NOT NULL, + text TEXT NOT NULL, + metadata_json TEXT, + PRIMARY KEY(provider, message_id) + );", + ) + .await + .map_err(|err| format!("create test session schema: {err}"))?; + schema::ensure_lcm_schema(&conn) + .await + .map_err(|err| format!("create test LCM schema: {err}"))?; + + for ordinal in 1..=MAX_SAMPLES + 1 { + let session_id = format!("session-{ordinal:02}"); + let message_id = format!("message-{ordinal:02}"); + insert_test_clean_candidate(&conn, temp.path(), &session_id, &message_id).await?; + } + conn.execute( + "UPDATE lcm_raw_messages + SET index_text = printf('%.*c', ?1, 'x') + WHERE provider = 'cursor' AND session_id = 'session-21'", + params![RETENTION_HEAVY_CHARS + 1], + ) + .await + .map_err(|err| format!("make late retention candidate largest: {err}"))?; + + let result = doctor( + &conn, + DoctorRequest { + storage_root: temp.path(), + db_path: &db_path, + provider: "cursor", + session_id: None, + mode: "diagnose", + apply: false, + clean_config: LcmCleanConfig::default(), + gc_config: LcmGcConfig::default(), + }, + ) + .await + .map_err(|err| format!("run diagnose doctor: {err}"))?; + + assert!( + result["diagnostics"]["retention"]["candidates"] + .as_array() + .is_some_and(|candidates| candidates + .iter() + .any(|candidate| candidate["session_id"] == "session-21")), + "non-retention diagnostics must retain the largest candidate beyond the bounded first-{MAX_SAMPLES} sessions" + ); + Ok(()) + } + #[tokio::test] async fn retention_mode_bounds_raw_message_scan_and_reports_truncation() -> Result<(), String> { let temp = tempfile::tempdir().map_err(|err| format!("create tempdir: {err}"))?; @@ -1726,14 +1887,40 @@ mod tests { .await .map_err(|err| format!("create multi-session retention fixture: {err}"))?; - let first = retention_candidates(&conn, "codex", None) + let first = retention_candidates_bounded(&conn, "codex", None) .await .map_err(|err| format!("first retention scan: {err}"))?; - let second = retention_candidates(&conn, "codex", None) + let second = retention_candidates_bounded(&conn, "codex", None) .await .map_err(|err| format!("second retention scan: {err}"))?; - assert_eq!(first, second, "bounded sampling must be deterministic"); + for field in [ + "sessions_analyzed", + "complete_sessions_analyzed", + "messages_analyzed", + "scan_limit", + "per_session_scan_limit", + "scan_truncated", + "session_scan_truncated", + "incomplete_session_count", + "incomplete_sessions", + "candidate_count", + ] { + assert_eq!(first[field], second[field], "{field} must be deterministic"); + } + let candidate_session_ids = |retention: &Value| { + retention["candidates"] + .as_array() + .into_iter() + .flatten() + .map(|candidate| candidate["session_id"].clone()) + .collect::>() + }; + assert_eq!( + candidate_session_ids(&first), + candidate_session_ids(&second), + "bounded candidate sampling must be deterministic" + ); assert_eq!(first["sessions_analyzed"], 2); assert_eq!(first["complete_sessions_analyzed"], 1); assert_eq!(first["incomplete_session_count"], 1); @@ -1805,7 +1992,7 @@ mod tests { .await .map_err(|err| format!("create bounded-discovery fixture: {err}"))?; - let retention = retention_candidates(&conn, "codex", None) + let retention = retention_candidates_bounded(&conn, "codex", None) .await .map_err(|err| format!("run bounded retention scan: {err}"))?;