From 2f8603f07547247e698748884542ba60a157621c Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Mon, 14 Sep 2026 12:41:01 +0000 Subject: [PATCH] Extract Guardian sampler execution into a dedicated module (#45420) ## What changed Move request execution from `LunaSampler` into `SamplingExecution` in `sampler/execution.rs`, keeping request preparation and active-request tracking in the sampler. Preserve the existing retry, authentication recovery, cancellation, streaming, connection reuse, and token accounting behavior. GitOrigin-RevId: 5e50116459fe9cc196c30c58e42773fc32be3d02 --- .../guardian-v2/src/async_scorer/sampler.rs | 279 +--------------- .../src/async_scorer/sampler/execution.rs | 302 ++++++++++++++++++ 2 files changed, 315 insertions(+), 266 deletions(-) create mode 100644 codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index c89cc49414..4ad0c0e2ea 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -2,9 +2,9 @@ //! Both transports share request identity, retry, cancellation, and output handling. mod connection_pool; +mod execution; use connection_pool::ConnectionPool; -use std::collections::HashMap; use std::collections::VecDeque; use std::sync::Arc; use std::sync::Mutex; @@ -14,14 +14,10 @@ use std::sync::atomic::Ordering; use codex_api::ApiError; use codex_api::Reasoning; use codex_api::ReasoningContext; -use codex_api::ResponseEvent; use codex_api::ResponsesApiRequest; -use codex_api::ResponsesEndpoint; -use codex_api::TransportError; use codex_extension_api::ExtensionMetrics; use codex_http_client::HttpClientFactory; use codex_login::AgentIdentityAuthPolicy; -use codex_login::UnauthorizedRecovery; use codex_model_provider::SharedModelProvider; use codex_protocol::ResponseItemId; use codex_protocol::error::CodexErr; @@ -29,9 +25,6 @@ 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::StatusCode; -use serde_json::json; use thiserror::Error; use tokio::sync::oneshot; use uuid::Uuid; @@ -42,10 +35,6 @@ pub(crate) const CLASSIFICATION_TOKEN_USAGE_METRIC: &str = const MAX_OUTPUT_BYTES: usize = 8 * 1024; pub(super) const INITIAL_WEBSOCKET_CONNECTIONS: usize = if cfg!(test) { 2 } else { 8 }; const MAX_CONCURRENT_REQUESTS: usize = 16; -const MAX_SAMPLING_RETRIES: usize = 2; -const RESPONSES_LITE_METADATA_KEY: &str = - "ws_request_header_x_openai_internal_codex_responses_lite"; -const TURN_METADATA_KEY: &str = "x-codex-turn-metadata"; /// Host-owned provider, authentication, and attribution for one Luna connection. pub struct LunaSamplerConfig { @@ -129,34 +118,6 @@ struct ActiveRequest { scored: Arc, } -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)], - ); - } -} - /// Runs bounded Luna classifications over pooled WebSockets or HTTP. pub struct LunaSampler { config: Arc, @@ -189,65 +150,6 @@ impl LunaSampler { } } - async fn retry_after_failure( - &self, - error: &LunaSamplerError, - auth_recovery: &mut Option, - retries: &mut usize, - ) -> bool { - let retryable = match error { - LunaSamplerError::ConnectionTimeout - | LunaSamplerError::Api( - ApiError::Retryable { .. } - | ApiError::RateLimitExceeded { .. } - | ApiError::Stream(_) - | ApiError::ServerOverloaded, - ) - | LunaSamplerError::Api(ApiError::Transport( - TransportError::RetryLimit - | TransportError::Timeout - | TransportError::Connection(_) - | TransportError::Network(_), - )) => true, - LunaSamplerError::Api(ApiError::Transport(TransportError::Http { status, .. })) - | LunaSamplerError::Api(ApiError::Api { status, .. }) => { - if *status == StatusCode::UNAUTHORIZED { - let Some(recovery) = auth_recovery.as_mut() else { - return false; - }; - if !recovery.has_next() || recovery.next().await.is_err() { - return false; - } - self.connections.clear(); - return true; - } else { - status.is_server_error() || *status == StatusCode::TOO_MANY_REQUESTS - } - } - LunaSamplerError::Provider(_) - | LunaSamplerError::MissingOutput - | LunaSamplerError::OutputTooLarge - | LunaSamplerError::Superseded - | LunaSamplerError::IncompatibleCompaction - | LunaSamplerError::InputTooLarge - | LunaSamplerError::Api( - ApiError::Transport(TransportError::Build(_)) - | ApiError::ContextWindowExceeded - | ApiError::QuotaExceeded - | ApiError::UsageNotIncluded - | ApiError::RateLimit(_) - | ApiError::InvalidRequest { .. } - | ApiError::MisalignmentPolicyViolation { .. } - | ApiError::CyberPolicy { .. }, - ) => false, - }; - if retryable && *retries < MAX_SAMPLING_RETRIES { - *retries += 1; - return true; - } - false - } - /// Sends one tool-less classification request using an available transport. pub async fn sample(&self, request: LunaSamplingRequest) -> Result { if request.parent_compaction.is_some() @@ -320,7 +222,7 @@ impl LunaSampler { if total_tokens > self.config.max_input_tokens.saturating_sub(/*rhs*/ 256) { return Err(LunaSamplerError::InputTooLarge); } - let mut request = ResponsesApiRequest { + let request = ResponsesApiRequest { model: MODEL.to_owned(), instructions: String::new(), input, @@ -342,7 +244,7 @@ impl LunaSampler { client_metadata: None, access_programs: None, }; - let (supersede, mut superseded) = oneshot::channel(); + let (supersede, superseded) = oneshot::channel(); let scored = Arc::new(AtomicBool::new(false)); { let mut active_requests = self @@ -364,172 +266,17 @@ impl LunaSampler { scored: Arc::clone(&scored), }); } - let mut retries = 0; - let mut auth_recovery = self - .config - .provider - .auth_manager() - .map(|manager| manager.unauthorized_recovery()); - 'retry: loop { - let lease = match tokio::select! { - biased; - _ = &mut superseded => return Err(LunaSamplerError::Superseded), - lease = self.connections.lease() => lease, - } { - Ok(lease) => lease, - Err(error) => { - if self - .retry_after_failure(&error, &mut auth_recovery, &mut retries) - .await - { - continue; - } - return Err(error); - } - }; - request.service_tier = if lease.endpoint == ResponsesEndpoint::GuardianClassifier { - None - } else { - self.config.service_tier.clone() - }; - let thread_id = &lease.thread_id; - let mut turn_metadata = json!({ - "session_id": self.config.session_id, - "thread_id": thread_id, - "guardian_classifier_source_thread_id": self.config.thread_id, - "turn_id": turn_id, - "parent_turn_id": parent_turn_id, - "thread_source": "guardian_classifier", - "turn_trigger": "guardian_classifier", - }); - let mut client_metadata = HashMap::from([ - ("session_id".to_owned(), self.config.session_id.clone()), - ("thread_id".to_owned(), thread_id.clone()), - ("turn_id".to_owned(), turn_id.clone()), - ("parent_turn_id".to_owned(), parent_turn_id.clone()), - ("x-openai-subagent".to_owned(), "guardian".to_owned()), - // Classifier requests do not advance their own context window. - ("x-codex-window-id".to_owned(), format!("{thread_id}:0")), - (RESPONSES_LITE_METADATA_KEY.to_owned(), "true".to_owned()), - ]); - if let Some(root_turn_id) = &root_turn_id { - client_metadata.insert("root_turn_id".to_owned(), root_turn_id.clone()); - turn_metadata["root_turn_id"] = json!(root_turn_id); - } - client_metadata.insert(TURN_METADATA_KEY.to_owned(), turn_metadata.to_string()); - if lease.endpoint == ResponsesEndpoint::GuardianClassifier - && let Some(parent_response_id) = &parent_response_id - { - client_metadata.insert("parent_response_id".to_owned(), parent_response_id.clone()); - } - request.client_metadata = Some(client_metadata); - let mut stream = match tokio::select! { - biased; - _ = &mut superseded => return Err(LunaSamplerError::Superseded), - stream = lease.stream_request(&request) => stream, - } { - Ok(stream) => stream, - Err(error) => { - let error = LunaSamplerError::Api(error); - if self - .retry_after_failure(&error, &mut auth_recovery, &mut retries) - .await - { - continue; - } - return Err(error); - } - }; - - let mut output = String::new(); - while let Some(event) = tokio::select! { - biased; - _ = &mut superseded => { - return if scored.load(Ordering::Relaxed) && !output.is_empty() { - Ok(output) - } else { - Err(LunaSamplerError::Superseded) - }; - } - event = stream.rx_event.recv() => event, - } { - let event = match event { - Ok(event) => event, - Err(error) => { - let error = LunaSamplerError::Api(error); - if self - .retry_after_failure(&error, &mut auth_recovery, &mut retries) - .await - { - continue 'retry; - } - return Err(error); - } - }; - match event { - ResponseEvent::OutputTextDelta(delta) => { - if delta.is_empty() { - continue; - } - if delta.len() > MAX_OUTPUT_BYTES { - return Err(LunaSamplerError::OutputTooLarge); - } - // The first output token is the complete classification. - // Later output cannot revise that decision; drain it only - // to preserve connection reuse and token accounting. - 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; - _ = &mut superseded => None, - event = remaining_events.recv() => event, - } { - match event { - Ok(ResponseEvent::Completed { token_usage, .. }) => { - record_token_usage( - metrics.as_deref(), - token_usage.as_ref(), - ); - lease.reuse(); - break; - } - Err(_) => break, - _ => {} - } - } - }); - return Ok(delta); - } - ResponseEvent::OutputItemDone(ResponseItem::Message { - role, content, .. - }) if role == "assistant" => { - for item in content { - if let ContentItem::OutputText { text } = item { - output.push_str(&text); - } - } - } - 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); - } - return Err(LunaSamplerError::MissingOutput); - } - _ => {} - } - if output.len() > MAX_OUTPUT_BYTES { - return Err(LunaSamplerError::OutputTooLarge); - } - if !output.is_empty() { - scored.store(true, Ordering::Relaxed); - } - } - return Err(LunaSamplerError::MissingOutput); + execution::SamplingExecution { + config: Arc::clone(&self.config), + connections: Arc::clone(&self.connections), + request, + turn_id, + parent_response_id, + parent_turn_id, + root_turn_id, } + .run(superseded, scored) + .await } } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs new file mode 100644 index 0000000000..642587e94b --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs @@ -0,0 +1,302 @@ +//! Executes one prepared classifier request with the existing retry and streaming rules. +//! The first output returns immediately; a detached drain preserves reuse and token accounting. + +use super::CLASSIFICATION_TOKEN_USAGE_METRIC; +use super::ConnectionPool; +use super::LunaSamplerConfig; +use super::LunaSamplerError; +use super::MAX_OUTPUT_BYTES; +use codex_api::ApiError; +use codex_api::ResponseEvent; +use codex_api::ResponsesApiRequest; +use codex_api::ResponsesEndpoint; +use codex_api::TransportError; +use codex_extension_api::ExtensionMetrics; +use codex_login::UnauthorizedRecovery; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::TokenUsage; +use http::StatusCode; +use serde_json::json; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use tokio::sync::oneshot; + +const MAX_SAMPLING_RETRIES: usize = 2; +const RESPONSES_LITE_METADATA_KEY: &str = + "ws_request_header_x_openai_internal_codex_responses_lite"; +const TURN_METADATA_KEY: &str = "x-codex-turn-metadata"; + +pub(super) struct SamplingExecution { + pub(super) config: Arc, + pub(super) connections: Arc, + pub(super) request: ResponsesApiRequest, + pub(super) turn_id: String, + pub(super) parent_response_id: Option, + pub(super) parent_turn_id: String, + pub(super) root_turn_id: Option, +} + +impl SamplingExecution { + async fn retry_after_failure( + &self, + error: &LunaSamplerError, + auth_recovery: &mut Option, + retries: &mut usize, + ) -> bool { + let retryable = match error { + LunaSamplerError::ConnectionTimeout + | LunaSamplerError::Api( + ApiError::Retryable { .. } + | ApiError::RateLimitExceeded { .. } + | ApiError::Stream(_) + | ApiError::ServerOverloaded, + ) + | LunaSamplerError::Api(ApiError::Transport( + TransportError::RetryLimit + | TransportError::Timeout + | TransportError::Connection(_) + | TransportError::Network(_), + )) => true, + LunaSamplerError::Api(ApiError::Transport(TransportError::Http { status, .. })) + | LunaSamplerError::Api(ApiError::Api { status, .. }) => { + if *status == StatusCode::UNAUTHORIZED { + let Some(recovery) = auth_recovery.as_mut() else { + return false; + }; + if !recovery.has_next() || recovery.next().await.is_err() { + return false; + } + self.connections.clear(); + return true; + } else { + status.is_server_error() || *status == StatusCode::TOO_MANY_REQUESTS + } + } + LunaSamplerError::Provider(_) + | LunaSamplerError::MissingOutput + | LunaSamplerError::OutputTooLarge + | LunaSamplerError::Superseded + | LunaSamplerError::IncompatibleCompaction + | LunaSamplerError::InputTooLarge + | LunaSamplerError::Api( + ApiError::Transport(TransportError::Build(_)) + | ApiError::ContextWindowExceeded + | ApiError::QuotaExceeded + | ApiError::UsageNotIncluded + | ApiError::RateLimit(_) + | ApiError::InvalidRequest { .. } + | ApiError::MisalignmentPolicyViolation { .. } + | ApiError::CyberPolicy { .. }, + ) => false, + }; + if retryable && *retries < MAX_SAMPLING_RETRIES { + *retries += 1; + return true; + } + false + } + + pub(super) async fn run( + mut self, + mut superseded: oneshot::Receiver<()>, + scored: Arc, + ) -> Result { + let mut retries = 0; + let mut auth_recovery = self + .config + .provider + .auth_manager() + .map(|manager| manager.unauthorized_recovery()); + 'retry: loop { + let lease = match tokio::select! { + biased; + _ = &mut superseded => return Err(LunaSamplerError::Superseded), + lease = self.connections.lease() => lease, + } { + Ok(lease) => lease, + Err(error) => { + if self + .retry_after_failure(&error, &mut auth_recovery, &mut retries) + .await + { + continue; + } + return Err(error); + } + }; + self.request.service_tier = if lease.endpoint == ResponsesEndpoint::GuardianClassifier { + None + } else { + self.config.service_tier.clone() + }; + let thread_id = &lease.thread_id; + let mut turn_metadata = json!({ + "session_id": self.config.session_id, + "thread_id": thread_id, + "guardian_classifier_source_thread_id": self.config.thread_id, + "turn_id": self.turn_id, + "parent_turn_id": self.parent_turn_id, + "thread_source": "guardian_classifier", + "turn_trigger": "guardian_classifier", + }); + let mut client_metadata = HashMap::from([ + ("session_id".to_owned(), self.config.session_id.clone()), + ("thread_id".to_owned(), thread_id.clone()), + ("turn_id".to_owned(), self.turn_id.clone()), + ("parent_turn_id".to_owned(), self.parent_turn_id.clone()), + ("x-openai-subagent".to_owned(), "guardian".to_owned()), + // Classifier requests do not advance their own context window. + ("x-codex-window-id".to_owned(), format!("{thread_id}:0")), + (RESPONSES_LITE_METADATA_KEY.to_owned(), "true".to_owned()), + ]); + if let Some(root_turn_id) = &self.root_turn_id { + client_metadata.insert("root_turn_id".to_owned(), root_turn_id.clone()); + turn_metadata["root_turn_id"] = json!(root_turn_id); + } + client_metadata.insert(TURN_METADATA_KEY.to_owned(), turn_metadata.to_string()); + if lease.endpoint == ResponsesEndpoint::GuardianClassifier + && let Some(parent_response_id) = &self.parent_response_id + { + client_metadata.insert("parent_response_id".to_owned(), parent_response_id.clone()); + } + self.request.client_metadata = Some(client_metadata); + let mut stream = match tokio::select! { + biased; + _ = &mut superseded => return Err(LunaSamplerError::Superseded), + stream = lease.stream_request(&self.request) => stream, + } { + Ok(stream) => stream, + Err(error) => { + let error = LunaSamplerError::Api(error); + if self + .retry_after_failure(&error, &mut auth_recovery, &mut retries) + .await + { + continue; + } + return Err(error); + } + }; + + let mut output = String::new(); + while let Some(event) = tokio::select! { + biased; + _ = &mut superseded => { + return if scored.load(Ordering::Relaxed) && !output.is_empty() { + Ok(output) + } else { + Err(LunaSamplerError::Superseded) + }; + } + event = stream.rx_event.recv() => event, + } { + let event = match event { + Ok(event) => event, + Err(error) => { + let error = LunaSamplerError::Api(error); + if self + .retry_after_failure(&error, &mut auth_recovery, &mut retries) + .await + { + continue 'retry; + } + return Err(error); + } + }; + match event { + ResponseEvent::OutputTextDelta(delta) => { + if delta.is_empty() { + continue; + } + if delta.len() > MAX_OUTPUT_BYTES { + return Err(LunaSamplerError::OutputTooLarge); + } + // The first output token is the complete classification. + // Later output cannot revise that decision; drain it only + // to preserve connection reuse and token accounting. + 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; + _ = &mut superseded => None, + event = remaining_events.recv() => event, + } { + match event { + Ok(ResponseEvent::Completed { token_usage, .. }) => { + record_token_usage( + metrics.as_deref(), + token_usage.as_ref(), + ); + lease.reuse(); + break; + } + Err(_) => break, + _ => {} + } + } + }); + return Ok(delta); + } + ResponseEvent::OutputItemDone(ResponseItem::Message { + role, content, .. + }) if role == "assistant" => { + for item in content { + if let ContentItem::OutputText { text } = item { + output.push_str(&text); + } + } + } + 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); + } + return Err(LunaSamplerError::MissingOutput); + } + _ => {} + } + if output.len() > MAX_OUTPUT_BYTES { + return Err(LunaSamplerError::OutputTooLarge); + } + if !output.is_empty() { + scored.store(true, Ordering::Relaxed); + } + } + return Err(LunaSamplerError::MissingOutput); + } + } +} + +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)], + ); + } +}