From 1031a7b9aa74c1c74561a6133cbaf0349284ad5d Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Fri, 6 Nov 2026 02:14:34 +0000 Subject: [PATCH] Check extension budgets before model generation requests (#59794) ## What changed Add a default `ModelRequestContributor::budget_exhausted` hook and fail generation with `SessionBudgetExceeded` when any contributor reports an exhausted budget for the thread. Check budgets before HTTP requests and before WebSocket connection setup, then check again after the connection is established. WebSocket warmup remains available when the budget is exhausted. ## Testing Add coverage for budget changes during a tool call, verifying that exhaustion prevents the next model request while an available budget allows completion. Extend WebSocket coverage to verify that exhausted generation opens no connection and warmup still permits connection reuse once the budget becomes available. GitOrigin-RevId: 97ce02dc7f9bd54830e4b73c4f0334a30367eebf --- codex-rs/core/src/client.rs | 37 ++++++-- codex-rs/core/src/model_request.rs | 19 +++- .../core/tests/suite/client_websockets.rs | 38 +++++++- .../suite/client_websockets_runtime_auth.rs | 1 + codex-rs/core/tests/suite/model_request.rs | 94 +++++++++++++++++++ .../ext/extension-api/src/model_request.rs | 11 ++- 6 files changed, 186 insertions(+), 14 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index c43378900b..47c780a3a7 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1946,11 +1946,18 @@ impl ResponsesHttpExecutor<'_> { let request_session_telemetry = session_telemetry_for_request(session_telemetry, &request); 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, + )?; let interceptors = crate::model_request::prepare( &self.client.request_contributors, - &self.client.state.thread_id.to_string(), + &thread_id, &model_info.slug, - codex_extension_api::ModelRequestKind::Generation, + kind, &mut request.client_metadata, ); let inference_trace_attempt = inference_trace.start_attempt(); @@ -2057,6 +2064,12 @@ impl ModelClientSession { let after_prewarm = !warmup && self.websocket_session.last_response_from_untraced_warmup; let provider = Arc::clone(&self.client.state.provider); let auth_manager = provider.auth_manager(); + let thread_id = self.client.state.thread_id.to_string(); + let kind = if warmup { + codex_extension_api::ModelRequestKind::Warmup + } else { + codex_extension_api::ModelRequestKind::Generation + }; let mut auth_recovery = auth_manager .as_ref() @@ -2064,6 +2077,11 @@ 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, + )?; let client_setup = self .client .current_client_setup(ClientRouting::Workspace) @@ -2158,6 +2176,13 @@ impl ModelClientSession { Err(err) => return Err(provider.map_api_error(err)), } + // The budget can be exhausted while the connection is being established. + crate::model_request::check_budget( + &self.client.request_contributors, + &thread_id, + kind, + )?; + let mut client_metadata = self.client.build_ws_client_metadata( responses_metadata, include_internal, @@ -2248,13 +2273,9 @@ impl ModelClientSession { ); let interceptors = crate::model_request::prepare( &self.client.request_contributors, - &self.client.state.thread_id.to_string(), + &thread_id, &model_info.slug, - if warmup { - codex_extension_api::ModelRequestKind::Warmup - } else { - codex_extension_api::ModelRequestKind::Generation - }, + kind, &mut ws_payload.client_metadata, ); let mut ws_request = ResponsesWsRequest::ResponseCreate(ws_payload); diff --git a/codex-rs/core/src/model_request.rs b/codex-rs/core/src/model_request.rs index 8b13422bc7..a714806690 100644 --- a/codex-rs/core/src/model_request.rs +++ b/codex-rs/core/src/model_request.rs @@ -1,12 +1,29 @@ -//! Compose request-scoped response interceptors before tracing and response bookkeeping. +//! Check budgets and prepare response interceptors before sending model requests. use codex_extension_api::ModelRequestContributor; use codex_extension_api::ModelRequestInput; use codex_extension_api::ModelRequestKind; use codex_extension_api::ModelResponseInterceptor; use codex_extension_api::ModelResponseStream; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result; use std::collections::HashMap; use std::sync::Arc; +pub(crate) 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); + } + Ok(()) +} + pub(crate) fn prepare( contributors: &[Arc], thread_id: &str, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 1ff75b3cf4..560d5a504c 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -1,4 +1,5 @@ #![allow(clippy::unwrap_used)] +use super::model_request::BudgetCache; use codex_api::WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY; use codex_api::WS_REQUEST_HEADER_TRACESTATE_CLIENT_METADATA_KEY; use codex_core::CodexResponsesMetadata; @@ -12,6 +13,7 @@ use codex_core::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER; use codex_core::test_support::EmptyUserInstructionsProvider; use codex_core::test_support::with_parent_turn; use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ModelRequestContributor; use codex_features::Feature; use codex_http_client::OutboundProxyPolicy; use codex_login::AuthManager; @@ -30,6 +32,7 @@ use codex_protocol::ThreadId; use codex_protocol::account::PlanType; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::ServiceTier; +use codex_protocol::error::CodexErrorDetails; use codex_protocol::models::AudioReference; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; @@ -75,6 +78,7 @@ 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; @@ -738,15 +742,40 @@ 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 harness = websocket_harness_with_provider_options( + let budget = Arc::new(BudgetCache(true.into())); + let harness = websocket_harness_with_provider_factory( provider, /*runtime_metrics_enabled*/ true, /*concurrent_reasoning_summaries_enabled*/ true, /*enabled_features*/ &[], + /*auth*/ None, + vec![budget.clone()], + codex_model_provider::create_model_provider, ) .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 { + panic!("exhausted generation unexpectedly succeeded"); + }; + assert!( + matches!(error.details(), CodexErrorDetails::SessionBudgetExceeded), + "{error:?}" + ); + assert!(server.handshakes().is_empty()); + let responses_metadata = prewarm_metadata(&harness, /*turn_id*/ None); client_session .prewarm_websocket( @@ -760,6 +789,7 @@ async fn responses_websocket_request_prewarm_reuses_connection() { ) .await .expect("websocket prewarm failed"); + budget.0.store(false, Ordering::Relaxed); stream_until_complete(&mut client_session, &harness, &prompt).await; assert_eq!(server.handshakes().len(), 1); @@ -2870,7 +2900,8 @@ async fn websocket_harness_with_provider_options_and_auth( concurrent_reasoning_summaries_enabled, enabled_features, auth, - |provider, auth| codex_model_provider::create_model_provider(provider, auth), + Vec::new(), + codex_model_provider::create_model_provider, ) .await } @@ -2881,6 +2912,7 @@ async fn websocket_harness_with_provider_factory( concurrent_reasoning_summaries_enabled: bool, enabled_features: &[Feature], auth: Option, + request_contributors: Vec>, provider_factory: impl FnOnce( ModelProviderInfo, Option>, @@ -2976,7 +3008,7 @@ async fn websocket_harness_with_provider_factory( /*attestation_provider*/ None, http_client_factory, config.workspace_routing_context(), - Vec::new(), + request_contributors, ); WebsocketTestHarness { 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 6facf10eb1..73c4f8d803 100644 --- a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs +++ b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs @@ -66,6 +66,7 @@ async fn runtime_harness( /*concurrent_reasoning_summaries_enabled*/ false, &[], Some(initial.clone()), + Vec::new(), move |provider, credentials| { let credentials = credentials.expect("runtime fixture requires credentials"); let runtime = Arc::new(RuntimeAuthManager::new( diff --git a/codex-rs/core/tests/suite/model_request.rs b/codex-rs/core/tests/suite/model_request.rs index c72f7f884c..39b83c460d 100644 --- a/codex-rs/core/tests/suite/model_request.rs +++ b/codex-rs/core/tests/suite/model_request.rs @@ -10,6 +10,11 @@ use codex_extension_api::ModelResponseInterceptor; use codex_extension_api::ModelResponseStream; use codex_extension_api::ResponseEvent; use codex_features::Feature; +use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem; +use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; +use codex_protocol::dynamic_tools::DynamicToolResponse; +use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::user_input::UserInput; @@ -24,6 +29,8 @@ 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::Semaphore; use tokio::time::Duration; use tokio::time::timeout; @@ -122,6 +129,93 @@ impl Drop for Permit { } } +#[derive(Debug, Default)] +pub(super) struct BudgetCache(pub(super) AtomicBool); + +impl ModelRequestContributor for BudgetCache { + fn request(&self, _input: ModelRequestInput<'_>) -> Option> { + None + } + + fn budget_exhausted(&self, _thread_id: &str) -> bool { + self.0.load(Ordering::Relaxed) + } +} + +#[tokio::test] +#[test_case::test_case(false; "budget available")] +#[test_case::test_case(true; "budget exhausted")] +async fn budget_is_checked_again_after_tool_call(exhausted: bool) -> Result<()> { + const CALL_ID: &str = "status-call"; + const TOOL_NAME: &str = "check_status"; + const TOOL_OUTPUT: &str = "Release ready."; + + let server = responses::start_mock_server().await; + let model = + responses::mount_function_call_agent_response(&server, CALL_ID, "{}", TOOL_NAME).await; + let budget = Arc::new(BudgetCache::default()); + let mut registry = ExtensionRegistryBuilder::new(); + registry.model_request_contributor(budget.clone()); + let mut builder = test_codex().with_extensions(Arc::new(registry.build())); + let test = builder.build_with_auto_env(&server).await?; + let mut options = test.start_thread_options().await; + options.dynamic_tools = vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: TOOL_NAME.into(), + description: "Check the release status.".into(), + input_schema: json!({"type": "object", "properties": {}}), + defer_loading: false, + })]; + let thread = test.thread_manager.start_thread(options).await?.thread; + thread + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Check the release status.".into(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event( + &thread, + |event| matches!(event, EventMsg::DynamicToolCallRequest(call) if call.call_id == CALL_ID), + ) + .await; + budget.0.store(exhausted, Ordering::Relaxed); + thread + .submit(Op::DynamicToolResponse { + id: CALL_ID.into(), + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: TOOL_OUTPUT.into(), + }], + success: true, + }, + }) + .await?; + let EventMsg::TurnComplete(completed) = + wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await + else { + unreachable!(); + }; + let request_count = responses::received_responses_requests(&server).await.len(); + if exhausted { + assert_eq!( + completed.error.and_then(|error| error.codex_error_info), + Some(CodexErrorInfo::SessionBudgetExceeded) + ); + assert_eq!(request_count, 1); + } else { + assert_eq!(completed.error, None); + assert_eq!(completed.last_agent_message.as_deref(), Some("done")); + assert_eq!(request_count, 2); + assert_eq!( + model + .completion + .single_request() + .function_call_output_content_and_success(CALL_ID), + Some((Some(TOOL_OUTPUT.into()), None)) + ); + } + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn model_request_gates_output_and_cancellation() -> Result<()> { let server = responses::start_mock_server().await; diff --git a/codex-rs/ext/extension-api/src/model_request.rs b/codex-rs/ext/extension-api/src/model_request.rs index 0f477aa99e..a9be4fa25a 100644 --- a/codex-rs/ext/extension-api/src/model_request.rs +++ b/codex-rs/ext/extension-api/src/model_request.rs @@ -29,9 +29,16 @@ pub struct ModelRequestInput<'a> { pub model: &'a str, } -/// Creates request-scoped interceptors without delaying inference. -/// Return None for requests this host does not manage. +/// 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. + /// 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 + } + + /// Return None for requests this contributor does not intercept. fn request(&self, input: ModelRequestInput<'_>) -> Option>; }