Record Guardian v2 classification token usage (#39303)

## What changed

- Pass session-attributed extension metrics to the Guardian v2 sampler.
- Record `codex.guardian_v2.classification.token_usage` histograms for total, input, cached input, cache-write input, non-cached input, output, and reasoning output tokens.
- Preserve token accounting when streamed JSON produces an early score by recording usage from the terminal response in the background.

## Testing

- Cover token metrics through the Guardian v2 extension and the sampler's early-score path.

GitOrigin-RevId: 04fadc07f3b9cbe83f9fa3737536860b5cf5b4be

codex-rs/ext/guardian-v2/src/extension.rs

@@ -283,6 +283,7 @@ impl ThreadLifecycleContributor<Config> for GuardianV2Extension {
                     .map(|originator| originator.0.clone()),
                 service_tier: input.config.service_tier.clone(),
                 luna_compaction_hash,
+                metrics: input.extension_metrics.clone(),
             })
             .await;
 

codex-rs/ext/guardian-v2/src/extension_tests.rs

@@ -58,6 +58,7 @@ use super::TOOL_CALL_LAG_METRIC;
 use super::encrypted_parent_compaction;
 use crate::config::DEFAULT_MODEL_CONTEXT_ITEM_TOKENS;
 use crate::config::DEFAULT_PARENT_COMPACTION_TOKENS;
+use crate::sampler::CLASSIFICATION_TOKEN_USAGE_METRIC;
 use crate::sampler::MODEL;
 
 const TEST_GUARDIAN_POLICY: &str =
@@ -441,9 +442,17 @@ async fn sample_configured_conversation_history(
         builder
     };
     let test = builder.build_with_auto_env(&thread_server).await?;
