openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-06 20:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Check extension budgets before model generation requests (#59794)

Browse Source 
## 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
This commit is contained in:
rhan-oai
2026-11-06 02:31:54 +00:00
committed by copyberry
parent 2f02d8461d
commit 1031a7b9aa
6 changed files with 186 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/client.rs
+29 -8
View File
@@ -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);
codex-rs/core/src/model_request.rs
+18 -1
View File
@@ -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<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);
}
Ok(())
}
pub(crate) fn prepare(
contributors: &[Arc<dyn ModelRequestContributor>],
thread_id: &str,
codex-rs/core/tests/suite/client_websockets.rs
+35 -3
View File
@@ -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<CodexAuth>,
request_contributors: Vec<Arc<dyn ModelRequestContributor>>,
provider_factory: impl FnOnce(
ModelProviderInfo,
Option<Arc<AuthManager>>,
@@ -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 {
codex-rs/core/tests/suite/client_websockets_runtime_auth.rs
+1
View File
@@ -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(
codex-rs/core/tests/suite/model_request.rs
+94
View File
@@ -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<Box<dyn ModelResponseInterceptor>> {
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;
codex-rs/ext/extension-api/src/model_request.rs
+9 -2
View File
@@ -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<Box<dyn ModelResponseInterceptor>>;
}