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