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

@@ -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,

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<dyn ModelRequestContributor>],
     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(())
 }

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);

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<CodexAuth>);
 
+#[derive(Debug, Default)]
+struct BudgetChecks(Mutex<VecDeque<Option<Arc<Barrier>>>>);
+
+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<Box<dyn ModelResponseInterceptor>> {
+        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<Arc<dyn ModelRequestContributor>>,
 ) -> (WebsocketTestHarness, Arc<Credentials>) {
     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(&registration, provider, mode).await;
+    let budget = Arc::new(BudgetChecks::default());
+    let (harness, credentials) =
+        runtime_harness(&registration, 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();

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<bool>);
 
 impl ModelRequestContributor for BudgetCache {
     fn request(&self, _input: ModelRequestInput<'_>) -> Option<Box<dyn ModelResponseInterceptor>> {
         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(),

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.