diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index a4078c720c..640e0d4f29 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -30,6 +30,7 @@ use codex_guardian_reviewer::ConversationState; use codex_guardian_reviewer::ReviewModel; use codex_guardian_reviewer::ReviewSessionResult; use codex_guardian_reviewer::SessionDisposition; +use codex_guardian_reviewer::token_usage_delta; use codex_history::InitialHistory; use codex_history::RolloutItem; use codex_protocol::ThreadId; @@ -49,7 +50,6 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; -use codex_protocol::protocol::TokenUsage; use futures::future::BoxFuture; use serde_json::Value; use tokio::sync::Mutex; @@ -161,20 +161,6 @@ fn had_prior_review_context(prompt_mode: &GuardianPromptMode) -> bool { matches!(prompt_mode, GuardianPromptMode::Delta { .. }) } -fn token_usage_delta(start: &TokenUsage, end: &TokenUsage) -> TokenUsage { - TokenUsage { - input_tokens: (end.input_tokens - start.input_tokens).max(0), - cached_input_tokens: (end.cached_input_tokens - start.cached_input_tokens).max(0), - cache_write_input_tokens: (end.cache_write_input_tokens - start.cache_write_input_tokens) - .max(0), - output_tokens: (end.output_tokens - start.output_tokens).max(0), - reasoning_output_tokens: (end.reasoning_output_tokens - start.reasoning_output_tokens) - .max(0), - total_tokens: (end.total_tokens - start.total_tokens).max(0), - codex_rollout_budget_units: None, - } -} - type GuardianReviewForkSnapshot = ConversationCheckpoint; /// Host-owned history and admitted evidence used to seed a private reviewer fork. diff --git a/codex-rs/core/src/guardian/review_session_tests.rs b/codex-rs/core/src/guardian/review_session_tests.rs index 6e71ad6b91..93c913764c 100644 --- a/codex-rs/core/src/guardian/review_session_tests.rs +++ b/codex-rs/core/src/guardian/review_session_tests.rs @@ -673,41 +673,6 @@ async fn guardian_review_session_config_resolves_policy_and_template( ); } -#[test] -fn token_usage_delta_never_reports_negative_usage() { - let start = TokenUsage { - input_tokens: 10, - cached_input_tokens: 8, - cache_write_input_tokens: 8, - output_tokens: 6, - reasoning_output_tokens: 4, - total_tokens: 28, - codex_rollout_budget_units: None, - }; - let end = TokenUsage { - input_tokens: 15, - cached_input_tokens: 7, - cache_write_input_tokens: 7, - output_tokens: 10, - reasoning_output_tokens: 2, - total_tokens: 34, - codex_rollout_budget_units: None, - }; - - assert_eq!( - token_usage_delta(&start, &end), - TokenUsage { - input_tokens: 5, - cached_input_tokens: 0, - cache_write_input_tokens: 0, - output_tokens: 4, - reasoning_output_tokens: 0, - total_tokens: 6, - codex_rollout_budget_units: None, - } - ); -} - #[tokio::test] async fn run_review_on_reused_session_waits_for_submitted_turn() { let (review_session, tx_event, rx_sub) = test_review_session().await; diff --git a/codex-rs/ext/guardian-reviewer/src/lib.rs b/codex-rs/ext/guardian-reviewer/src/lib.rs index 50e8c8fcce..f2ff22fe2d 100644 --- a/codex-rs/ext/guardian-reviewer/src/lib.rs +++ b/codex-rs/ext/guardian-reviewer/src/lib.rs @@ -32,6 +32,7 @@ pub(crate) use circuit_breaker::AUTO_REVIEW_DENIAL_WINDOW_SIZE; pub(crate) use circuit_breaker::GuardianRejectionCircuitBreaker; pub(crate) use circuit_breaker::GuardianRejectionCircuitBreakerAction; pub(crate) use circuit_breaker::GuardianRejectionCircuitBreakerPolicy; +pub use metrics::token_usage_delta; pub use model::ReviewModel; pub use model::select_review_model; pub use outcome::GuardianReviewError; diff --git a/codex-rs/ext/guardian-reviewer/src/metrics.rs b/codex-rs/ext/guardian-reviewer/src/metrics.rs index c55dc49044..9bed7e42df 100644 --- a/codex-rs/ext/guardian-reviewer/src/metrics.rs +++ b/codex-rs/ext/guardian-reviewer/src/metrics.rs @@ -1,4 +1,5 @@ -//! Emits Guardian metrics with bounded tags and the existing telemetry schema. +//! Calculates per-review token usage and emits Guardian metrics with bounded tags +//! and the existing telemetry schema. use std::time::Duration; @@ -21,6 +22,20 @@ use codex_protocol::protocol::GuardianRiskLevel; use codex_protocol::protocol::GuardianUserAuthorization; use codex_protocol::protocol::TokenUsage; +pub fn token_usage_delta(start: &TokenUsage, end: &TokenUsage) -> TokenUsage { + TokenUsage { + input_tokens: (end.input_tokens - start.input_tokens).max(0), + cached_input_tokens: (end.cached_input_tokens - start.cached_input_tokens).max(0), + cache_write_input_tokens: (end.cache_write_input_tokens - start.cache_write_input_tokens) + .max(0), + output_tokens: (end.output_tokens - start.output_tokens).max(0), + reasoning_output_tokens: (end.reasoning_output_tokens - start.reasoning_output_tokens) + .max(0), + total_tokens: (end.total_tokens - start.total_tokens).max(0), + codex_rollout_budget_units: None, + } +} + pub(crate) fn emit_guardian_review_metrics( session_telemetry: &SessionTelemetry, result: &GuardianReviewAnalyticsResult, @@ -295,6 +310,41 @@ mod tests { use std::collections::BTreeMap; use test_case::test_case; + #[test] + fn token_usage_delta_never_reports_negative_usage() { + let start = TokenUsage { + input_tokens: 10, + cached_input_tokens: 8, + cache_write_input_tokens: 8, + output_tokens: 6, + reasoning_output_tokens: 4, + total_tokens: 28, + codex_rollout_budget_units: None, + }; + let end = TokenUsage { + input_tokens: 15, + cached_input_tokens: 7, + cache_write_input_tokens: 7, + output_tokens: 10, + reasoning_output_tokens: 2, + total_tokens: 34, + codex_rollout_budget_units: None, + }; + + assert_eq!( + token_usage_delta(&start, &end), + TokenUsage { + input_tokens: 5, + cached_input_tokens: 0, + cache_write_input_tokens: 0, + output_tokens: 4, + reasoning_output_tokens: 0, + total_tokens: 6, + codex_rollout_budget_units: None, + } + ); + } + fn test_session_telemetry() -> SessionTelemetry { let exporter = InMemoryMetricExporter::default(); let metrics = MetricsClient::new(