mirror of
https://github.com/openai/codex.git
synced 2026-11-07 20:30:33 +00:00
rhan-oai
1031a7b9aa
## 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
69 lines
2.1 KiB
Rust
69 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) 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,
|
|
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;
|