From 7a13564a8bef2ca774b7ec75f2a7a4a8fb92e6fc Mon Sep 17 00:00:00 2001 From: Affaan Mustafa Date: Fri, 10 Apr 2026 05:49:43 -0700 Subject: [PATCH] feat: add ecc2 graph recall memory ranking --- ecc2/src/main.rs | 149 ++++++++++++++++++++++++ ecc2/src/session/mod.rs | 8 ++ ecc2/src/session/store.rs | 232 +++++++++++++++++++++++++++++++++++++- ecc2/src/tui/dashboard.rs | 144 +++++++++++++++++++---- 4 files changed, 506 insertions(+), 27 deletions(-) diff --git a/ecc2/src/main.rs b/ecc2/src/main.rs index 24038d34..7b8d8e7a 100644 --- a/ecc2/src/main.rs +++ b/ecc2/src/main.rs @@ -457,6 +457,20 @@ enum GraphCommands { #[arg(long)] json: bool, }, + /// Recall relevant context graph entities for a query + Recall { + /// Filter by source session ID or alias + #[arg(long)] + session_id: Option, + /// Natural-language query used for recall scoring + query: String, + /// Maximum entities to return + #[arg(long, default_value_t = 8)] + limit: usize, + /// Emit machine-readable JSON instead of the human summary + #[arg(long)] + json: bool, + }, /// Show one entity plus its incoming and outgoing relations Show { /// Entity ID @@ -1229,6 +1243,27 @@ async fn main() -> Result<()> { println!("{}", format_graph_relations_human(&relations)); } } + GraphCommands::Recall { + session_id, + query, + limit, + json, + } => { + let resolved_session_id = session_id + .as_deref() + .map(|value| resolve_session_id(&db, value)) + .transpose()?; + let entries = + db.recall_context_entities(resolved_session_id.as_deref(), &query, limit)?; + if json { + println!("{}", serde_json::to_string_pretty(&entries)?); + } else { + println!( + "{}", + format_graph_recall_human(&entries, resolved_session_id.as_deref(), &query) + ); + } + } GraphCommands::Show { entity_id, limit, @@ -2214,6 +2249,49 @@ fn format_graph_relations_human(relations: &[session::ContextGraphRelation]) -> lines.join("\n") } +fn format_graph_recall_human( + entries: &[session::ContextGraphRecallEntry], + session_id: Option<&str>, + query: &str, +) -> String { + if entries.is_empty() { + return format!("No relevant context graph entities found for query: {query}"); + } + + let scope = session_id + .map(short_session) + .unwrap_or_else(|| "all sessions".to_string()); + let mut lines = vec![format!( + "Relevant memory: {} entries for \"{}\" ({scope})", + entries.len(), + query + )]; + for entry in entries { + let mut line = format!( + "- #{} [{}] {} | score {} | relations {}", + entry.entity.id, + entry.entity.entity_type, + entry.entity.name, + entry.score, + entry.relation_count + ); + if let Some(session_id) = entry.entity.session_id.as_deref() { + line.push_str(&format!(" | {}", short_session(session_id))); + } + lines.push(line); + if !entry.matched_terms.is_empty() { + lines.push(format!(" matches {}", entry.matched_terms.join(", "))); + } + if let Some(path) = entry.entity.path.as_deref() { + lines.push(format!(" path {path}")); + } + if !entry.entity.summary.is_empty() { + lines.push(format!(" summary {}", entry.entity.summary)); + } + } + lines.join("\n") +} + fn format_graph_entity_detail_human(detail: &session::ContextGraphEntityDetail) -> String { let mut lines = vec![format_graph_entity_human(&detail.entity)]; lines.push(String::new()); @@ -4114,6 +4192,40 @@ mod tests { } } + #[test] + fn cli_parses_graph_recall_command() { + let cli = Cli::try_parse_from([ + "ecc", + "graph", + "recall", + "--session-id", + "latest", + "--limit", + "4", + "--json", + "auth callback recovery", + ]) + .expect("graph recall should parse"); + + match cli.command { + Some(Commands::Graph { + command: + GraphCommands::Recall { + session_id, + query, + limit, + json, + }, + }) => { + assert_eq!(session_id.as_deref(), Some("latest")); + assert_eq!(query, "auth callback recovery"); + assert_eq!(limit, 4); + assert!(json); + } + _ => panic!("expected graph recall subcommand"), + } + } + #[test] fn format_decisions_human_renders_details() { let text = format_decisions_human( @@ -4196,6 +4308,43 @@ mod tests { assert!(text.contains("[contains] #6 dashboard.rs -> render_metrics")); } + #[test] + fn format_graph_recall_human_renders_scores_and_matches() { + let text = format_graph_recall_human( + &[session::ContextGraphRecallEntry { + entity: session::ContextGraphEntity { + id: 11, + session_id: Some("sess-12345678".to_string()), + entity_type: "file".to_string(), + name: "callback.ts".to_string(), + path: Some("src/routes/auth/callback.ts".to_string()), + summary: "Handles auth callback recovery".to_string(), + metadata: BTreeMap::new(), + created_at: chrono::DateTime::parse_from_rfc3339("2026-04-10T01:02:03Z") + .unwrap() + .with_timezone(&chrono::Utc), + updated_at: chrono::DateTime::parse_from_rfc3339("2026-04-10T01:02:03Z") + .unwrap() + .with_timezone(&chrono::Utc), + }, + score: 319, + matched_terms: vec![ + "auth".to_string(), + "callback".to_string(), + "recovery".to_string(), + ], + relation_count: 2, + }], + Some("sess-12345678"), + "auth callback recovery", + ); + + assert!(text.contains("Relevant memory: 1 entries")); + assert!(text.contains("[file] callback.ts | score 319 | relations 2")); + assert!(text.contains("matches auth, callback, recovery")); + assert!(text.contains("path src/routes/auth/callback.ts")); + } + #[test] fn format_graph_sync_stats_human_renders_counts() { let text = format_graph_sync_stats_human( diff --git a/ecc2/src/session/mod.rs b/ecc2/src/session/mod.rs index 583d8bde..7bd380f1 100644 --- a/ecc2/src/session/mod.rs +++ b/ecc2/src/session/mod.rs @@ -190,6 +190,14 @@ pub struct ContextGraphEntityDetail { pub incoming: Vec, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ContextGraphRecallEntry { + pub entity: ContextGraphEntity, + pub score: u64, + pub matched_terms: Vec, + pub relation_count: usize, +} + #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] pub struct ContextGraphSyncStats { pub sessions_scanned: usize, diff --git a/ecc2/src/session/store.rs b/ecc2/src/session/store.rs index b32bb0ea..c0f465d3 100644 --- a/ecc2/src/session/store.rs +++ b/ecc2/src/session/store.rs @@ -14,9 +14,9 @@ use crate::observability::{ToolCallEvent, ToolLogEntry, ToolLogPage}; use super::output::{OutputLine, OutputStream, OUTPUT_BUFFER_LIMIT}; use super::{ default_project_label, default_task_group_label, normalize_group_label, ContextGraphEntity, - ContextGraphEntityDetail, ContextGraphRelation, ContextGraphSyncStats, DecisionLogEntry, - FileActivityAction, FileActivityEntry, Session, SessionAgentProfile, SessionMessage, - SessionMetrics, SessionState, WorktreeInfo, + ContextGraphEntityDetail, ContextGraphRecallEntry, ContextGraphRelation, ContextGraphSyncStats, + DecisionLogEntry, FileActivityAction, FileActivityEntry, Session, SessionAgentProfile, + SessionMessage, SessionMetrics, SessionState, WorktreeInfo, }; pub struct StateStore { @@ -2024,6 +2024,82 @@ impl StateStore { Ok(entries) } + pub fn recall_context_entities( + &self, + session_id: Option<&str>, + query: &str, + limit: usize, + ) -> Result> { + if limit == 0 { + return Ok(Vec::new()); + } + + let terms = context_graph_recall_terms(query); + if terms.is_empty() { + return Ok(Vec::new()); + } + + let candidate_limit = (limit.saturating_mul(12)).clamp(24, 512); + let mut stmt = self.conn.prepare( + "SELECT e.id, e.session_id, e.entity_type, e.name, e.path, e.summary, e.metadata_json, + e.created_at, e.updated_at, + ( + SELECT COUNT(*) + FROM context_graph_relations r + WHERE r.from_entity_id = e.id OR r.to_entity_id = e.id + ) AS relation_count + FROM context_graph_entities e + WHERE (?1 IS NULL OR e.session_id = ?1) + ORDER BY e.updated_at DESC, e.id DESC + LIMIT ?2", + )?; + + let candidates = stmt + .query_map( + rusqlite::params![session_id, candidate_limit as i64], + |row| { + let entity = map_context_graph_entity(row)?; + let relation_count = row.get::<_, i64>(9)?.max(0) as usize; + Ok((entity, relation_count)) + }, + )? + .collect::, _>>()?; + + let now = chrono::Utc::now(); + let mut entries = candidates + .into_iter() + .filter_map(|(entity, relation_count)| { + let matched_terms = context_graph_matched_terms(&entity, &terms); + if matched_terms.is_empty() { + return None; + } + + Some(ContextGraphRecallEntry { + score: context_graph_recall_score( + matched_terms.len(), + relation_count, + entity.updated_at, + now, + ), + entity, + matched_terms, + relation_count, + }) + }) + .collect::>(); + + entries.sort_by(|left, right| { + right + .score + .cmp(&left.score) + .then_with(|| right.entity.updated_at.cmp(&left.entity.updated_at)) + .then_with(|| right.entity.id.cmp(&left.entity.id)) + }); + entries.truncate(limit); + + Ok(entries) + } + pub fn get_context_entity_detail( &self, entity_id: i64, @@ -3071,6 +3147,65 @@ fn map_context_graph_relation(row: &rusqlite::Row<'_>) -> rusqlite::Result Vec { + let mut terms = Vec::new(); + for raw_term in + query.split(|c: char| !(c.is_ascii_alphanumeric() || matches!(c, '_' | '-' | '.' | '/'))) + { + let term = raw_term.trim().to_ascii_lowercase(); + if term.len() < 3 || terms.iter().any(|existing| existing == &term) { + continue; + } + terms.push(term); + } + terms +} + +fn context_graph_matched_terms(entity: &ContextGraphEntity, terms: &[String]) -> Vec { + let mut haystacks = vec![ + entity.entity_type.to_ascii_lowercase(), + entity.name.to_ascii_lowercase(), + entity.summary.to_ascii_lowercase(), + ]; + if let Some(path) = entity.path.as_ref() { + haystacks.push(path.to_ascii_lowercase()); + } + for (key, value) in &entity.metadata { + haystacks.push(key.to_ascii_lowercase()); + haystacks.push(value.to_ascii_lowercase()); + } + + let mut matched = Vec::new(); + for term in terms { + if haystacks.iter().any(|value| value.contains(term)) { + matched.push(term.clone()); + } + } + matched +} + +fn context_graph_recall_score( + matched_term_count: usize, + relation_count: usize, + updated_at: chrono::DateTime, + now: chrono::DateTime, +) -> u64 { + let recency_bonus = { + let age = now.signed_duration_since(updated_at); + if age <= chrono::Duration::hours(1) { + 9 + } else if age <= chrono::Duration::hours(24) { + 6 + } else if age <= chrono::Duration::days(7) { + 3 + } else { + 0 + } + }; + + (matched_term_count as u64 * 100) + (relation_count.min(9) as u64 * 10) + recency_bonus +} + fn parse_store_timestamp( raw: String, column: usize, @@ -3855,6 +3990,89 @@ mod tests { Ok(()) } + #[test] + fn recall_context_entities_ranks_matching_entities() -> Result<()> { + let tempdir = TestDir::new("store-context-recall")?; + let db = StateStore::open(&tempdir.path().join("state.db"))?; + let now = Utc::now(); + + db.insert_session(&Session { + id: "session-1".to_string(), + task: "Investigate auth callback recovery".to_string(), + project: "ecc-tools".to_string(), + task_group: "incident".to_string(), + agent_type: "claude".to_string(), + working_dir: PathBuf::from("/tmp"), + state: SessionState::Running, + pid: None, + worktree: None, + created_at: now, + updated_at: now, + last_heartbeat_at: now, + metrics: SessionMetrics::default(), + })?; + + let callback = db.upsert_context_entity( + Some("session-1"), + "file", + "callback.ts", + Some("src/routes/auth/callback.ts"), + "Handles auth callback recovery and billing portal fallback", + &BTreeMap::from([("area".to_string(), "auth".to_string())]), + )?; + let recovery = db.upsert_context_entity( + Some("session-1"), + "decision", + "Use recovery-first callback routing", + None, + "Auth callback recovery should prefer the billing portal", + &BTreeMap::new(), + )?; + let unrelated = db.upsert_context_entity( + Some("session-1"), + "file", + "dashboard.rs", + Some("ecc2/src/tui/dashboard.rs"), + "Renders the TUI dashboard", + &BTreeMap::new(), + )?; + + db.upsert_context_relation( + Some("session-1"), + callback.id, + recovery.id, + "supports", + "Callback route supports recovery-first routing", + )?; + db.upsert_context_relation( + Some("session-1"), + callback.id, + unrelated.id, + "references", + "Callback route references the dashboard summary", + )?; + + let results = + db.recall_context_entities(Some("session-1"), "Investigate auth callback recovery", 3)?; + + assert_eq!(results.len(), 2); + assert_eq!(results[0].entity.id, callback.id); + assert!(results[0].matched_terms.iter().any(|term| term == "auth")); + assert!(results[0] + .matched_terms + .iter() + .any(|term| term == "callback")); + assert!(results[0] + .matched_terms + .iter() + .any(|term| term == "recovery")); + assert_eq!(results[0].relation_count, 2); + assert_eq!(results[1].entity.id, recovery.id); + assert!(!results.iter().any(|entry| entry.entity.id == unrelated.id)); + + Ok(()) + } + #[test] fn context_graph_detail_includes_incoming_and_outgoing_relations() -> Result<()> { let tempdir = TestDir::new("store-context-relations")?; @@ -4139,8 +4357,12 @@ mod tests { .expect("session entity should exist"); let relations = db.list_context_relations(Some(session_entity.id), 10)?; assert_eq!(relations.len(), 3); - assert!(relations.iter().any(|relation| relation.relation_type == "decided")); - assert!(relations.iter().any(|relation| relation.relation_type == "modify")); + assert!(relations + .iter() + .any(|relation| relation.relation_type == "decided")); + assert!(relations + .iter() + .any(|relation| relation.relation_type == "modify")); assert!(relations .iter() .any(|relation| relation.relation_type == "delegates_to")); diff --git a/ecc2/src/tui/dashboard.rs b/ecc2/src/tui/dashboard.rs index 396026fb..824691a9 100644 --- a/ecc2/src/tui/dashboard.rs +++ b/ecc2/src/tui/dashboard.rs @@ -844,7 +844,8 @@ impl Dashboard { self.render_searchable_graph(&lines) } else { Text::from( - lines.into_iter() + lines + .into_iter() .map(|line| Line::from(line.text)) .collect::>(), ) @@ -1228,7 +1229,7 @@ impl Dashboard { self.theme_palette(), ) }) - .collect::>(), + .collect::>(), ) } @@ -3296,7 +3297,10 @@ impl Dashboard { return; } - if !matches!(self.output_mode, OutputMode::SessionOutput | OutputMode::ContextGraph) { + if !matches!( + self.output_mode, + OutputMode::SessionOutput | OutputMode::ContextGraph + ) { self.set_operator_note( "search is only available in session output or graph view".to_string(), ); @@ -4914,8 +4918,12 @@ impl Dashboard { .selected_agent_type() .unwrap_or(self.cfg.default_agent.as_str()) .to_string(); - self.selected_route_preview = - self.build_route_preview(&session_id, &selected_agent_type, team.total, &route_candidates); + self.selected_route_preview = self.build_route_preview( + &session_id, + &selected_agent_type, + team.total, + &route_candidates, + ); delegated.sort_by_key(|delegate| { ( delegate_attention_priority(delegate), @@ -5027,8 +5035,7 @@ impl Dashboard { if message.to_session != session_id || message.msg_type != "task_handoff" { return None; } - manager::parse_task_handoff_task(&message.content) - .or_else(|| Some(message.content)) + manager::parse_task_handoff_task(&message.content).or_else(|| Some(message.content)) }) } @@ -5289,6 +5296,60 @@ impl Dashboard { lines } + fn session_graph_recall_lines(&self, session: &Session) -> Vec { + let query = session.task.trim(); + if query.is_empty() { + return Vec::new(); + } + + let Ok(entries) = self.db.recall_context_entities(None, query, 4) else { + return Vec::new(); + }; + + let entries = entries + .into_iter() + .filter(|entry| { + !(entry.entity.entity_type == "session" && entry.entity.name == session.id) + }) + .take(3) + .collect::>(); + if entries.is_empty() { + return Vec::new(); + } + + let mut lines = vec!["Relevant memory".to_string()]; + for entry in entries { + let mut line = format!( + "- #{} [{}] {} | score {} | relations {}", + entry.entity.id, + entry.entity.entity_type, + truncate_for_dashboard(&entry.entity.name, 60), + entry.score, + entry.relation_count + ); + if let Some(session_id) = entry.entity.session_id.as_deref() { + if session_id != session.id { + line.push_str(&format!(" | {}", format_session_id(session_id))); + } + } + lines.push(line); + if !entry.matched_terms.is_empty() { + lines.push(format!(" matches {}", entry.matched_terms.join(", "))); + } + if let Some(path) = entry.entity.path.as_deref() { + lines.push(format!(" path {}", truncate_for_dashboard(path, 72))); + } + if !entry.entity.summary.is_empty() { + lines.push(format!( + " summary {}", + truncate_for_dashboard(&entry.entity.summary, 72) + )); + } + } + + lines + } + fn visible_git_status_lines(&self) -> Vec> { self.selected_git_status_entries .iter() @@ -6254,6 +6315,7 @@ impl Dashboard { } } } + lines.extend(self.session_graph_recall_lines(session)); lines.extend(self.session_graph_metrics_lines(&session.id)); let file_overlaps = self .db @@ -10213,8 +10275,12 @@ diff --git a/src/lib.rs b/src/lib.rs\n\ let mut dashboard = test_dashboard(vec![focus.clone(), review.clone()], 0); dashboard.db.insert_session(&focus)?; dashboard.db.insert_session(&review)?; - dashboard.db.insert_decision(&focus.id, "Alpha graph path", &[], "planner path")?; - dashboard.db.insert_decision(&review.id, "Beta graph path", &[], "review path")?; + dashboard + .db + .insert_decision(&focus.id, "Alpha graph path", &[], "planner path")?; + dashboard + .db + .insert_decision(&review.id, "Beta graph path", &[], "review path")?; dashboard.toggle_context_graph_mode(); dashboard.toggle_search_scope(); @@ -10254,8 +10320,12 @@ diff --git a/src/lib.rs b/src/lib.rs\n\ let mut dashboard = test_dashboard(vec![focus.clone(), review.clone()], 0); dashboard.db.insert_session(&focus)?; dashboard.db.insert_session(&review)?; - dashboard.db.insert_decision(&focus.id, "alpha local graph", &[], "planner path")?; - dashboard.db.insert_decision(&review.id, "alpha remote graph", &[], "review path")?; + dashboard + .db + .insert_decision(&focus.id, "alpha local graph", &[], "planner path")?; + dashboard + .db + .insert_decision(&review.id, "alpha remote graph", &[], "review path")?; dashboard.toggle_context_graph_mode(); dashboard.toggle_search_scope(); @@ -10274,7 +10344,10 @@ diff --git a/src/lib.rs b/src/lib.rs\n\ dashboard.operator_note.as_deref(), Some("graph search /alpha.* match 2/2 | all sessions") ); - assert_ne!(dashboard.selected_session_id().map(str::to_string), first_session); + assert_ne!( + dashboard.selected_session_id().map(str::to_string), + first_session + ); Ok(()) } @@ -10322,14 +10395,7 @@ diff --git a/src/lib.rs b/src/lib.rs\n\ 1, 1, ); - let delegate = sample_session( - "delegate-87654321", - "coder", - SessionState::Idle, - None, - 1, - 1, - ); + let delegate = sample_session("delegate-87654321", "coder", SessionState::Idle, None, 1, 1); let dashboard = test_dashboard(vec![focus.clone(), delegate.clone()], 0); dashboard.db.insert_session(&focus)?; dashboard.db.insert_session(&delegate)?; @@ -10354,6 +10420,38 @@ diff --git a/src/lib.rs b/src/lib.rs\n\ Ok(()) } + #[test] + fn selected_session_metrics_text_includes_relevant_memory() -> Result<()> { + let mut focus = sample_session( + "focus-12345678", + "planner", + SessionState::Running, + None, + 1, + 1, + ); + focus.task = "Investigate auth callback recovery".to_string(); + let mut memory = sample_session("memory-87654321", "coder", SessionState::Idle, None, 1, 1); + memory.task = "Auth callback recovery notes".to_string(); + let dashboard = test_dashboard(vec![focus.clone(), memory.clone()], 0); + dashboard.db.insert_session(&focus)?; + dashboard.db.insert_session(&memory)?; + dashboard.db.upsert_context_entity( + Some(&memory.id), + "file", + "callback.ts", + Some("src/routes/auth/callback.ts"), + "Handles auth callback recovery and billing fallback", + &BTreeMap::from([("area".to_string(), "auth".to_string())]), + )?; + + let text = dashboard.selected_session_metrics_text(); + assert!(text.contains("Relevant memory")); + assert!(text.contains("[file] callback.ts")); + assert!(text.contains("matches auth, callback, recovery")); + Ok(()) + } + #[test] fn worktree_diff_columns_split_removed_and_added_lines() { let patch = "\ @@ -11178,8 +11276,10 @@ diff --git a/src/lib.rs b/src/lib.rs 24, ); - let mut dashboard = - test_dashboard(vec![lead.clone(), older_worker.clone(), auth_worker.clone()], 0); + let mut dashboard = test_dashboard( + vec![lead.clone(), older_worker.clone(), auth_worker.clone()], + 0, + ); dashboard.db.insert_session(&lead).unwrap(); dashboard.db.insert_session(&older_worker).unwrap(); dashboard.db.insert_session(&auth_worker).unwrap();