openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-10 21:00:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Record Guardian v2 classification token usage (#39303)

Browse Source 
## 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
This commit is contained in:
felixxia-oai
2026-08-18 22:33:28 +00:00
committed by copyberry
parent fe50b61689
commit b473c4e6ab
4 changed files with 147 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/ext/guardian-v2/src/extension.rs
+1
View File
@@ -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
+32 -5
View File
@@ -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
+39 -2
View File
@@ -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
+75
View File
@@ -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?;