From c4f0e09912fd4c1173258ccd17e33284a876e826 Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Sat, 7 Nov 2026 07:35:28 +0000 Subject: [PATCH] Allow asynchronous model request budget checks (#60080) ## What changed Make `ModelRequestContributor::budget_exhausted` return an `ExtensionFuture` and await budget checks before HTTP and WebSocket generation. WebSocket warmup continues to skip budget checks. If the account changes during the post-connect budget check, rebuild the request setup and reconnect before sending the WebSocket request. ## Testing Update tests to cover a pending budget check that resolves to exhaustion, warmup connection reuse, and an account switch while the post-connect budget check is awaiting completion. GitOrigin-RevId: 462c74b86a46a4b2276f901c69187fc7cbc3ff81 --- codex-rs/core/src/client.rs | 31 +++++----- codex-rs/core/src/model_request.rs | 15 ++--- .../core/tests/suite/client_websockets.rs | 37 +++++++----- .../suite/client_websockets_runtime_auth.rs | 58 ++++++++++++++++--- codex-rs/core/tests/suite/model_request.rs | 12 ++-- .../ext/extension-api/src/model_request.rs | 7 ++- 6 files changed, 105 insertions(+), 55 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index de72f5e1a9..7a06e6e90e 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -2068,11 +2068,8 @@ impl ResponsesHttpExecutor<'_> { options.extra_headers.extend(responses_headers); let thread_id = self.client.state.thread_id.to_string(); let kind = codex_extension_api::ModelRequestKind::Generation; - crate::model_request::check_budget( - &self.client.request_contributors, - &thread_id, - kind, - )?; + crate::model_request::check_budget(&self.client.request_contributors, &thread_id, kind) + .await?; let interceptors = crate::model_request::prepare( &self.client.request_contributors, &thread_id, @@ -2197,11 +2194,8 @@ impl ModelClientSession { let mut provider_auth_recovery_attempted = false; let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - crate::model_request::check_budget( - &self.client.request_contributors, - &thread_id, - kind, - )?; + crate::model_request::check_budget(&self.client.request_contributors, &thread_id, kind) + .await?; let client_setup = self .client .current_client_setup(ClientRouting::Workspace) @@ -2297,11 +2291,18 @@ impl ModelClientSession { } // The budget can be exhausted while the connection is being established. - crate::model_request::check_budget( - &self.client.request_contributors, - &thread_id, - kind, - )?; + if !warmup && !self.client.request_contributors.is_empty() { + crate::model_request::check_budget( + &self.client.request_contributors, + &thread_id, + kind, + ) + .await?; + // An account switch during the check requires fresh request setup and a new socket. + if self.websocket_session.auth_owner != self.client.auth_owner() { + continue; + } + } let mut client_metadata = self.client.build_ws_client_metadata( responses_metadata, diff --git a/codex-rs/core/src/model_request.rs b/codex-rs/core/src/model_request.rs index a714806690..353794945b 100644 --- a/codex-rs/core/src/model_request.rs +++ b/codex-rs/core/src/model_request.rs @@ -9,17 +9,18 @@ use codex_protocol::error::Result; use std::collections::HashMap; use std::sync::Arc; -pub(crate) fn check_budget( +pub(crate) async fn check_budget( contributors: &[Arc], thread_id: &str, kind: ModelRequestKind, ) -> Result<()> { - if kind == ModelRequestKind::Generation - && contributors - .iter() - .any(|contributor| contributor.budget_exhausted(thread_id)) - { - return Err(CodexErr::SessionBudgetExceeded); + if kind == ModelRequestKind::Warmup { + return Ok(()); + } + for contributor in contributors { + if contributor.budget_exhausted(thread_id).await { + return Err(CodexErr::SessionBudgetExceeded); + } } Ok(()) } diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 620bd2f231..ae69634c9a 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -72,13 +72,13 @@ use core_test_support::test_codex::RecordingUserInstructionsProvider; use core_test_support::test_codex::test_codex; use core_test_support::tracing::install_test_tracing; use core_test_support::wait_for_event; +use futures::FutureExt; use futures::StreamExt; use opentelemetry_sdk::metrics::InMemoryMetricExporter; use opentelemetry_sdk::metrics::data::ResourceMetrics; use pretty_assertions::assert_eq; use serde_json::json; use std::sync::Arc; -use std::sync::atomic::Ordering; use std::time::Duration; use tempfile::TempDir; use tracing::Instrument; @@ -742,7 +742,7 @@ async fn responses_websocket_request_prewarm_reuses_connection() { let mut provider = websocket_provider(&server); provider.name = ModelProviderInfo::create_openai_provider(/*base_url*/ None).name; - let budget = Arc::new(BudgetCache(true.into())); + let budget = Arc::new(BudgetCache::default()); let harness = websocket_harness_with_provider_factory( provider, /*runtime_metrics_enabled*/ true, @@ -755,19 +755,24 @@ async fn responses_websocket_request_prewarm_reuses_connection() { .await; let mut client_session = harness.client.new_session(); let prompt = prompt_with_input(vec![message_item("hello")]); - let Err(error) = client_session - .stream( - &prompt, - &harness.model_info, - &harness.session_telemetry, - harness.effort.clone(), - harness.summary, - /*service_tier*/ None, - &turn_metadata(&harness, /*turn_id*/ None), - &InferenceTraceContext::disabled(), - ) - .await - else { + let responses_metadata = turn_metadata(&harness, /*turn_id*/ None); + let inference_trace = InferenceTraceContext::disabled(); + let mut generation = Box::pin(client_session.stream( + &prompt, + &harness.model_info, + &harness.session_telemetry, + harness.effort.clone(), + harness.summary, + /*service_tier*/ None, + &responses_metadata, + &inference_trace, + )); + { + let mut exhausted = budget.0.lock().await; + assert!(generation.as_mut().now_or_never().is_none()); + *exhausted = true; + } + let Err(error) = generation.await else { panic!("exhausted generation unexpectedly succeeded"); }; assert!( @@ -789,7 +794,7 @@ async fn responses_websocket_request_prewarm_reuses_connection() { ) .await .expect("websocket prewarm failed"); - budget.0.store(false, Ordering::Relaxed); + *budget.0.lock().await = false; stream_until_complete(&mut client_session, &harness, &prompt).await; assert_eq!(server.handshakes().len(), 1); diff --git a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs index 4bd2193fd3..9e0d588ce8 100644 --- a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs +++ b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs @@ -2,6 +2,9 @@ use super::*; use base64::Engine; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ModelRequestInput; +use codex_extension_api::ModelResponseInterceptor; use codex_features::RuntimeAuthMode; use codex_login::ExternalAuth; use codex_login::ExternalAuthFuture; @@ -16,8 +19,11 @@ use codex_runtime_agent_instance::RuntimeRole; use core_test_support::responses; use pretty_assertions::assert_eq; use serde_json::Value; +use std::collections::VecDeque; +use std::sync::Mutex; use std::sync::RwLock; use test_case::test_case; +use tokio::sync::Barrier; use uuid::Uuid; use wiremock::Mock; use wiremock::MockServer; @@ -28,6 +34,26 @@ use wiremock::matchers::path; struct Credentials(RwLock); +#[derive(Debug, Default)] +struct BudgetChecks(Mutex>>>); + +impl ModelRequestContributor for BudgetChecks { + fn budget_exhausted<'a>(&'a self, _thread_id: &'a str) -> ExtensionFuture<'a, bool> { + let gate = self.0.lock().unwrap().pop_front().flatten(); + Box::pin(async move { + if let Some(gate) = gate { + gate.wait().await; // Signal that Core reached this check. + gate.wait().await; // Wait for the account switch. + } + false + }) + } + + fn request(&self, _: ModelRequestInput<'_>) -> Option> { + None + } +} + impl ExternalAuth for Credentials { fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { Box::pin(async { Ok(self.0.read().unwrap().clone()) }) @@ -56,6 +82,7 @@ async fn runtime_harness( registration: &MockServer, mut provider: ModelProviderInfo, mode: RuntimeAuthMode, + request_contributors: Vec>, ) -> (WebsocketTestHarness, Arc) { provider.requires_openai_auth = true; let initial = CodexAuth::from_api_key("key-a"); @@ -66,7 +93,7 @@ async fn runtime_harness( /*concurrent_reasoning_summaries_enabled*/ false, &[], Some(initial.clone()), - Vec::new(), + request_contributors, move |provider, credentials| { let credentials = credentials.expect("runtime fixture requires credentials"); let runtime = Arc::new(RuntimeAuthManager::new( @@ -218,7 +245,7 @@ async fn http_requests_keep_matching_runtime_and_credentials_across_retries_and_ }, )); } - let (harness, credentials) = runtime_harness(&server, provider, mode).await; + let (harness, credentials) = runtime_harness(&server, provider, mode, Vec::new()).await; let mut session = harness.client.new_session(); complete(&mut session, &harness).await.unwrap(); *credentials.0.write().unwrap() = CodexAuth::from_api_key("key-other-user"); @@ -336,7 +363,9 @@ async fn websocket_reconnects_for_replacement_and_account_switch_within_a_turn( ("x-openai-fedramp".into(), "true".into()), ("x-test-provider-header".into(), "preserved".into()), ])); - let (harness, credentials) = runtime_harness(®istration, provider, mode).await; + let budget = Arc::new(BudgetChecks::default()); + let (harness, credentials) = + runtime_harness(®istration, provider, mode, vec![budget.clone()]).await; let mut session = harness.client.new_session(); let mut input = vec![message_item("first")]; stream_until_complete(&mut session, &harness, &prompt_with_input(input.clone())).await; @@ -363,9 +392,22 @@ async fn websocket_reconnects_for_replacement_and_account_switch_within_a_turn( stream_until_complete(&mut session, &harness, &prompt_with_input(input.clone())).await; assert_eq!(server.handshakes().len(), 3); - *credentials.0.write().unwrap() = CodexAuth::from_api_key("key-other-user"); input.push(message_item("fourth")); - stream_until_complete(&mut session, &harness, &prompt_with_input(input)).await; + let prompt = prompt_with_input(input); + let gate = Arc::new(Barrier::new(2)); + // Allow the pre-connect check, then pause after Core selects the existing socket. + budget.0.lock().unwrap().extend([None, Some(gate.clone())]); + let generation = stream_until_complete(&mut session, &harness, &prompt); + tokio::pin!(generation); + tokio::select! { + () = &mut generation => panic!("generation finished before the post-connect budget check"), + _ = gate.wait() => {} + } + let previous_owner = harness.auth_manager.auth_owner(); + *credentials.0.write().unwrap() = CodexAuth::from_api_key("key-other-user"); + harness.auth_manager.reload().await; + assert!(!previous_owner.is_current()); + tokio::join!(generation, gate.wait()); let handshakes = server.handshakes(); assert_eq!(handshakes.len(), 4); let runtime_ids: Vec<_> = handshakes @@ -439,7 +481,7 @@ async fn mismatched_organization_never_sends_model_request(mode: RuntimeAuthMode provider.supports_websockets = false; provider.http_headers = Some([(String::from("OpenAI-Organization"), "another-org".into())].into()); - let (harness, _) = runtime_harness(&server, provider, mode).await; + let (harness, _) = runtime_harness(&server, provider, mode, Vec::new()).await; let error = complete(&mut harness.client.new_session(), &harness) .await .unwrap_err(); @@ -479,7 +521,7 @@ async fn gateway_refresh_failure_blocks_runtime_inference(mode: RuntimeAuthMode) scheme: "Bearer".into(), }, )); - let (harness, _) = runtime_harness(&server, provider.clone(), mode).await; + let (harness, _) = runtime_harness(&server, provider.clone(), mode, Vec::new()).await; let _gateway = codex_model_provider::test_support::seed_gateway_auth( &provider, &harness.auth_manager, @@ -542,7 +584,7 @@ async fn model_runtime_auth_strips_credentials_on_cross_origin_redirects(mode: R let mut provider = ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri()))); provider.supports_websockets = false; - let (harness, _) = runtime_harness(&server, provider, mode).await; + let (harness, _) = runtime_harness(&server, provider, mode, Vec::new()).await; complete(&mut harness.client.new_session(), &harness) .await .unwrap(); diff --git a/codex-rs/core/tests/suite/model_request.rs b/codex-rs/core/tests/suite/model_request.rs index efd43bcee9..02299217b6 100644 --- a/codex-rs/core/tests/suite/model_request.rs +++ b/codex-rs/core/tests/suite/model_request.rs @@ -1,6 +1,7 @@ //! Output gates allow inference immediately, hold complete output, and cancel with the turn. use anyhow::Result; use codex_core::TurnInputRequest; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::GuardianAssessment; use codex_extension_api::ModelRequestContributor; @@ -29,8 +30,7 @@ use futures::StreamExt; use pretty_assertions::assert_eq; use serde_json::json; use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; +use tokio::sync::Mutex; use tokio::sync::Semaphore; use tokio::time::Duration; use tokio::time::timeout; @@ -130,15 +130,15 @@ impl Drop for Permit { } #[derive(Debug, Default)] -pub(super) struct BudgetCache(pub(super) AtomicBool); +pub(super) struct BudgetCache(pub(super) Mutex); impl ModelRequestContributor for BudgetCache { fn request(&self, _input: ModelRequestInput<'_>) -> Option> { None } - fn budget_exhausted(&self, _thread_id: &str) -> bool { - self.0.load(Ordering::Relaxed) + fn budget_exhausted<'a>(&'a self, _thread_id: &'a str) -> ExtensionFuture<'a, bool> { + Box::pin(async { *self.0.lock().await }) } } @@ -177,7 +177,7 @@ async fn budget_is_checked_again_after_tool_call(exhausted: bool) -> Result<()> |event| matches!(event, EventMsg::DynamicToolCallRequest(call) if call.call_id == CALL_ID), ) .await; - budget.0.store(exhausted, Ordering::Relaxed); + *budget.0.lock().await = exhausted; thread .submit(Op::DynamicToolResponse { id: CALL_ID.into(), diff --git a/codex-rs/ext/extension-api/src/model_request.rs b/codex-rs/ext/extension-api/src/model_request.rs index a9be4fa25a..aaf1bbf79b 100644 --- a/codex-rs/ext/extension-api/src/model_request.rs +++ b/codex-rs/ext/extension-api/src/model_request.rs @@ -1,4 +1,5 @@ //! Request-scoped interception of model response streams. +use crate::ExtensionFuture; pub use codex_api::ApiError as ModelResponseError; pub use codex_api::GuardianAssessment; pub use codex_api::ResponseEvent; @@ -31,11 +32,11 @@ pub struct ModelRequestInput<'a> { /// Checks model requests and creates request-scoped response interceptors. pub trait ModelRequestContributor: Send + Sync + std::fmt::Debug { - /// Reads the local budget state for a managed thread before submitting generation to the transport. + /// Core awaits this check before submitting generation to the transport. /// Return true to fail with `SessionBudgetExceeded`. Requests already handed to the /// transport can still complete or retry. WebSocket warmup skips this check. - fn budget_exhausted(&self, _thread_id: &str) -> bool { - false + fn budget_exhausted<'a>(&'a self, _thread_id: &'a str) -> ExtensionFuture<'a, bool> { + Box::pin(async { false }) } /// Return None for requests this contributor does not intercept.