lean_ctx/core/
proxy_value_metrics.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Per-session value metrics for proxied completions.
3//!
4//! The proxy records prompt-compression measurements and a per-request input
5//! cost estimate for each completion.
6
7use std::sync::{Mutex, PoisonError};
8
9use lean_ctx_protocol::AcceptanceState;
10
11/// Aggregate value metrics for the proxy process's current session.
12#[derive(Clone, Debug, Default, PartialEq, Eq)]
13pub struct ProxyValueMetrics {
14    pub request_count: u64,
15    pub total_tokens_pruned: u64,
16    pub total_original_tokens: u64,
17    pub last_task_class: Option<String>,
18    pub cost_micros_estimate: u64,
19    pub session_cpao_micros: Option<u64>,
20    pub accepted_outcome_count: u64,
21    pub unknown_outcome_count: u64,
22}
23
24static PROXY_VALUE_METRICS: Mutex<ProxyValueMetrics> = Mutex::new(ProxyValueMetrics {
25    request_count: 0,
26    total_tokens_pruned: 0,
27    total_original_tokens: 0,
28    last_task_class: None,
29    cost_micros_estimate: 0,
30    session_cpao_micros: None,
31    accepted_outcome_count: 0,
32    unknown_outcome_count: 0,
33});
34
35/// Records a completed proxy request in the current process session.
36pub fn record_completion(
37    tokens_pruned: usize,
38    original_tokens: usize,
39    task_class: &str,
40    acceptance: AcceptanceState,
41) {
42    let mut metrics = PROXY_VALUE_METRICS
43        .lock()
44        .unwrap_or_else(PoisonError::into_inner);
45    metrics.request_count = metrics.request_count.saturating_add(1);
46    metrics.total_tokens_pruned = metrics
47        .total_tokens_pruned
48        .saturating_add(tokens_pruned as u64);
49    metrics.total_original_tokens = metrics
50        .total_original_tokens
51        .saturating_add(original_tokens as u64);
52    metrics.last_task_class = Some(task_class.to_owned());
53    let cost_micros = (original_tokens as u64).saturating_mul(3) / 1_000;
54    metrics.cost_micros_estimate = metrics.cost_micros_estimate.saturating_add(cost_micros);
55    match acceptance {
56        AcceptanceState::Accepted => {
57            metrics.accepted_outcome_count = metrics.accepted_outcome_count.saturating_add(1);
58            metrics.session_cpao_micros =
59                Some(metrics.cost_micros_estimate / metrics.accepted_outcome_count);
60        }
61        AcceptanceState::Unknown => {
62            metrics.unknown_outcome_count = metrics.unknown_outcome_count.saturating_add(1);
63            metrics.session_cpao_micros = None;
64        }
65        AcceptanceState::Rejected => metrics.session_cpao_micros = None,
66    }
67    tracing::debug!(
68        task_class,
69        tokens_pruned,
70        original_tokens,
71        "recorded proxy value metrics"
72    );
73}
74
75/// Clears all value metrics for the current proxy session.
76pub fn reset_session() {
77    *PROXY_VALUE_METRICS
78        .lock()
79        .unwrap_or_else(PoisonError::into_inner) = ProxyValueMetrics::default();
80}
81
82/// Returns a snapshot of value metrics accumulated for the current session.
83pub fn session_metrics() -> ProxyValueMetrics {
84    PROXY_VALUE_METRICS
85        .lock()
86        .unwrap_or_else(PoisonError::into_inner)
87        .clone()
88}
89
90/// Returns the fraction of original tokens removed during the current session.
91pub fn compression_ratio() -> f64 {
92    let metrics = session_metrics();
93    if metrics.total_original_tokens == 0 {
94        return 0.0;
95    }
96    metrics.total_tokens_pruned as f64 / metrics.total_original_tokens as f64
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    static TEST_LOCK: Mutex<()> = Mutex::new(());
104
105    #[test]
106    fn record_completion_accumulates_session_totals() {
107        let _guard = TEST_LOCK.lock().unwrap_or_else(PoisonError::into_inner);
108        reset_session();
109
110        record_completion(20, 100, "coding_fix", AcceptanceState::Unknown);
111        record_completion(30, 200, "coding_new", AcceptanceState::Unknown);
112
113        let metrics = session_metrics();
114        assert_eq!(metrics.request_count, 2);
115        assert_eq!(metrics.total_tokens_pruned, 50);
116        assert_eq!(metrics.total_original_tokens, 300);
117        assert_eq!(metrics.last_task_class.as_deref(), Some("coding_new"));
118        assert_eq!(metrics.cost_micros_estimate, 0);
119        assert_eq!(metrics.session_cpao_micros, None);
120        assert_eq!(metrics.unknown_outcome_count, 2);
121    }
122
123    #[test]
124    fn session_metrics_returns_a_clone() {
125        let _guard = TEST_LOCK.lock().unwrap_or_else(PoisonError::into_inner);
126        reset_session();
127        record_completion(20, 100, "coding_fix", AcceptanceState::Unknown);
128
129        let mut snapshot = session_metrics();
130        snapshot.request_count = 0;
131
132        assert_eq!(snapshot.request_count, 0);
133        assert_eq!(session_metrics().request_count, 1);
134    }
135
136    #[test]
137    fn compression_ratio_is_zero_without_original_tokens() {
138        let _guard = TEST_LOCK.lock().unwrap_or_else(PoisonError::into_inner);
139        reset_session();
140
141        assert_eq!(compression_ratio(), 0.0);
142    }
143
144    #[test]
145    fn compression_ratio_uses_session_totals() {
146        let _guard = TEST_LOCK.lock().unwrap_or_else(PoisonError::into_inner);
147        reset_session();
148        record_completion(25, 100, "coding_fix", AcceptanceState::Unknown);
149        record_completion(50, 200, "coding_new", AcceptanceState::Unknown);
150
151        assert_eq!(compression_ratio(), 0.25);
152    }
153
154    #[test]
155    fn record_completion_estimates_session_cpao() {
156        let _guard = TEST_LOCK.lock().unwrap_or_else(PoisonError::into_inner);
157        reset_session();
158
159        record_completion(0, 1_000, "coding_fix", AcceptanceState::Accepted);
160        record_completion(0, 2_000, "coding_new", AcceptanceState::Accepted);
161
162        let metrics = session_metrics();
163        assert_eq!(metrics.cost_micros_estimate, 9);
164        assert_eq!(metrics.session_cpao_micros, Some(4));
165    }
166
167    #[test]
168    fn reset_session_clears_all_metrics() {
169        let _guard = TEST_LOCK.lock().unwrap_or_else(PoisonError::into_inner);
170        record_completion(20, 1_000, "coding_fix", AcceptanceState::Unknown);
171
172        reset_session();
173
174        assert_eq!(session_metrics(), ProxyValueMetrics::default());
175    }
176}
lean_ctx/core/proxy_value_metrics.rs

lean_ctx/core/
proxy_value_metrics.rs
