mirror of
https://github.com/openai/codex.git
synced 2026-11-10 21:00:47 +00:00
rhan-oai
c4f0e09912
## 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
70 lines
2.1 KiB
Rust
70 lines
2.1 KiB
Rust
//! 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) async fn check_budget(
|
|
contributors: &[Arc<dyn ModelRequestContributor>],
|
|
thread_id: &str,
|
|
kind: ModelRequestKind,
|
|
) -> Result<()> {
|
|
if kind == ModelRequestKind::Warmup {
|
|
return Ok(());
|
|
}
|
|
for contributor in contributors {
|
|
if contributor.budget_exhausted(thread_id).await {
|
|
return Err(CodexErr::SessionBudgetExceeded);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub(crate) fn prepare(
|
|
contributors: &[Arc<dyn ModelRequestContributor>],
|
|
thread_id: &str,
|
|
model: &str,
|
|
kind: ModelRequestKind,
|
|
metadata: &mut Option<HashMap<String, String>>,
|
|
) -> Vec<Box<dyn ModelResponseInterceptor>> {
|
|
contributors
|
|
.iter()
|
|
.filter_map(|contributor| {
|
|
let mut additions = None;
|
|
let interceptor = contributor.request(ModelRequestInput {
|
|
kind,
|
|
thread_id,
|
|
client_metadata: &mut additions,
|
|
model,
|
|
});
|
|
if let Some(additions) = additions {
|
|
for (key, value) in crate::responses_metadata::filter_extra_metadata(additions) {
|
|
if key != "parent_response_id" {
|
|
metadata.get_or_insert_default().entry(key).or_insert(value);
|
|
}
|
|
}
|
|
}
|
|
interceptor
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
pub(crate) fn intercept_stream(
|
|
mut stream: ModelResponseStream,
|
|
interceptors: Vec<Box<dyn ModelResponseInterceptor>>,
|
|
) -> ModelResponseStream {
|
|
for interceptor in interceptors {
|
|
stream = interceptor.intercept(stream);
|
|
}
|
|
stream
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "model_request_tests.rs"]
|
|
mod tests;
|