+    let mut completed = ev_completed("response-1");
+    completed["response"]["usage"] = json!({
+        "input_tokens": 120,
+        "input_tokens_details": {"cached_tokens": 40, "cache_write_tokens": 20},
+        "output_tokens": 30,
+        "output_tokens_details": {"reasoning_tokens": 10},
+        "total_tokens": 150,
+    });
     let events = vec![
         ev_assistant_message("sample", r#"{"scores":{"action_risk":0.8}}"#),
-        ev_completed("response-1"),
+        completed,
     ];
     let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
     let provider_info = ModelProviderInfo::create_openai_provider(Some(format!(
@@ -675,7 +684,7 @@ max_recent_non_user_entries = 8
             .latest_scored_tool_call
             .load(Ordering::Acquire)
             == 0
-            || metrics.0.lock().unwrap().len() < 2
+            || metrics.0.lock().unwrap().len() < 9
         {
             tokio::task::yield_now().await;
         }
@@ -785,7 +794,7 @@ max_recent_non_user_entries = 8
     );
 
     let samples = initial_metrics.0.lock().unwrap();
-    let classification_duration_ms = match &samples[1] {
+    let classification_duration_ms = match &samples[8] {
         RecordedMetric::Histogram(name, duration_ms, _)
             if name == CLASSIFICATION_DURATION_METRIC =>
         {
@@ -795,7 +804,24 @@ max_recent_non_user_entries = 8
     };
     assert_eq!(
         *samples,
-        vec![
+        [
+            ("total", 150),
+            ("input", 120),
+            ("cached_input", 40),
+            ("cache_write_input", 20),
+            ("non_cached_input", 80),
+            ("output", 30),
+            ("reasoning_output", 10),
+        ]
+        .into_iter()
+        .map(|(token_type, value)| {
+            RecordedMetric::Histogram(
+                CLASSIFICATION_TOKEN_USAGE_METRIC.to_owned(),
+                value,
+                vec![("token_type".to_owned(), token_type.to_owned())],
+            )
+        })
+        .chain([
             RecordedMetric::Counter(
                 CLASSIFICATION_METRIC.to_owned(),
                 1,
@@ -809,7 +835,8 @@ max_recent_non_user_entries = 8
             RecordedMetric::Histogram(TOOL_CALL_LAG_METRIC.to_owned(), 0, vec![]),
             RecordedMetric::Histogram(TOOL_CALL_LAG_METRIC.to_owned(), 0, vec![]),
             RecordedMetric::Histogram(TOOL_CALL_LAG_METRIC.to_owned(), 2, vec![]),
-        ]
+        ])
+        .collect::<Vec<_>>()
     );
     let metrics = thread_store.get::<RecordingMetrics>().unwrap();
     assert_eq!(

codex-rs/ext/guardian-v2/src/sampler.rs

@@ -18,6 +18,7 @@ use codex_api::ResponsesWsRequest;
 use codex_api::TransportError;
 use codex_api::build_session_headers;
 use codex_api::create_text_param_for_request;
+use codex_extension_api::ExtensionMetrics;
 use codex_http_client::HttpClientFactory;
 use codex_login::AgentIdentityAuthPolicy;
 use codex_login::CodexAuth;
@@ -33,6 +34,7 @@ use codex_protocol::models::ContentItem;
 use codex_protocol::models::ResponseItem;
 use codex_protocol::openai_models::ReasoningEffort;
 use codex_protocol::protocol::SessionSource;
+use codex_protocol::protocol::TokenUsage;
 use http::HeaderValue;
 use http::StatusCode;
 use serde_json::Value;
@@ -43,6 +45,8 @@ use tokio::sync::Semaphore;
 use tokio::sync::oneshot;
 
 pub(crate) const MODEL: &str = "gpt-5.6-luna";
+pub(crate) const CLASSIFICATION_TOKEN_USAGE_METRIC: &str =
+    "codex.guardian_v2.classification.token_usage";
 const MAX_OUTPUT_BYTES: usize = 8 * 1024;
 const INITIAL_WEBSOCKET_CONNECTIONS: usize = 2;
 const MAX_WEBSOCKET_CONNECTIONS: usize = 16;
@@ -73,6 +77,8 @@ pub struct LunaSamplerConfig {
     pub service_tier: Option<String>,
     /// Luna model's host-resolved encrypted-compaction compatibility hash.
     pub luna_compaction_hash: Option<String>,
+    /// Host-provided metrics capability with the owning session's attribution.
+    pub metrics: Option<Arc<dyn ExtensionMetrics>>,
 }
 
 /// One tool-less structured Luna request over an already-open connection.
@@ -146,6 +152,34 @@ struct ActiveRequest {
     scored: Arc<AtomicBool>,
 }
 
+fn record_token_usage(metrics: Option<&dyn ExtensionMetrics>, token_usage: Option<&TokenUsage>) {
+    let (Some(metrics), Some(token_usage)) = (metrics, token_usage) else {
+        return;
+    };
+
+    for (token_type, value) in [
+        ("total", token_usage.total_tokens.max(0)),
+        ("input", token_usage.input_tokens.max(0)),
+        ("cached_input", token_usage.cached_input()),
+        (
+            "cache_write_input",
+            token_usage.cache_write_input_tokens.max(0),
+        ),
+        ("non_cached_input", token_usage.non_cached_input()),
+        ("output", token_usage.output_tokens.max(0)),
+        (
+            "reasoning_output",
+            token_usage.reasoning_output_tokens.max(0),
+        ),
+    ] {
+        metrics.histogram(
+            CLASSIFICATION_TOKEN_USAGE_METRIC,
+            value,
+            &[("token_type", token_type)],
+        );
+    }
+}
+
 /// A bounded pool of authenticated Responses WebSockets dedicated to Luna sampling.
 pub struct LunaSampler {
     config: LunaSamplerConfig,
@@ -568,7 +602,8 @@ impl LunaSampler {
                             }
                         }
                     }
-                    ResponseEvent::Completed { .. } => {
+                    ResponseEvent::Completed { token_usage, .. } => {
+                        record_token_usage(self.config.metrics.as_deref(), token_usage.as_ref());
                         lease.reuse();
                         if !output.is_empty() {
                             return Ok(output);
@@ -592,6 +627,7 @@ impl LunaSampler {
                 if serde_json::from_str::<serde_json::Map<String, Value>>(&deltas).is_ok() {
                     scored.store(true, Ordering::Relaxed);
                     let mut remaining_events = stream.rx_event;
+                    let metrics = self.config.metrics.clone();
                     tokio::spawn(async move {
                         while let Some(event) = tokio::select! {
                             biased;
@@ -599,7 +635,8 @@ impl LunaSampler {
                             event = remaining_events.recv() => event,
                         } {
                             match event {
-                                Ok(ResponseEvent::Completed { .. }) => {
+                                Ok(ResponseEvent::Completed { token_usage, .. }) => {
+                                    record_token_usage(metrics.as_deref(), token_usage.as_ref());
                                     lease.reuse();
                                     break;
                                 }

codex-rs/ext/guardian-v2/src/sampler_tests.rs

@@ -1,4 +1,5 @@
 use anyhow::Result;
+use codex_extension_api::ExtensionMetrics;
 use codex_http_client::HttpClientFactory;
 use codex_http_client::OutboundProxyPolicy;
 use codex_login::AgentIdentityAuthPolicy;
@@ -18,6 +19,7 @@ use core_test_support::responses;
 use core_test_support::responses::WebSocketConnectionConfig;
 use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
+use core_test_support::responses::ev_completed_with_tokens;
 use core_test_support::responses::ev_model_verification_metadata;
 use core_test_support::responses::ev_output_text_delta;
 use core_test_support::skip_if_no_network;
@@ -25,12 +27,14 @@ use pretty_assertions::assert_eq;
 use serde_json::json;
 use std::collections::HashSet;
 use std::sync::Arc;
+use std::sync::Mutex;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
 use tokio::net::TcpListener;
 use tokio::net::TcpStream;
 
+use super::CLASSIFICATION_TOKEN_USAGE_METRIC;
 use super::LunaSampler;
 use super::LunaSamplerConfig;
 use super::LunaSamplingRequest;
@@ -124,6 +128,7 @@ fn sampler_config(base_url: String) -> LunaSamplerConfig {
         originator: Some("guardian-v2-test".to_owned()),
         service_tier: None,
         luna_compaction_hash: None,
+        metrics: None,
     }
 }
 
@@ -145,6 +150,74 @@ fn sample_request(turn_id: &str) -> LunaSamplingRequest {
     }
 }
 
+type RecordedMetric = (String, i64, Vec<(String, String)>);
+
+#[derive(Default)]
+struct RecordingMetrics(Mutex<Vec<RecordedMetric>>);
+
+impl ExtensionMetrics for RecordingMetrics {
+    fn counter(&self, _name: &str, _inc: i64, _tags: &[(&str, &str)]) {}
+
+    fn histogram(&self, name: &str, value: i64, tags: &[(&str, &str)]) {
+        self.0.lock().unwrap().push((
+            name.to_owned(),
+            value,
+            tags.iter()
+                .map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
+                .collect(),
+        ));
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn sampler_records_token_usage_after_returning_an_early_score() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let events = vec![
+        ev_output_text_delta(r#"{"score":0.25}"#),
+        ev_completed_with_tokens("response-1", /*total_tokens*/ 37),
+    ];
+    let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
+    let metrics = Arc::new(RecordingMetrics::default());
+    let mut config = sampler_config(format!(
+        "http://{}/v1",
+        server.uri().trim_start_matches("ws://")
+    ));
+    config.metrics = Some(metrics.clone());
+    let sampler = LunaSampler::connect(config).await?;
+
+    assert_eq!(
+        sampler.sample(sample_request("turn-1")).await?,
+        r#"{"score":0.25}"#
+    );
+    tokio::time::timeout(Duration::from_secs(2), async {
+        while metrics.0.lock().unwrap().len() < 7 {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await?;
+
+    assert_eq!(
+        *metrics.0.lock().unwrap(),
+        [
+            ("total", 37),
+            ("input", 37),
+            ("cached_input", 0),
+            ("cache_write_input", 0),
+            ("non_cached_input", 37),
+            ("output", 0),
+            ("reasoning_output", 0),
+        ]
+        .map(|(token_type, value)| (
+            CLASSIFICATION_TOKEN_USAGE_METRIC.to_owned(),
+            value,
+            vec![("token_type".to_owned(), token_type.to_owned())],
+        ))
+    );
+
+    Ok(())
+}
+
 struct RefreshableAuth(std::sync::Mutex<&'static str>);
 impl ExternalAuth for RefreshableAuth {
     fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
@@ -200,6 +273,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ
         originator: Some("guardian-v2-test".to_owned()),
         service_tier: None,
         luna_compaction_hash: None,
+        metrics: None,
     })
     .await?;
 
@@ -415,6 +489,7 @@ async fn sampler_returns_complete_json_before_terminal_response_events() -> Resu
         originator: None,
         service_tier: None,
         luna_compaction_hash: None,
+        metrics: None,
     })
     .await?;