lean_ctx_server/server/
schema_hook.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Server bridge for Context Kernel list-tools schema optimization.
3
4use std::borrow::Cow;
5
6use lean_ctx::core::context_kernel::list_tools_opt;
7
8/// Optimizes tool descriptions when schema optimization supports the client.
9#[must_use]
10pub fn optimize_tools(tools: Vec<rmcp::model::Tool>, client_name: &str) -> Vec<rmcp::model::Tool> {
11    if !list_tools_opt::should_optimize_for_client(client_name) {
12        return tools;
13    }
14
15    let tuples = tools
16        .iter()
17        .map(|tool| {
18            let parameter_count = tool
19                .input_schema
20                .get("properties")
21                .and_then(serde_json::Value::as_object)
22                .map_or(0, serde_json::Map::len);
23            (
24                tool.name.to_string(),
25                tool.description.as_deref().unwrap_or("").to_owned(),
26                parameter_count,
27            )
28        })
29        .collect();
30    // The optimizer re-sorts its entries and may drop some to meet the budget,
31    // so its descriptions are matched back by tool name; a tool it dropped
32    // keeps its own description and stays in the list.
33    let mut optimized: std::collections::HashMap<String, String> =
34        list_tools_opt::optimize_descriptions(tuples, client_name)
35            .into_iter()
36            .map(|(name, description, _)| (name, description))
37            .collect();
38
39    tools
40        .into_iter()
41        .map(|mut tool| {
42            if let Some(description) = optimized.remove(tool.name.as_ref()) {
43                tool.description = Some(Cow::Owned(description));
44            }
45            tool
46        })
47        .collect()
48}
49
50/// Returns cumulative list-tools schema optimization statistics.
51#[must_use]
52pub fn summary() -> list_tools_opt::SchemaOptSummary {
53    list_tools_opt::schema_opt_summary()
54}
55
56#[cfg(test)]
57mod tests {
58    use std::sync::MutexGuard;
59
60    use super::list_tools_opt;
61    use lean_ctx::core::context_kernel::kernel_config::{self, KERNEL_TEST_LOCK, KernelFeatures};
62
63    fn setup() -> MutexGuard<'static, ()> {
64        let guard = KERNEL_TEST_LOCK
65            .lock()
66            .unwrap_or_else(std::sync::PoisonError::into_inner);
67        kernel_config::reset_features();
68        list_tools_opt::reset();
69        guard
70    }
71
72    #[test]
73    fn disabled_returns_same_count() {
74        let _guard = setup();
75        let features = KernelFeatures {
76            enabled: false,
77            ..KernelFeatures::default()
78        };
79        kernel_config::update_features(features);
80        let tools = vec![("tool".to_owned(), "description".to_owned(), 1)];
81
82        assert_eq!(
83            list_tools_opt::optimize_descriptions(tools, "cursor").len(),
84            1
85        );
86    }
87
88    #[test]
89    fn summary_works() {
90        let _guard = setup();
91        let summary = list_tools_opt::schema_opt_summary();
92
93        assert_eq!(summary.optimizations_applied, 0);
94        assert_eq!(summary.total_tokens_saved, 0);
95        assert_eq!(summary.avg_reduction_percent, 0.0);
96    }
97
98    /// F7: the optimizer sorts by size and drops entries over budget; every
99    /// tool must survive and keep a description that belongs to it.
100    #[test]
101    fn descriptions_are_paired_by_name_and_no_tool_is_lost() {
102        let _guard = setup();
103        let tools: Vec<rmcp::model::Tool> = (0..40)
104            .map(|i| {
105                // Descending size, so the optimizer's ascending sort reverses
106                // the order, and large enough that the budget drops some.
107                let description = format!("tool-{i}: {}", "x".repeat(8_000 - i * 100));
108                rmcp::model::Tool::new(
109                    format!("tool_{i}"),
110                    description,
111                    std::sync::Arc::new(serde_json::Map::new()),
112                )
113            })
114            .collect();
115
116        let optimized = super::optimize_tools(tools, "cursor");
117
118        assert_eq!(optimized.len(), 40, "no tool may be dropped from the list");
119        for (i, tool) in optimized.iter().enumerate() {
120            assert_eq!(tool.name, format!("tool_{i}"));
121            let description = tool.description.as_deref().unwrap_or("");
122            assert!(
123                description.starts_with(&format!("tool-{i}: ")),
124                "{} got a foreign description: {}",
125                tool.name,
126                &description[..description.len().min(20)]
127            );
128        }
129    }
130
131    #[test]
132    fn optimization_flag_check() {
133        let _guard = setup();
134
135        assert!(list_tools_opt::should_optimize_for_client("cursor"));
136        assert!(!list_tools_opt::should_optimize_for_client(
137            "unknown-client"
138        ));
139    }
140}
lean_ctx_server/server/schema_hook.rs

lean_ctx_server/server/
schema_hook.rs
