Skip to main content

lean_ctx/core/context_kernel/
hotpath_wiring.rs

1//! Unified Context Kernel integration for tool hot-paths.
2
3use super::accounting_fix::{PostDeliveryAccounting, compute_honest_accounting};
4use super::activation::{
5    KernelModeConfig, load_config, should_suppress_in_mode, supplement_budget,
6};
7use super::bridge::{KernelEnrichment, kernel_enrich};
8
9/// Result of the unified kernel integration for a single hot-path call.
10#[derive(Debug, Clone)]
11pub struct KernelIntegration {
12    /// Text to append to the tool output (budget-capped, deduplicated).
13    pub supplement: Option<String>,
14    /// Whether the content should be suppressed (already delivered, enforce mode).
15    pub suppress: bool,
16    /// Honest token accounting for this call.
17    pub accounting: PostDeliveryAccounting,
18    /// Budget tokens consumed by the kernel.
19    pub budget_used: usize,
20}
21
22/// Unified kernel integration for all hot-paths.
23///
24/// Call this instead of manually orchestrating kernel enrichment, activation,
25/// and accounting. The returned supplement is already bounded by the active
26/// kernel budget.
27pub fn kernel_integrate(
28    query: &str,
29    project_root: &str,
30    original_tokens: usize,
31    compressed_tokens: usize,
32) -> KernelIntegration {
33    let config = load_config(project_root);
34    let budget = supplement_budget(&config);
35    let enrichment = kernel_enrich(query, project_root, budget);
36
37    finish_integration(config.mode, original_tokens, compressed_tokens, enrichment)
38}
39
40/// Returns the stable delimiter placed before appended kernel context.
41pub fn format_integration_header() -> &'static str {
42    "\n--- kernel context ---\n"
43}
44
45/// Returns the kernel token overhead recorded for an integration.
46pub fn integration_overhead(integration: &KernelIntegration) -> usize {
47    integration.budget_used
48}
49
50fn finish_integration(
51    mode: KernelModeConfig,
52    original_tokens: usize,
53    compressed_tokens: usize,
54    enrichment: Option<KernelEnrichment>,
55) -> KernelIntegration {
56    let (supplement, budget_used) = match enrichment {
57        Some(enrichment) => {
58            let budget_used = crate::core::tokens::count_tokens(&enrichment.blocks);
59            (Some(enrichment.blocks), budget_used)
60        }
61        None => (None, 0),
62    };
63    let accounting = compute_honest_accounting(original_tokens, compressed_tokens, budget_used, 0);
64
65    KernelIntegration {
66        supplement,
67        suppress: should_suppress_in_mode(mode),
68        accounting,
69        budget_used,
70    }
71}
72
73/// Bridge: integrate kernel context for an MCP tool call.
74///
75/// Unlike [`kernel_integrate`], which takes raw token counts, this function
76/// takes headers and determines the optimal integration strategy based on
77/// the client's coverage class and profile.
78pub fn integrate_for_mcp(
79    query: &str,
80    project_root: &str,
81    headers: &[(String, String)],
82    original_tokens: usize,
83    compressed_tokens: usize,
84) -> KernelIntegration {
85    let profile = super::client_profile::detect_from_headers(headers);
86    let coverage = profile.coverage;
87
88    if !super::coverage_class::is_addressable(coverage) {
89        return KernelIntegration {
90            supplement: None,
91            suppress: false,
92            accounting: compute_honest_accounting(original_tokens, compressed_tokens, 0, 0),
93            budget_used: 0,
94        };
95    }
96
97    kernel_integrate(query, project_root, original_tokens, compressed_tokens)
98}
99
100/// Returns the coverage-aware kernel budget for an MCP request.
101pub fn mcp_kernel_budget(headers: &[(String, String)]) -> usize {
102    let profile = super::client_profile::detect_from_headers(headers);
103    let broker = super::context_broker::ContextBroker::new(profile);
104    broker.compute_budget().kernel_tokens
105}
106
107#[cfg(test)]
108mod tests {
109    use std::collections::HashMap;
110    use std::fs;
111
112    use super::{
113        KernelIntegration, finish_integration, integrate_for_mcp, integration_overhead,
114        kernel_integrate, mcp_kernel_budget,
115    };
116    use crate::core::context_kernel::accounting_fix::detect_negative_savings;
117    use crate::core::context_kernel::activation::KernelModeConfig;
118    use crate::core::context_kernel::bridge::{KernelEnrichment, KernelVerdict};
119    use crate::core::context_kernel::enforce::KernelMode;
120    use crate::core::context_kernel::types::{ContextPlanV1, PlanBudget};
121
122    fn enrichment(blocks: String) -> KernelEnrichment {
123        let budget_used = crate::core::tokens::count_tokens(&blocks);
124        KernelEnrichment {
125            plan: ContextPlanV1 {
126                plan_id: "test-plan".to_owned(),
127                intent: "test".to_owned(),
128                budget: PlanBudget {
129                    total_tokens: 150,
130                    used_tokens: budget_used,
131                    remaining_tokens: 150_usize.saturating_sub(budget_used),
132                },
133                selected: Vec::new(),
134                excluded: Vec::new(),
135                deferred: Vec::new(),
136                provider_stats: HashMap::new(),
137            },
138            blocks: blocks.clone(),
139            verdict: KernelVerdict {
140                supplement: Some(blocks),
141                suppress: Vec::new(),
142                budget_used,
143            },
144            enforced_mode: KernelMode::Shadow,
145        }
146    }
147
148    fn integration(
149        mode: KernelModeConfig,
150        original_tokens: usize,
151        compressed_tokens: usize,
152        blocks: Option<String>,
153    ) -> KernelIntegration {
154        finish_integration(
155            mode,
156            original_tokens,
157            compressed_tokens,
158            blocks.map(enrichment),
159        )
160    }
161
162    #[test]
163    fn integrate_caps_budget_at_150() {
164        let root = tempfile::tempdir().expect("temporary project root");
165        fs::write(
166            root.path().join(".lean-ctx.toml"),
167            "[kernel]\nmax_supplement_tokens = 500\n",
168        )
169        .expect("kernel config should be writable");
170
171        let result = kernel_integrate(
172            "query with no project candidates",
173            root.path().to_str().expect("UTF-8 project path"),
174            1_000,
175            500,
176        );
177
178        assert!(result.budget_used <= 150);
179    }
180
181    #[test]
182    fn integrate_accounting_honest() {
183        let result = integration(
184            KernelModeConfig::Shadow,
185            100,
186            40,
187            Some("kernel context".to_owned()),
188        );
189
190        assert_eq!(result.accounting.kernel_overhead_tokens, result.budget_used);
191        assert_eq!(
192            result.accounting.delivered_tokens,
193            40 + integration_overhead(&result)
194        );
195    }
196
197    #[test]
198    fn integrate_shadow_no_suppress() {
199        let result = integration(KernelModeConfig::Shadow, 100, 50, None);
200
201        assert!(!result.suppress);
202    }
203
204    #[test]
205    fn integrate_no_enrichment_zero_overhead() {
206        let result = integration(KernelModeConfig::Enforce, 100, 50, None);
207
208        assert_eq!(result.budget_used, 0);
209        assert_eq!(result.accounting.kernel_overhead_tokens, 0);
210        assert!(result.supplement.is_none());
211    }
212
213    #[test]
214    fn integrate_negative_savings_detected() {
215        let result = integration(
216            KernelModeConfig::Shadow,
217            10,
218            10,
219            Some("kernel overhead exceeds original input size".repeat(20)),
220        );
221
222        assert!(detect_negative_savings(&result.accounting));
223    }
224
225    #[test]
226    fn mcp_unmanaged_no_enrichment() {
227        let root = tempfile::tempdir().expect("temporary project root");
228        let headers = vec![("x-coverage-class".to_owned(), "unmanaged".to_owned())];
229
230        let result = integrate_for_mcp(
231            "query with no project candidates",
232            root.path().to_str().expect("UTF-8 project path"),
233            &headers,
234            100,
235            50,
236        );
237
238        assert_eq!(result.budget_used, 0);
239        assert!(result.supplement.is_none());
240    }
241
242    #[test]
243    fn mcp_full_inline_enriches() {
244        let root = tempfile::tempdir().expect("temporary project root");
245        let project_root = root.path().to_str().expect("UTF-8 project path");
246        let headers = vec![("x-coverage-class".to_owned(), "full_inline".to_owned())];
247        let expected = kernel_integrate("bridge query", project_root, 100, 50);
248
249        let result = integrate_for_mcp("bridge query", project_root, &headers, 100, 50);
250
251        assert_eq!(result.supplement, expected.supplement);
252        assert_eq!(result.budget_used, expected.budget_used);
253        assert_eq!(result.suppress, expected.suppress);
254    }
255
256    #[test]
257    fn mcp_budget_from_profile() {
258        let headers = vec![("x-context-window".to_owned(), "64000".to_owned())];
259
260        assert_eq!(mcp_kernel_budget(&headers), 6_400);
261    }
262}