lean_ctx/proxy/
sticky_tools.rs

1//! Sticky CCR tool injection for provider prefix-cache stability.
2//!
3//! Once a conversation has seen at least one CCR handle, the `ctx_expand` tool
4//! definition stays in `tools[]` for the remainder of the session, even on turns
5//! with no active markers. This prevents the tool-list change from busting the
6//! provider prompt-cache prefix.
7
8use std::collections::HashSet;
9use std::sync::{Mutex, OnceLock};
10
11use serde_json::{Value, json};
12
13const MAX_TRACKED: usize = 4096;
14const EXPAND_TOOL_NAME: &str = "ctx_expand";
15
16fn active_sessions() -> &'static Mutex<HashSet<u64>> {
17    static SESSIONS: OnceLock<Mutex<HashSet<u64>>> = OnceLock::new();
18    SESSIONS.get_or_init(|| Mutex::new(HashSet::new()))
19}
20
21/// Mark a conversation as having used CCR.
22pub fn mark_ccr_active(conv_id: u64) {
23    if let Ok(mut guard) = active_sessions().lock() {
24        if guard.len() >= MAX_TRACKED
25            && !guard.contains(&conv_id)
26            && let Some(&oldest) = guard.iter().next()
27        {
28            guard.remove(&oldest);
29        }
30        guard.insert(conv_id);
31    }
32}
33
34/// Returns `true` if this conversation has ever used CCR.
35pub fn is_ccr_active(conv_id: u64) -> bool {
36    active_sessions().lock().is_ok_and(|g| g.contains(&conv_id))
37}
38
39fn expand_tool_definition() -> Value {
40    json!({
41        "name": EXPAND_TOOL_NAME,
42        "description": "Retrieve the original uncompressed content of a compressed tool result by its tee handle or hash.",
43        "input_schema": {
44            "type": "object",
45            "properties": {
46                "id": {
47                    "type": "string",
48                    "description": "The tee path or hash handle from the compressed output."
49                }
50            },
51            "required": ["id"]
52        }
53    })
54}
55
56/// Ensure `ctx_expand` is present in the `tools` array when CCR is active.
57/// Returns `true` if the tool was injected.
58pub fn ensure_tool_present(conv_id: u64, doc: &mut Value) -> bool {
59    if !is_ccr_active(conv_id) {
60        return false;
61    }
62
63    let tools = match doc.get_mut("tools") {
64        Some(Value::Array(arr)) => arr,
65        Some(_) => return false,
66        None => {
67            doc["tools"] = Value::Array(Vec::new());
68            doc["tools"].as_array_mut().unwrap()
69        }
70    };
71
72    let already_present = tools.iter().any(|t| {
73        t.get("name")
74            .and_then(Value::as_str)
75            .is_some_and(|n| n == EXPAND_TOOL_NAME)
76    });
77
78    if already_present {
79        return false;
80    }
81
82    tools.push(expand_tool_definition());
83    true
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use serde_json::json;
90
91    // High conv IDs reserved for these unit tests. Parallel tests in compress.rs
92    // and history_prune.rs call mark_ccr_active(0); clearing the global set
93    // races with them, so each test marks its own ID instead.
94    const NO_CCR: u64 = 0xAA01;
95    const INJECT: u64 = 0xAA02;
96    const DEDUP: u64 = 0xAA03;
97    const STABLE: u64 = 0xAA04;
98    const STICKY: u64 = 0xAA05;
99
100    #[test]
101    fn tool_not_injected_without_ccr() {
102        let mut doc = json!({"tools": [], "messages": []});
103        assert!(!ensure_tool_present(NO_CCR, &mut doc));
104        assert!(doc["tools"].as_array().unwrap().is_empty());
105    }
106
107    #[test]
108    fn tool_injected_after_ccr_activation() {
109        mark_ccr_active(INJECT);
110        assert!(is_ccr_active(INJECT));
111        let mut doc = json!({"tools": [], "messages": []});
112        assert!(ensure_tool_present(INJECT, &mut doc));
113        assert_eq!(doc["tools"].as_array().unwrap().len(), 1);
114        assert_eq!(doc["tools"][0]["name"], "ctx_expand");
115    }
116
117    #[test]
118    fn tool_not_duplicated() {
119        mark_ccr_active(DEDUP);
120        let mut doc = json!({"tools": [expand_tool_definition()], "messages": []});
121        assert!(!ensure_tool_present(DEDUP, &mut doc));
122        assert_eq!(doc["tools"].as_array().unwrap().len(), 1);
123    }
124
125    #[test]
126    fn tools_array_stays_stable_after_ccr_activation() {
127        mark_ccr_active(STABLE);
128        let existing = json!({"name": "other_tool", "description": "test"});
129
130        let mut doc1 = json!({"tools": [existing.clone()], "messages": []});
131        ensure_tool_present(STABLE, &mut doc1);
132        let snap1 = serde_json::to_string(&doc1["tools"]).unwrap();
133
134        let mut doc2 = json!({"tools": [existing], "messages": []});
135        ensure_tool_present(STABLE, &mut doc2);
136        let snap2 = serde_json::to_string(&doc2["tools"]).unwrap();
137
138        assert_eq!(
139            snap1, snap2,
140            "tool list must be byte-identical across turns"
141        );
142    }
143
144    #[test]
145    fn sticky_survives_turn_without_markers() {
146        mark_ccr_active(STICKY);
147        assert!(is_ccr_active(STICKY));
148        let mut doc = json!({"tools": [], "messages": []});
149        assert!(ensure_tool_present(STICKY, &mut doc));
150    }
151
152    #[test]
153    fn max_tracked_does_not_panic() {
154        for i in 0..(MAX_TRACKED + 100) {
155            mark_ccr_active(i as u64);
156        }
157        assert!(active_sessions().lock().unwrap().len() <= MAX_TRACKED);
158    }
159}
lean_ctx/proxy/sticky_tools.rs

lean_ctx/proxy/
sticky_tools.rs
