lean_ctx/core/context_kernel/
a2a_fixes.rs

1//! Standalone fixes for avoidable agent-to-agent context overhead.
2
3use std::collections::HashSet;
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7
8const DEFAULT_AGENT_BUDGET: usize = 1_000_000;
9const BUDGET_WARNING_PERCENT: f64 = 80.0;
10
11/// Lightweight representation of an agent scratchpad message.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct MessageEntry {
14    pub category: String,
15    pub body: String,
16    pub from_agent: String,
17    pub timestamp_epoch: u64,
18}
19
20/// Compatibility name for a lightweight scratchpad message.
21pub type ScratchpadEntry = MessageEntry;
22
23/// Compute the stable content ID used to track whether a message was read.
24pub fn message_id(message: &MessageEntry) -> String {
25    let mut hasher = blake3::Hasher::new();
26    for field in [&message.category, &message.body, &message.from_agent] {
27        hasher.update(&(field.len() as u64).to_le_bytes());
28        hasher.update(field.as_bytes());
29    }
30    hasher.finalize().to_hex()[..16].to_owned()
31}
32
33/// Properly filter messages to only return genuinely unread ones.
34pub fn filter_truly_unread<'a>(
35    messages: &'a [ScratchpadEntry],
36    read_ids: &HashSet<String>,
37) -> Vec<&'a ScratchpadEntry> {
38    messages
39        .iter()
40        .filter(|message| !read_ids.contains(&message_id(message)))
41        .collect()
42}
43
44/// Compute a real agent budget from configuration instead of `usize::MAX`.
45pub fn real_agent_budget(config_limit: Option<usize>) -> usize {
46    config_limit.unwrap_or(DEFAULT_AGENT_BUDGET)
47}
48
49/// Result of checking a requested token consumption against a real budget.
50#[derive(Debug, Clone, PartialEq)]
51pub enum BudgetCheckResult {
52    Allowed {
53        remaining: usize,
54    },
55    Denied {
56        over_by: usize,
57    },
58    Warning {
59        remaining: usize,
60        threshold_pct: f64,
61    },
62}
63
64/// Validate that a token consumption will not exceed the budget.
65pub fn check_budget_with_real_limit(
66    current_used: usize,
67    budget_limit: usize,
68    tokens_to_consume: usize,
69) -> BudgetCheckResult {
70    let projected = current_used.saturating_add(tokens_to_consume);
71    if projected > budget_limit {
72        return BudgetCheckResult::Denied {
73            over_by: projected - budget_limit,
74        };
75    }
76
77    let remaining = budget_limit - projected;
78    let threshold_pct = if budget_limit == 0 {
79        100.0
80    } else {
81        projected as f64 * 100.0 / budget_limit as f64
82    };
83    if threshold_pct >= BUDGET_WARNING_PERCENT {
84        BudgetCheckResult::Warning {
85            remaining,
86            threshold_pct,
87        }
88    } else {
89        BudgetCheckResult::Allowed { remaining }
90    }
91}
92
93/// Convert JSON to compact format after stripping null and empty values.
94pub fn compact_json(pretty: &Value) -> String {
95    match serde_json::to_string(&strip_nulls(pretty)) {
96        Ok(compact) => compact,
97        Err(_) => "null".to_owned(),
98    }
99}
100
101/// Estimate a JSON string's token count using four characters per token.
102pub fn json_token_estimate(json: &str) -> usize {
103    json.chars().count().div_ceil(4)
104}
105
106/// Strip null and empty values from a JSON value recursively.
107pub fn strip_nulls(value: &Value) -> Value {
108    prune_value(value).unwrap_or(Value::Null)
109}
110
111fn prune_value(value: &Value) -> Option<Value> {
112    match value {
113        Value::Null => None,
114        Value::String(text) if text.is_empty() => None,
115        Value::Array(values) => {
116            let values: Vec<_> = values.iter().filter_map(prune_value).collect();
117            (!values.is_empty()).then_some(Value::Array(values))
118        }
119        Value::Object(fields) => {
120            let fields: serde_json::Map<String, Value> = fields
121                .iter()
122                .filter_map(|(key, value)| prune_value(value).map(|value| (key.clone(), value)))
123                .collect();
124            (!fields.is_empty()).then_some(Value::Object(fields))
125        }
126        _ => Some(value.clone()),
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use std::collections::HashSet;
133
134    use serde_json::json;
135
136    use super::{
137        BudgetCheckResult, MessageEntry, check_budget_with_real_limit, compact_json,
138        filter_truly_unread, json_token_estimate, message_id, real_agent_budget, strip_nulls,
139    };
140
141    fn message(body: &str) -> MessageEntry {
142        MessageEntry {
143            category: "status".to_owned(),
144            body: body.to_owned(),
145            from_agent: "agent-a".to_owned(),
146            timestamp_epoch: 42,
147        }
148    }
149
150    #[test]
151    fn filter_removes_already_read() {
152        let messages: Vec<_> = (0..5).map(|index| message(&index.to_string())).collect();
153        let read_ids = [message_id(&messages[1]), message_id(&messages[3])]
154            .into_iter()
155            .collect::<HashSet<_>>();
156
157        let unread = filter_truly_unread(&messages, &read_ids);
158
159        assert_eq!(unread.len(), 3);
160        assert_eq!(unread[0].body, "0");
161        assert_eq!(unread[1].body, "2");
162        assert_eq!(unread[2].body, "4");
163    }
164
165    #[test]
166    fn message_id_is_deterministic() {
167        let first = message("same content");
168        let mut second = first.clone();
169        second.timestamp_epoch = 99;
170        assert_eq!(message_id(&first), message_id(&second));
171        assert_eq!(message_id(&first).len(), 16);
172    }
173
174    #[test]
175    fn real_budget_not_max() {
176        assert_eq!(real_agent_budget(None), 1_000_000);
177        assert_ne!(real_agent_budget(None), usize::MAX);
178        assert_eq!(real_agent_budget(Some(250_000)), 250_000);
179    }
180
181    #[test]
182    fn budget_denied_when_exceeded() {
183        assert_eq!(
184            check_budget_with_real_limit(900_000, 1_000_000, 200_000),
185            BudgetCheckResult::Denied { over_by: 100_000 }
186        );
187    }
188
189    #[test]
190    fn budget_warning_at_80_percent() {
191        assert_eq!(
192            check_budget_with_real_limit(800_000, 1_000_000, 10_000),
193            BudgetCheckResult::Warning {
194                remaining: 190_000,
195                threshold_pct: 81.0,
196            }
197        );
198    }
199
200    #[test]
201    fn compact_json_smaller() {
202        let value = json!({
203            "messages": [
204                {"body": "short", "metadata": null, "tags": []},
205                {"body": "reply", "metadata": null, "tags": []}
206            ],
207            "unused": null
208        });
209        let pretty = serde_json::to_string_pretty(&value).expect("test JSON must serialize");
210        let compact = compact_json(&value);
211        assert!(compact.len() * 2 < pretty.len());
212    }
213
214    #[test]
215    fn strip_nulls_removes_null_fields() {
216        assert_eq!(strip_nulls(&json!({"a": 1, "b": null})), json!({"a": 1}));
217    }
218
219    #[test]
220    fn token_estimate_reasonable() {
221        let estimate = json_token_estimate(&"x".repeat(1_000));
222        assert!((200..=300).contains(&estimate));
223    }
224}
lean_ctx/core/context_kernel/a2a_fixes.rs

lean_ctx/core/context_kernel/
a2a_fixes.rs
