Skip to main content

lean_ctx/core/
chain_compression.rs

1//! Chain Compression — context deduplication between agent hops (P11 / DIM 4).
2//!
3//! When context flows through a chain of agents, each hop typically carries
4//! overlapping references. This module identifies shared context and produces
5//! minimal deltas, ensuring only novel information crosses agent boundaries.
6
7use std::collections::{BTreeMap, BTreeSet};
8
9use serde::{Deserialize, Serialize};
10
11pub const CHAIN_COMPRESSION_SCHEMA_VERSION: u16 = 1;
12const MAX_CHAIN_HISTORY: usize = 64;
13
14/// A content-addressed context item tracked across hops.
15#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
16pub struct ChainContextItem {
17    pub content_ref: String,
18    pub freshness_ref: String,
19    pub hop_introduced: u16,
20    pub last_referenced_hop: u16,
21}
22
23/// Delta produced when forwarding context to the next hop.
24#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
25pub struct ChainDelta {
26    pub schema_version: u16,
27    pub chain_id: String,
28    pub from_hop: u16,
29    pub to_hop: u16,
30    pub added_refs: Vec<String>,
31    pub removed_refs: Vec<String>,
32    pub unchanged_count: usize,
33    pub total_refs_at_target: usize,
34    pub compression_ratio: f64,
35}
36
37/// Tracks context across hops in a single chain, computes minimal deltas.
38pub struct ChainCompressionTracker {
39    chain_id: String,
40    items_by_hop: BTreeMap<u16, BTreeSet<String>>,
41    all_items: BTreeMap<String, ChainContextItem>,
42    current_hop: u16,
43}
44
45impl ChainCompressionTracker {
46    #[must_use]
47    pub fn new(chain_id: String) -> Self {
48        Self {
49            chain_id,
50            items_by_hop: BTreeMap::new(),
51            all_items: BTreeMap::new(),
52            current_hop: 0,
53        }
54    }
55
56    /// Register context refs at a given hop.
57    pub fn register_hop(
58        &mut self,
59        hop: u16,
60        content_refs: Vec<String>,
61        freshness_refs: Vec<String>,
62    ) -> Result<(), ChainCompressionError> {
63        if content_refs.len() != freshness_refs.len() {
64            return Err(ChainCompressionError::MismatchedLengths);
65        }
66        if self.items_by_hop.len() >= MAX_CHAIN_HISTORY {
67            let oldest = *self.items_by_hop.keys().next().unwrap_or(&0);
68            self.items_by_hop.remove(&oldest);
69        }
70        let mut hop_refs = BTreeSet::new();
71        for (content_ref, freshness_ref) in content_refs.into_iter().zip(freshness_refs) {
72            hop_refs.insert(content_ref.clone());
73            self.all_items
74                .entry(content_ref.clone())
75                .and_modify(|item| {
76                    item.last_referenced_hop = hop;
77                    item.freshness_ref.clone_from(&freshness_ref);
78                })
79                .or_insert(ChainContextItem {
80                    content_ref,
81                    freshness_ref,
82                    hop_introduced: hop,
83                    last_referenced_hop: hop,
84                });
85        }
86        self.items_by_hop.insert(hop, hop_refs);
87        self.current_hop = self.current_hop.max(hop);
88        Ok(())
89    }
90
91    /// Compute the delta between two hops.
92    pub fn compute_delta(
93        &self,
94        from_hop: u16,
95        to_hop: u16,
96    ) -> Result<ChainDelta, ChainCompressionError> {
97        let from_refs = self
98            .items_by_hop
99            .get(&from_hop)
100            .ok_or(ChainCompressionError::HopNotFound(from_hop))?;
101        let to_refs = self
102            .items_by_hop
103            .get(&to_hop)
104            .ok_or(ChainCompressionError::HopNotFound(to_hop))?;
105
106        let added: Vec<String> = to_refs.difference(from_refs).cloned().collect();
107        let removed: Vec<String> = from_refs.difference(to_refs).cloned().collect();
108        let unchanged_count = to_refs.intersection(from_refs).count();
109        let total_refs_at_target = to_refs.len();
110
111        let full_transfer_size = total_refs_at_target.max(1);
112        let delta_size = (added.len() + removed.len()).max(1);
113        let compression_ratio = 1.0 - (delta_size as f64 / full_transfer_size as f64);
114
115        Ok(ChainDelta {
116            schema_version: CHAIN_COMPRESSION_SCHEMA_VERSION,
117            chain_id: self.chain_id.clone(),
118            from_hop,
119            to_hop,
120            added_refs: added,
121            removed_refs: removed,
122            unchanged_count,
123            total_refs_at_target,
124            compression_ratio: compression_ratio.max(0.0),
125        })
126    }
127
128    /// Compute what a new hop needs vs the parent hop (forward delta).
129    pub fn forward_delta(
130        &self,
131        parent_hop: u16,
132        child_refs: &[String],
133    ) -> Result<ChainDelta, ChainCompressionError> {
134        let parent_refs = self
135            .items_by_hop
136            .get(&parent_hop)
137            .ok_or(ChainCompressionError::HopNotFound(parent_hop))?;
138        let child_set: BTreeSet<&String> = child_refs.iter().collect();
139        let parent_set: BTreeSet<&String> = parent_refs.iter().collect();
140
141        let added: Vec<String> = child_set
142            .difference(&parent_set)
143            .map(|s| (*s).clone())
144            .collect();
145        let removed: Vec<String> = parent_set
146            .difference(&child_set)
147            .map(|s| (*s).clone())
148            .collect();
149        let unchanged_count = child_set.intersection(&parent_set).count();
150        let total = child_refs.len();
151
152        let full_size = total.max(1);
153        let delta_size = (added.len() + removed.len()).max(1);
154        let ratio = 1.0 - (delta_size as f64 / full_size as f64);
155
156        Ok(ChainDelta {
157            schema_version: CHAIN_COMPRESSION_SCHEMA_VERSION,
158            chain_id: self.chain_id.clone(),
159            from_hop: parent_hop,
160            to_hop: self.current_hop + 1,
161            added_refs: added,
162            removed_refs: removed,
163            unchanged_count,
164            total_refs_at_target: total,
165            compression_ratio: ratio.max(0.0),
166        })
167    }
168
169    /// Items that haven't been referenced since `stale_threshold_hop`.
170    pub fn stale_items(&self, stale_threshold_hop: u16) -> Vec<&ChainContextItem> {
171        self.all_items
172            .values()
173            .filter(|item| item.last_referenced_hop < stale_threshold_hop)
174            .collect()
175    }
176
177    pub fn current_hop(&self) -> u16 {
178        self.current_hop
179    }
180
181    pub fn total_tracked_items(&self) -> usize {
182        self.all_items.len()
183    }
184}
185
186// ─── Errors ──────────────────────────────────────────────────────────────────
187
188#[derive(Debug, thiserror::Error)]
189pub enum ChainCompressionError {
190    #[error("content_refs and freshness_refs have different lengths")]
191    MismatchedLengths,
192    #[error("hop {0} not found in chain history")]
193    HopNotFound(u16),
194}
195
196// ─── Tests ───────────────────────────────────────────────────────────────────
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201
202    #[test]
203    fn delta_shows_only_differences() {
204        let mut tracker = ChainCompressionTracker::new("chain:1".into());
205        tracker
206            .register_hop(
207                0,
208                vec!["blake3:a".into(), "blake3:b".into(), "blake3:c".into()],
209                vec!["fresh:1".into(), "fresh:1".into(), "fresh:1".into()],
210            )
211            .unwrap();
212        tracker
213            .register_hop(
214                1,
215                vec!["blake3:b".into(), "blake3:c".into(), "blake3:d".into()],
216                vec!["fresh:1".into(), "fresh:1".into(), "fresh:1".into()],
217            )
218            .unwrap();
219
220        let delta = tracker.compute_delta(0, 1).unwrap();
221        assert_eq!(delta.added_refs, vec!["blake3:d"]);
222        assert_eq!(delta.removed_refs, vec!["blake3:a"]);
223        assert_eq!(delta.unchanged_count, 2);
224        assert!(delta.compression_ratio > 0.0);
225    }
226
227    #[test]
228    fn forward_delta_for_child() {
229        let mut tracker = ChainCompressionTracker::new("chain:2".into());
230        tracker
231            .register_hop(
232                0,
233                vec!["blake3:x".into(), "blake3:y".into(), "blake3:z".into()],
234                vec!["fresh:1".into(), "fresh:1".into(), "fresh:1".into()],
235            )
236            .unwrap();
237
238        let child_refs = vec!["blake3:y".into(), "blake3:z".into(), "blake3:new".into()];
239        let delta = tracker.forward_delta(0, &child_refs).unwrap();
240        assert_eq!(delta.added_refs, vec!["blake3:new"]);
241        assert_eq!(delta.removed_refs, vec!["blake3:x"]);
242        assert_eq!(delta.unchanged_count, 2);
243    }
244
245    #[test]
246    fn identical_hops_yield_perfect_compression() {
247        let mut tracker = ChainCompressionTracker::new("chain:3".into());
248        let refs = vec!["blake3:a".into(), "blake3:b".into()];
249        let fresh = vec!["fresh:1".into(), "fresh:1".into()];
250        tracker
251            .register_hop(0, refs.clone(), fresh.clone())
252            .unwrap();
253        tracker.register_hop(1, refs, fresh).unwrap();
254
255        let delta = tracker.compute_delta(0, 1).unwrap();
256        assert!(delta.added_refs.is_empty());
257        assert!(delta.removed_refs.is_empty());
258        assert_eq!(delta.unchanged_count, 2);
259    }
260
261    #[test]
262    fn stale_detection() {
263        let mut tracker = ChainCompressionTracker::new("chain:4".into());
264        tracker
265            .register_hop(
266                0,
267                vec!["blake3:old".into(), "blake3:active".into()],
268                vec!["fresh:1".into(), "fresh:1".into()],
269            )
270            .unwrap();
271        tracker
272            .register_hop(
273                3,
274                vec!["blake3:active".into(), "blake3:new".into()],
275                vec!["fresh:2".into(), "fresh:2".into()],
276            )
277            .unwrap();
278
279        let stale = tracker.stale_items(2);
280        assert_eq!(stale.len(), 1);
281        assert_eq!(stale[0].content_ref, "blake3:old");
282    }
283
284    #[test]
285    fn mismatched_lengths_rejected() {
286        let mut tracker = ChainCompressionTracker::new("chain:5".into());
287        assert!(
288            tracker
289                .register_hop(0, vec!["blake3:a".into()], vec![])
290                .is_err()
291        );
292    }
293}