feat: add ecc2 graph recall memory ranking

This commit is contained in:
Affaan Mustafa
2026-04-10 05:49:43 -07:00
parent 23348a21a6
commit 7a13564a8b
4 changed files with 506 additions and 27 deletions

View File

@@ -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<String>,
/// 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(

View File

@@ -190,6 +190,14 @@ pub struct ContextGraphEntityDetail {
pub incoming: Vec<ContextGraphRelation>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ContextGraphRecallEntry {
pub entity: ContextGraphEntity,
pub score: u64,
pub matched_terms: Vec<String>,
pub relation_count: usize,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct ContextGraphSyncStats {
pub sessions_scanned: usize,

View File

@@ -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<Vec<ContextGraphRecallEntry>> {
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::<Result<Vec<_>, _>>()?;
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::<Vec<_>>();
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<Conte
})
}
fn context_graph_recall_terms(query: &str) -> Vec<String> {
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<String> {
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<chrono::Utc>,
now: chrono::DateTime<chrono::Utc>,
) -> 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"));

View File

@@ -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::<Vec<_>>(),
)
@@ -1228,7 +1229,7 @@ impl Dashboard {
self.theme_palette(),
)
})
.collect::<Vec<_>>(),
.collect::<Vec<_>>(),
)
}
@@ -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<String> {
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::<Vec<_>>();
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<Line<'static>> {
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();