Skip to main content

lean_ctx/core/context_kernel/
evidence_wiring.rs

1//! Dispatch-aware evidence recording for MCP tools and proxy requests.
2
3use std::sync::atomic::{AtomicUsize, Ordering};
4
5static TOOL_DISPATCHES: AtomicUsize = AtomicUsize::new(0);
6static PROXY_DISPATCHES: AtomicUsize = AtomicUsize::new(0);
7static TOTAL_TOKENS_SAVED: AtomicUsize = AtomicUsize::new(0);
8static TOTAL_INPUT_TOKENS: AtomicUsize = AtomicUsize::new(0);
9static TOTAL_OUTPUT_TOKENS: AtomicUsize = AtomicUsize::new(0);
10
11/// Cumulative evidence totals observed at dispatch boundaries.
12#[derive(Debug, Clone, Default, serde::Serialize)]
13pub struct DispatchSummary {
14    /// Number of dispatched MCP tool calls.
15    pub tool_dispatches: usize,
16    /// Number of dispatched proxy requests.
17    pub proxy_dispatches: usize,
18    /// Tokens saved across all dispatches.
19    pub total_tokens_saved: usize,
20    /// Input tokens across all dispatches.
21    pub total_input_tokens: usize,
22    /// Output tokens across all dispatches.
23    pub total_output_tokens: usize,
24}
25
26fn record_dispatch(input_tokens: usize, output_tokens: usize, tokens_saved: usize) {
27    TOTAL_INPUT_TOKENS.fetch_add(input_tokens, Ordering::Relaxed);
28    TOTAL_OUTPUT_TOKENS.fetch_add(output_tokens, Ordering::Relaxed);
29    TOTAL_TOKENS_SAVED.fetch_add(tokens_saved, Ordering::Relaxed);
30}
31
32/// Records evidence for one dispatched MCP tool call.
33pub fn record_from_tool_dispatch(
34    tool_name: &str,
35    input_tokens: usize,
36    output_tokens: usize,
37    tokens_saved: usize,
38) {
39    TOOL_DISPATCHES.fetch_add(1, Ordering::Relaxed);
40    record_dispatch(input_tokens, output_tokens, tokens_saved);
41    super::evidence_hook::record_tool_call(tool_name, input_tokens, output_tokens);
42}
43
44/// Records evidence for one proxy-forwarded request.
45pub fn record_from_proxy_dispatch(
46    input_tokens: usize,
47    output_tokens: usize,
48    tokens_saved: usize,
49    model: Option<&str>,
50    provider: Option<&str>,
51) {
52    PROXY_DISPATCHES.fetch_add(1, Ordering::Relaxed);
53    record_dispatch(input_tokens, output_tokens, tokens_saved);
54    super::evidence_hook::record_proxy_call(
55        input_tokens,
56        output_tokens,
57        tokens_saved,
58        model,
59        provider,
60    );
61}
62
63/// Returns cumulative dispatch-boundary evidence totals.
64#[must_use]
65pub fn dispatch_summary() -> DispatchSummary {
66    DispatchSummary {
67        tool_dispatches: TOOL_DISPATCHES.load(Ordering::Relaxed),
68        proxy_dispatches: PROXY_DISPATCHES.load(Ordering::Relaxed),
69        total_tokens_saved: TOTAL_TOKENS_SAVED.load(Ordering::Relaxed),
70        total_input_tokens: TOTAL_INPUT_TOKENS.load(Ordering::Relaxed),
71        total_output_tokens: TOTAL_OUTPUT_TOKENS.load(Ordering::Relaxed),
72    }
73}
74
75/// Clears dispatch counters and delegated evidence state.
76pub fn reset() {
77    TOOL_DISPATCHES.store(0, Ordering::Relaxed);
78    PROXY_DISPATCHES.store(0, Ordering::Relaxed);
79    TOTAL_TOKENS_SAVED.store(0, Ordering::Relaxed);
80    TOTAL_INPUT_TOKENS.store(0, Ordering::Relaxed);
81    TOTAL_OUTPUT_TOKENS.store(0, Ordering::Relaxed);
82    super::evidence_hook::reset();
83}
84
85#[cfg(test)]
86mod tests {
87    use super::{dispatch_summary, record_from_proxy_dispatch, record_from_tool_dispatch, reset};
88    use crate::core::context_kernel::{evidence_hook, kernel_config};
89
90    fn isolated() -> std::sync::MutexGuard<'static, ()> {
91        let guard = kernel_config::KERNEL_TEST_LOCK
92            .lock()
93            .unwrap_or_else(std::sync::PoisonError::into_inner);
94        kernel_config::reset_features();
95        reset();
96        guard
97    }
98
99    #[test]
100    fn tool_dispatch_records() {
101        let _guard = isolated();
102        for _ in 0..3 {
103            record_from_tool_dispatch("ctx_read", 10, 4, 6);
104        }
105        let summary = dispatch_summary();
106        assert_eq!(summary.tool_dispatches, 3);
107        assert_eq!(summary.total_tokens_saved, 18);
108    }
109
110    #[test]
111    fn proxy_dispatch_records() {
112        let _guard = isolated();
113        for _ in 0..2 {
114            record_from_proxy_dispatch(20, 5, 8, Some("gpt-test"), Some("openai"));
115        }
116        let summary = dispatch_summary();
117        assert_eq!(summary.proxy_dispatches, 2);
118        assert_eq!(summary.total_input_tokens, 40);
119    }
120
121    #[test]
122    fn disabled_kernel_no_records() {
123        let _guard = isolated();
124        let mut features = kernel_config::features();
125        features.enabled = false;
126        kernel_config::update_features(features);
127        record_from_tool_dispatch("ctx_read", 10, 4, 6);
128        record_from_proxy_dispatch(20, 5, 8, None, None);
129        assert_eq!(evidence_hook::evidence_report().tool_calls, 0);
130        assert_eq!(evidence_hook::evidence_report().proxy_calls, 0);
131        assert_eq!(dispatch_summary().tool_dispatches, 1);
132        assert_eq!(dispatch_summary().proxy_dispatches, 1);
133    }
134
135    #[test]
136    fn reset_clears_all() {
137        let _guard = isolated();
138        record_from_tool_dispatch("ctx_read", 10, 4, 6);
139        record_from_proxy_dispatch(20, 5, 8, None, None);
140        reset();
141        let summary = dispatch_summary();
142        assert_eq!(summary.tool_dispatches, 0);
143        assert_eq!(summary.proxy_dispatches, 0);
144        assert_eq!(summary.total_tokens_saved, 0);
145        assert_eq!(summary.total_input_tokens, 0);
146        assert_eq!(summary.total_output_tokens, 0);
147    }
148}