Display memory scores in status column

Move score display from name (via label()) to status column for cleaner
layout. Score now appears right of tokens for all memory nodes.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>

src/mind/mod.rs

@@ -33,6 +33,36 @@ pub use unconscious::{UnconsciousSnapshot, Unconscious};
 
 use crate::agent::context::{AstNode, NodeBody, Section, Ast, ContextState};
 
+fn match_scores(
+    nodes: &[AstNode],
+    scores: &std::collections::BTreeMap<String, f64>,
+) -> Vec<(usize, f64)> {
+    nodes.iter().enumerate()
+        .filter_map(|(i, node)| {
+            if let AstNode::Leaf(leaf) = node {
+                if let NodeBody::Memory { key, .. } = leaf.body() {
+                    return scores.get(key.as_str()).map(|&s| (i, s));
+                }
+            }
+            None
+        }).collect()
+}
+
+fn find_memory_by_key(ctx: &ContextState, key: &str) -> Option<(Section, usize)> {
+    [(Section::Identity, ctx.identity()), (Section::Conversation, ctx.conversation())]
+        .into_iter()
+        .find_map(|(section, nodes)| {
+            nodes.iter().enumerate().find_map(|(i, node)| {
+                if let AstNode::Leaf(leaf) = node {
+                    if let NodeBody::Memory { key: k, .. } = leaf.body() {
+                        if k == key { return Some((section, i)); }
+                    }
+                }
+                None
+            })
+        })
+}
+
 fn load_memory_scores(ctx: &mut ContextState, path: &std::path::Path) {
     let data = match std::fs::read_to_string(path) {
         Ok(d) => d,
@@ -42,25 +72,24 @@ fn load_memory_scores(ctx: &mut ContextState, path: &std::path::Path) {
         Ok(s) => s,
         Err(_) => return,
     };
-    let mut applied = 0;
-    for i in 0..ctx.conversation().len() {
-        if let AstNode::Leaf(leaf) = &ctx.conversation()[i] {
-            if let NodeBody::Memory { key, .. } = leaf.body() {
-                if let Some(&s) = scores.get(key.as_str()) {
-                    ctx.set_score(Section::Conversation, i, Some(s));
-                    applied += 1;
-                }
-            }
-        }
+    let identity_scores = match_scores(ctx.identity(), &scores);
+    let conv_scores = match_scores(ctx.conversation(), &scores);
+    let applied = identity_scores.len() + conv_scores.len();
+    for (i, s) in identity_scores {
+        ctx.set_score(Section::Identity, i, Some(s));
+    }
+    for (i, s) in conv_scores {
+        ctx.set_score(Section::Conversation, i, Some(s));
     }
     if applied > 0 {
         dbglog!("[scoring] loaded {} scores from {}", applied, path.display());
     }
 }
 
-/// Collect scored memory keys from conversation entries.
+/// Collect scored memory keys from identity and conversation entries.
 fn collect_memory_scores(ctx: &ContextState) -> std::collections::BTreeMap<String, f64> {
-    ctx.conversation().iter()
+    ctx.identity().iter()
+        .chain(ctx.conversation().iter())
         .filter_map(|node| {
             if let AstNode::Leaf(leaf) = node {
                 if let NodeBody::Memory { key, score: Some(s), .. } = leaf.body() {
@@ -531,14 +560,10 @@ impl Mind {
                     async move {
                         let scores_snapshot = {
                             let mut ctx = agent.context.lock().await;
-                            for i in 0..ctx.conversation().len() {
-                                if let AstNode::Leaf(leaf) = &ctx.conversation()[i] {
-                                    if let NodeBody::Memory { key: k, .. } = leaf.body() {
-                                        if *k == key {
-                                            ctx.set_score(Section::Conversation, i, Some(score));
-                                        }
-                                    }
-                                }
+                            // Find memory by key in identity or conversation
+                            let found = find_memory_by_key(&ctx, &key);
+                            if let Some((section, i)) = found {
+                                ctx.set_score(section, i, Some(score));
                             }
                             let snapshot = collect_memory_scores(&ctx);
                             drop(ctx);

src/subconscious/learn.rs

@@ -62,8 +62,16 @@ fn build_token_ids(
     for node in context.system() {
         ids.extend(node.token_ids());
     }
+    // Identity nodes can be filtered by key for scoring
     for node in context.identity() {
-        ids.extend(node.token_ids());
+        let skip = match &filter {
+            Filter::SkipKey(key) => memory_key(node) == Some(*key),
+            Filter::SkipAllMemories => is_memory(node),
+            _ => false,
+        };
+        if !skip {
+            ids.extend(node.token_ids());
+        }
     }
     for node in context.journal() {
         ids.extend(node.token_ids());
@@ -175,7 +183,9 @@ pub async fn score_memories(
     // Collect memory keys and response indices under a brief lock
     let (memory_keys, response_indices) = {
         let ctx = agent.context.lock().await;
-        let mut keys: Vec<String> = ctx.conversation().iter()
+        // Include identity nodes and conversation memories
+        let mut keys: Vec<String> = ctx.identity().iter()
+            .chain(ctx.conversation().iter())
             .filter_map(|node| memory_key(node).map(String::from))
             .collect();
         keys.dedup();
@@ -331,7 +341,10 @@ where
 
     {
         let store = &*store_arc;
-        for (i, node) in context.conversation().iter().enumerate() {
+        // Identity nodes always score at position 0; conversation nodes at their index
+        let identity_nodes = context.identity().iter().map(|n| (0, n));
+        let conv_nodes = context.conversation().iter().enumerate();
+        for (pos, node) in identity_nodes.chain(conv_nodes) {
             if let Some(key) = memory_key(node) {
                 if !seen.insert(key.to_owned()) { continue; }
                 let last_scored = store.get_node(key)
@@ -340,7 +353,7 @@ where
                     .map(|n| n.last_scored)
                     .unwrap_or(0);
                 if now - last_scored >= max_age_secs {
-                    candidates.push((i, key.to_owned(), last_scored));
+                    candidates.push((pos, key.to_owned(), last_scored));
                 }
             }
         }

src/user/context.rs

@@ -38,16 +38,13 @@ impl ConsciousScreen {
         for node in ctx.conversation() {
             if let AstNode::Leaf(leaf) = node {
                 if let NodeBody::Memory { key, score, text } = leaf.body() {
-                    let status = match score {
-                        Some(s) => { scored += 1; format!("{:.2}", s) }
-                        None => { unscored += 1; String::new() }
-                    };
+                    if score.is_some() { scored += 1; } else { unscored += 1; }
                     mem_children.push(SectionView {
-                        name: key.clone(),
+                        name: format!("mem: {}", key),
                         tokens: node.tokens(),
                         content: text.clone(),
                         children: Vec::new(),
-                        status,
+                        status: score.map(|s| format!("{:.2}", s)).unwrap_or_default(),
                     });
                 }
             }

src/user/widgets.rs

@@ -6,7 +6,7 @@ use ratatui::{
     widgets::{Block, Borders},
     crossterm::event::KeyCode,
 };
-use crate::agent::context::{AstNode, Ast};
+use crate::agent::context::{AstNode, Ast, NodeBody};
 
 #[derive(Debug, Clone)]
 pub struct SectionView {
@@ -20,13 +20,22 @@ pub struct SectionView {
 
 fn node_to_view(node: &AstNode) -> SectionView {
     match node {
-        AstNode::Leaf(leaf) => SectionView {
-            name: node.label(),
-            tokens: node.tokens(),
-            content: leaf.body().text().to_string(),
-            children: Vec::new(),
-            status: String::new(),
-        },
+        AstNode::Leaf(leaf) => {
+            let (name, status) = match leaf.body() {
+                NodeBody::Memory { key, score, .. } => {
+                    let s = score.map(|v| format!("{:.2}", v)).unwrap_or_default();
+                    (format!("mem: {}", key), s)
+                }
+                _ => (node.label(), String::new()),
+            };
+            SectionView {
+                name,
+                tokens: node.tokens(),
+                content: leaf.body().text().to_string(),
+                children: Vec::new(),
+                status,
+            }
+        }
         AstNode::Branch { children, .. } => {
             let child_views: Vec<SectionView> = children.iter()
                 .map(|c| node_to_view(c))