Skip to main content

talon_cli/mcp/session/
ledger.rs

1use std::collections::{HashMap, HashSet, VecDeque};
2
3pub const MAX_TURNS: usize = 32;
4
5#[derive(Debug, Clone)]
6pub struct TurnRecord {
7    pub turn_id: String,
8    pub query_fingerprint: String,
9    pub injected: Vec<InjectedChunk>,
10    pub suppressed: Vec<SuppressedRecall>,
11    pub skipped: bool,
12}
13
14#[derive(Debug, Clone)]
15pub struct InjectedChunk {
16    /// Deterministic chunk ID — see `chunk_id.rs` for how this is derived.
17    pub chunk_id: String,
18    pub path: String,
19    pub score: f64,
20}
21
22#[derive(Debug, Clone)]
23pub struct SuppressedRecall {
24    pub chunk_id: String,
25    pub path: String,
26    pub score: f64,
27    pub reason: SuppressionReason,
28}
29
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub enum SuppressionReason {
32    SameChunkRecentlyInjected,
33    SameNoteRecentlyInjected,
34    BelowConfidenceGate,
35}
36
37/// Bounded turn history.
38#[derive(Debug)]
39pub struct TurnLedger {
40    turns: VecDeque<TurnRecord>,
41    /// `chunk_id` → set of `turn_id`s where it was injected.
42    injected_chunks: HashMap<String, HashSet<String>>,
43    /// `path` → set of `turn_id`s where it was injected.
44    injected_notes: HashMap<String, HashSet<String>>,
45}
46
47impl Default for TurnLedger {
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53impl TurnLedger {
54    #[must_use]
55    pub fn new() -> Self {
56        Self {
57            turns: VecDeque::with_capacity(MAX_TURNS),
58            injected_chunks: HashMap::new(),
59            injected_notes: HashMap::new(),
60        }
61    }
62
63    /// Record a completed turn with injected and suppressed chunks.
64    pub fn record_turn(&mut self, record: TurnRecord) {
65        // If we're at capacity, remove the oldest turn and clean up its chunk/note refs.
66        if self.turns.len() >= MAX_TURNS
67            && let Some(evicted) = self.turns.pop_front()
68        {
69            self.remove_turn_refs(&evicted.turn_id);
70        }
71
72        for chunk in &record.injected {
73            self.injected_chunks
74                .entry(chunk.chunk_id.clone())
75                .or_default()
76                .insert(record.turn_id.clone());
77            self.injected_notes
78                .entry(chunk.path.clone())
79                .or_default()
80                .insert(record.turn_id.clone());
81        }
82        self.turns.push_back(record);
83    }
84
85    /// Removes `turn_id` references from chunk/note maps when a turn is evicted.
86    fn remove_turn_refs(&mut self, turn_id: &str) {
87        for ids in self.injected_chunks.values_mut() {
88            ids.remove(turn_id);
89        }
90        for ids in self.injected_notes.values_mut() {
91            ids.remove(turn_id);
92        }
93    }
94
95    /// Returns how many of the last N turns contained this `chunk_id`.
96    #[must_use]
97    pub fn chunk_injected_in_last_n(&self, chunk_id: &str, n: usize) -> usize {
98        let recent_ids: HashSet<&str> = self
99            .turns
100            .iter()
101            .rev()
102            .take(n)
103            .map(|t| t.turn_id.as_str())
104            .collect();
105        self.injected_chunks.get(chunk_id).map_or(0, |ids| {
106            ids.iter()
107                .filter(|id| recent_ids.contains(id.as_str()))
108                .count()
109        })
110    }
111
112    /// Returns how many of the last N turns contained this note `path`.
113    #[must_use]
114    pub fn note_injected_in_last_n(&self, path: &str, n: usize) -> usize {
115        let recent_ids: HashSet<&str> = self
116            .turns
117            .iter()
118            .rev()
119            .take(n)
120            .map(|t| t.turn_id.as_str())
121            .collect();
122        self.injected_notes.get(path).map_or(0, |ids| {
123            ids.iter()
124                .filter(|id| recent_ids.contains(id.as_str()))
125                .count()
126        })
127    }
128
129    /// Returns how many turns have elapsed since this chunk was last injected.
130    /// Returns `None` if the chunk has never been injected.
131    #[must_use]
132    pub fn turns_since_chunk_last_injected(&self, chunk_id: &str) -> Option<usize> {
133        let turn_ids = self.injected_chunks.get(chunk_id)?;
134        // Find the position from the end of the most recent matching turn.
135        self.turns
136            .iter()
137            .rev()
138            .position(|t| turn_ids.contains(t.turn_id.as_str()))
139    }
140
141    /// Returns how many turns have elapsed since any chunk from this note path
142    /// was last injected. Returns `None` if never injected.
143    #[must_use]
144    pub fn turns_since_note_last_injected(&self, path: &str) -> Option<usize> {
145        let turn_ids = self.injected_notes.get(path)?;
146        self.turns
147            .iter()
148            .rev()
149            .position(|t| turn_ids.contains(t.turn_id.as_str()))
150    }
151
152    /// Returns the fingerprint of the most recent turn, if any.
153    #[must_use]
154    pub fn last_fingerprint(&self) -> Option<&str> {
155        self.turns.back().map(|t| t.query_fingerprint.as_str())
156    }
157}