Skip to main content

lean_ctx_proxy/
break_even.rs

1// SPDX-License-Identifier: Apache-2.0
2use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
3
4/// Tracks whether MCP schema overhead is justified by proxy compression savings.
5pub struct BreakEvenCalculator {
6    proxy_tokens_saved: AtomicI64,
7    schema_overhead_per_turn: u64,
8    turn_count: AtomicU64,
9    has_used_ctx_tools: AtomicBool,
10}
11
12/// Snapshot of a session's MCP break-even state.
13#[derive(Debug, Clone)]
14pub struct BreakEvenSummary {
15    pub proxy_savings: i64,
16    pub estimated_mcp_overhead: i64,
17    pub net: i64,
18    pub mcp_recommended: bool,
19    pub reason: &'static str,
20}
21
22impl BreakEvenCalculator {
23    /// Create with the estimated schema overhead per turn (in tokens).
24    /// Typical values: 1500 (minimal profile) to 3500 (full profile).
25    pub fn new(schema_overhead_per_turn: u64) -> Self {
26        Self {
27            proxy_tokens_saved: AtomicI64::new(0),
28            schema_overhead_per_turn,
29            turn_count: AtomicU64::new(0),
30            has_used_ctx_tools: AtomicBool::new(false),
31        }
32    }
33
34    /// Record tokens saved by proxy-level compression this turn.
35    pub fn record_proxy_savings(&self, tokens: i64) {
36        self.proxy_tokens_saved.fetch_add(tokens, Ordering::Relaxed);
37    }
38
39    /// Advance the turn counter. Call once per API request.
40    pub fn record_turn(&self) {
41        self.turn_count.fetch_add(1, Ordering::Relaxed);
42    }
43
44    /// Mark that the agent has used a ctx_* MCP tool in this session.
45    /// Once set, MCP is never disabled mid-session.
46    pub fn mark_ctx_tool_used(&self) {
47        self.has_used_ctx_tools.store(true, Ordering::Relaxed);
48    }
49
50    /// Whether MCP tools should be enabled for this session.
51    pub fn should_enable_mcp(&self) -> bool {
52        if self.has_used_ctx_tools.load(Ordering::Relaxed) {
53            return true;
54        }
55
56        let turns = self.turn_count.load(Ordering::Relaxed);
57        if turns < 2 {
58            return false;
59        }
60
61        let savings = self.proxy_tokens_saved.load(Ordering::Relaxed);
62        let overhead = self.schema_overhead_per_turn as i64 * turns as i64;
63        savings > overhead
64    }
65
66    /// Snapshot of the current break-even state.
67    pub fn summary(&self) -> BreakEvenSummary {
68        let turns = self.turn_count.load(Ordering::Relaxed);
69        let savings = self.proxy_tokens_saved.load(Ordering::Relaxed);
70        let overhead = self.schema_overhead_per_turn as i64 * turns as i64;
71        let net = savings - overhead;
72        let mcp_recommended = self.should_enable_mcp();
73
74        let reason = if self.has_used_ctx_tools.load(Ordering::Relaxed) {
75            "ctx_tools used — MCP locked on"
76        } else if turns < 2 {
77            "too few turns — proxy-only"
78        } else if mcp_recommended {
79            "savings exceed schema overhead"
80        } else {
81            "schema overhead exceeds savings"
82        };
83
84        BreakEvenSummary {
85            proxy_savings: savings,
86            estimated_mcp_overhead: overhead,
87            net,
88            mcp_recommended,
89            reason,
90        }
91    }
92
93    /// Current turn count.
94    pub fn turn_count(&self) -> u64 {
95        self.turn_count.load(Ordering::Relaxed)
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::BreakEvenCalculator;
102
103    #[test]
104    fn new_calculator_starts_proxy_only() {
105        let calc = BreakEvenCalculator::new(3500);
106        assert!(!calc.should_enable_mcp());
107        assert_eq!(calc.turn_count(), 0);
108    }
109
110    #[test]
111    fn first_two_turns_always_proxy_only() {
112        let calc = BreakEvenCalculator::new(1500);
113        calc.record_proxy_savings(50000);
114        calc.record_turn();
115        assert!(!calc.should_enable_mcp());
116    }
117
118    #[test]
119    fn high_savings_enables_mcp() {
120        let calc = BreakEvenCalculator::new(1500);
121        for _ in 0..5 {
122            calc.record_proxy_savings(5000);
123            calc.record_turn();
124        }
125        assert!(calc.should_enable_mcp());
126    }
127
128    #[test]
129    fn low_savings_stays_proxy_only() {
130        let calc = BreakEvenCalculator::new(3500);
131        for _ in 0..3 {
132            calc.record_proxy_savings(500);
133            calc.record_turn();
134        }
135        assert!(!calc.should_enable_mcp());
136    }
137
138    #[test]
139    fn ctx_tool_usage_locks_mcp_on() {
140        let calc = BreakEvenCalculator::new(3500);
141        calc.mark_ctx_tool_used();
142        assert!(calc.should_enable_mcp());
143        assert!(calc.should_enable_mcp());
144    }
145
146    #[test]
147    fn summary_reflects_state() {
148        let calc = BreakEvenCalculator::new(2000);
149        calc.record_turn();
150        calc.record_turn();
151        calc.record_proxy_savings(3000);
152        let summary = calc.summary();
153        assert_eq!(summary.proxy_savings, 3000);
154        assert_eq!(summary.estimated_mcp_overhead, 4000);
155        assert_eq!(summary.net, -1000);
156        assert!(!summary.mcp_recommended);
157    }
158
159    #[test]
160    fn summary_reason_for_ctx_tools() {
161        let calc = BreakEvenCalculator::new(2000);
162        calc.mark_ctx_tool_used();
163        let summary = calc.summary();
164        assert_eq!(summary.reason, "ctx_tools used — MCP locked on");
165        assert!(summary.mcp_recommended);
166    }
167}