lean_ctx/core/context_kernel/
schema_wiring.rs

1//! Wiring between MCP tool registries and kernel schema optimization.
2
3use std::sync::{Mutex, MutexGuard, OnceLock};
4
5use super::coverage_class;
6use super::kernel_config;
7use super::mcp_coverage;
8use super::mcp_schema_opt::{self, SchemaEntry};
9
10const ESTIMATED_TOKENS_PER_PARAMETER: usize = 8;
11
12/// Result of optimizing the MCP tool list.
13#[derive(Debug, Clone)]
14pub struct OptimizedToolList {
15    /// Optimized tool entries (name, description, parameter count).
16    pub tools: Vec<(String, String, usize)>,
17    /// Tokens before optimization.
18    pub tokens_before: usize,
19    /// Tokens after optimization.
20    pub tokens_after: usize,
21    /// Number of tools dropped.
22    pub dropped_count: usize,
23    /// Number of descriptions compressed.
24    pub compressed_count: usize,
25    /// Whether optimization was applied (false if disabled).
26    pub optimized: bool,
27}
28
29/// Cumulative process-wide schema optimization savings.
30#[derive(Debug, Clone, Copy, Default)]
31pub struct SchemaSavings {
32    /// Number of tool-list optimizations applied.
33    pub optimizations_applied: usize,
34    /// Total estimated tokens removed by optimization.
35    pub total_tokens_saved: usize,
36    /// Saved tokens divided by tokens seen before optimization.
37    pub avg_compression_ratio: f64,
38}
39
40#[derive(Debug, Default)]
41struct SavingsState {
42    optimizations_applied: usize,
43    total_tokens_before: usize,
44    total_tokens_saved: usize,
45}
46
47static SAVINGS: OnceLock<Mutex<SavingsState>> = OnceLock::new();
48
49fn savings_guard() -> MutexGuard<'static, SavingsState> {
50    SAVINGS
51        .get_or_init(|| Mutex::new(SavingsState::default()))
52        .lock()
53        .unwrap_or_else(std::sync::PoisonError::into_inner)
54}
55
56fn estimated_schema_tokens(name: &str, description: &str, param_count: usize) -> usize {
57    mcp_schema_opt::estimate_tokens(name)
58        .saturating_add(mcp_schema_opt::estimate_tokens(description))
59        .saturating_add(param_count.saturating_mul(ESTIMATED_TOKENS_PER_PARAMETER))
60}
61
62fn unchanged(tools: &[(String, String, usize)]) -> OptimizedToolList {
63    let tokens = tools.iter().fold(0usize, |total, (name, desc, count)| {
64        total.saturating_add(estimated_schema_tokens(name, desc, *count))
65    });
66    OptimizedToolList {
67        tools: tools.to_vec(),
68        tokens_before: tokens,
69        tokens_after: tokens,
70        dropped_count: 0,
71        compressed_count: 0,
72        optimized: false,
73    }
74}
75
76fn optimize_tool_list_with_feature(
77    tools: &[(String, String, usize)],
78    client_name: &str,
79    enabled: bool,
80) -> OptimizedToolList {
81    if !enabled {
82        return unchanged(tools);
83    }
84
85    let coverage = mcp_coverage::detect_mcp_coverage(client_name, false, false);
86    let budget = mcp_schema_opt::budget_for_coverage(coverage);
87    let entries = tools
88        .iter()
89        .map(|(name, description, param_count)| SchemaEntry {
90            name: name.clone(),
91            description: description.clone(),
92            param_count: *param_count,
93            estimated_tokens: estimated_schema_tokens(name, description, *param_count),
94            essential: false,
95        })
96        .collect::<Vec<_>>();
97    let result = mcp_schema_opt::optimize_schemas(&entries, &budget);
98    let tokens_saved = result.tokens_before.saturating_sub(result.tokens_after);
99
100    {
101        let mut savings = savings_guard();
102        savings.optimizations_applied = savings.optimizations_applied.saturating_add(1);
103        savings.total_tokens_before = savings
104            .total_tokens_before
105            .saturating_add(result.tokens_before);
106        savings.total_tokens_saved = savings.total_tokens_saved.saturating_add(tokens_saved);
107    }
108
109    OptimizedToolList {
110        tools: result
111            .entries
112            .into_iter()
113            .map(|entry| (entry.name, entry.description, entry.param_count))
114            .collect(),
115        tokens_before: result.tokens_before,
116        tokens_after: result.tokens_after,
117        dropped_count: result.dropped_count,
118        compressed_count: result.compressed_count,
119        optimized: true,
120    }
121}
122
123/// Optimizes MCP tool descriptions and count for the client's coverage budget.
124#[must_use]
125pub fn optimize_tool_list(
126    tools: &[(String, String, usize)],
127    client_name: &str,
128) -> OptimizedToolList {
129    optimize_tool_list_with_feature(
130        tools,
131        client_name,
132        kernel_config::features().schema_optimization,
133    )
134}
135
136/// Returns cumulative savings produced by enabled schema optimization.
137#[must_use]
138pub fn schema_savings() -> SchemaSavings {
139    let savings = savings_guard();
140    let avg_compression_ratio = if savings.total_tokens_before == 0 {
141        0.0
142    } else {
143        savings.total_tokens_saved as f64 / savings.total_tokens_before as f64
144    };
145    SchemaSavings {
146        optimizations_applied: savings.optimizations_applied,
147        total_tokens_saved: savings.total_tokens_saved,
148        avg_compression_ratio,
149    }
150}
151
152fn should_optimize_with_feature(client_name: &str, enabled: bool) -> bool {
153    let coverage = mcp_coverage::detect_mcp_coverage(client_name, false, false);
154    enabled && coverage_class::is_addressable(coverage)
155}
156
157/// Returns whether schema optimization is enabled and the client is addressable.
158#[must_use]
159pub fn should_optimize(client_name: &str) -> bool {
160    should_optimize_with_feature(client_name, kernel_config::features().schema_optimization)
161}
162
163/// Clears cumulative schema optimization metrics.
164pub fn reset_schema_state() {
165    *savings_guard() = SavingsState::default();
166}
167
168#[cfg(test)]
169mod tests {
170    use std::sync::MutexGuard;
171
172    use super::{
173        optimize_tool_list_with_feature, reset_schema_state, schema_savings,
174        should_optimize_with_feature,
175    };
176
177    fn setup() -> MutexGuard<'static, ()> {
178        let guard = crate::core::context_kernel::kernel_config::KERNEL_TEST_LOCK
179            .lock()
180            .unwrap_or_else(std::sync::PoisonError::into_inner);
181        reset_schema_state();
182        guard
183    }
184
185    fn tools(description_len: usize) -> Vec<(String, String, usize)> {
186        (0..20)
187            .map(|index| (format!("tool_{index}"), "x".repeat(description_len), 3))
188            .collect()
189    }
190
191    #[test]
192    fn optimize_reduces_tokens() {
193        let _guard = setup();
194        let optimized = optimize_tool_list_with_feature(&tools(4_000), "cursor", true);
195        assert!(optimized.tokens_after < optimized.tokens_before);
196        assert!(optimized.compressed_count > 0);
197    }
198
199    #[test]
200    fn optimize_disabled_returns_unchanged() {
201        let _guard = setup();
202        let original = tools(100);
203        let optimized = optimize_tool_list_with_feature(&original, "cursor", false);
204        assert_eq!(optimized.tools, original);
205        assert!(!optimized.optimized);
206    }
207
208    #[test]
209    fn cursor_gets_full_budget() {
210        let _guard = setup();
211        let optimized = optimize_tool_list_with_feature(&tools(1_000), "cursor", true);
212        assert_eq!(optimized.tools.len(), 20);
213    }
214
215    #[test]
216    fn unknown_client_gets_small_budget() {
217        let _guard = setup();
218        let optimized = optimize_tool_list_with_feature(&tools(1_000), "random", true);
219        assert!(optimized.tools.len() < 20);
220    }
221
222    #[test]
223    fn should_optimize_respects_config() {
224        let _guard = setup();
225        assert!(!should_optimize_with_feature("cursor", false));
226        assert!(should_optimize_with_feature("cursor", true));
227        assert!(!should_optimize_with_feature("random", true));
228    }
229
230    #[test]
231    fn savings_accumulate() {
232        let _guard = setup();
233        let _ = optimize_tool_list_with_feature(&tools(4_000), "cursor", true);
234        let savings = schema_savings();
235        assert_eq!(savings.optimizations_applied, 1);
236        assert!(savings.total_tokens_saved > 0);
237        assert!(savings.avg_compression_ratio > 0.0);
238    }
239}
lean_ctx/core/context_kernel/schema_wiring.rs

lean_ctx/core/context_kernel/
schema_wiring.rs
