agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[codex] configure rollout budget reminder thresholds (#29423)

Browse Source 
## Summary

Instead of:

    reminder_interval_tokens = 65_536

allow users to configure explicit remaining-token reminder thresholds:

reminder_at_remaining_tokens = [65_536, 32_768, 16_384, 8_192, 4_096,
2_048, 1_024, 512]

## Validation

- CARGO_INCREMENTAL=0 just test -p codex-core rollout_budget: 9 passed
- just fix -p codex-core
- just fmt
This commit is contained in:
rka-oai
2026-06-22 13:25:48 -07:00
committed by GitHub
Unverified
parent 5a67d898a5
commit bd5bd953fb
8 changed files with 60 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/config.schema.json
+7 -4
View File
@@ -2674,10 +2674,13 @@
"minimum": 0.0,
"type": "number"
},
"reminder_interval_tokens": {
"format": "int64",
"minimum": 1.0,
"type": "integer"
"reminder_at_remaining_tokens": {
"description": "Remaining weighted-token values that trigger reminders when crossed.",
"items": {
"format": "int64",
"type": "integer"
},
"type": "array"
},
"sampling_token_weight": {
"format": "double",
codex-rs/core/src/config/config_tests.rs
+2 -2
View File
@@ -552,7 +552,7 @@ async fn load_config_resolves_rollout_budget() -> std::io::Result<()> {
[features.rollout_budget]
enabled = true
limit_tokens = 100000
reminder_interval_tokens = 10000
reminder_at_remaining_tokens = [50000, 25000, 10000]
sampling_token_weight = 1.0
prefill_token_weight = 0.1
"#,
@@ -571,7 +571,7 @@ prefill_token_weight = 0.1
config.rollout_budget,
Some(RolloutBudgetConfig {
limit_tokens: 100_000,
reminder_interval_tokens: 10_000,
reminder_at_remaining_tokens: vec![50_000, 25_000, 10_000],
sampling_token_weight: 1.0,
prefill_token_weight: 0.1,
})
codex-rs/core/src/config/mod.rs
+18 -8
View File
@@ -1098,10 +1098,10 @@ impl Default for TokenBudgetConfig {
}
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize)]
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct RolloutBudgetConfig {
pub limit_tokens: i64,
pub reminder_interval_tokens: i64,
pub reminder_at_remaining_tokens: Vec<i64>,
pub sampling_token_weight: f64,
pub prefill_token_weight: f64,
}
@@ -2606,13 +2606,23 @@ fn resolve_rollout_budget_config(
"features.rollout_budget.limit_tokens must be positive",
));
}
let reminder_interval_tokens = config
.reminder_interval_tokens
.unwrap_or_else(|| (limit_tokens / 10).max(1));
if reminder_interval_tokens <= 0 {
let reminder_at_remaining_tokens =
config
.reminder_at_remaining_tokens
.clone()
.ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.rollout_budget.reminder_at_remaining_tokens is required when rollout_budget is enabled",
)
})?;
if reminder_at_remaining_tokens
.iter()
.any(|&tokens| tokens <= 0 || tokens >= limit_tokens)
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.rollout_budget.reminder_interval_tokens must be positive",
"features.rollout_budget.reminder_at_remaining_tokens must contain only positive values below limit_tokens",
));
}
let sampling_token_weight = config.sampling_token_weight.unwrap_or(1.0);
@@ -2630,7 +2640,7 @@ fn resolve_rollout_budget_config(
}
Ok(Some(RolloutBudgetConfig {
limit_tokens,
reminder_interval_tokens,
reminder_at_remaining_tokens,
sampling_token_weight,
prefill_token_weight,
}))
codex-rs/core/src/rollout_budget.rs
+9 -5
View File
@@ -57,18 +57,22 @@ impl RolloutBudget {
window_id: &str,
) -> Option<RolloutBudgetReminder> {
let state = self.lock()?;
let reminder_index = (state.weighted_tokens_used
/ state.config.reminder_interval_tokens as f64)
let remaining_tokens = (state.config.limit_tokens as f64 - state.weighted_tokens_used)
.max(0.0)
.floor() as i64;
let reminder_index = state
.config
.reminder_at_remaining_tokens
.iter()
.filter(|&&threshold| remaining_tokens <= threshold)
.count() as i64;
if state.deliveries.get(&thread_id).is_some_and(|delivery| {
delivery.window_id.as_str() == window_id && delivery.reminder_index >= reminder_index
}) {
return None;
}
Some(RolloutBudgetReminder {
remaining_tokens: (state.config.limit_tokens as f64 - state.weighted_tokens_used)
.max(0.0)
.floor() as i64,
remaining_tokens,
reminder_index,
})
}
codex-rs/core/src/session/config_lock.rs
+2 -2
View File
@@ -255,7 +255,7 @@ mod tests {
.expect("token_budget should be enableable in tests");
config.rollout_budget = Some(crate::config::RolloutBudgetConfig {
limit_tokens: 100_000,
reminder_interval_tokens: 10_000,
reminder_at_remaining_tokens: vec![50_000, 25_000, 10_000],
sampling_token_weight: 1.0,
prefill_token_weight: 0.25,
});
@@ -348,7 +348,7 @@ mod tests {
Some(FeatureToml::Config(RolloutBudgetConfigToml {
enabled: Some(true),
limit_tokens: Some(100_000),
reminder_interval_tokens: Some(10_000),
reminder_at_remaining_tokens: Some(vec![50_000, 25_000, 10_000]),
sampling_token_weight: Some(1.0),
prefill_token_weight: Some(0.25),
}))
codex-rs/core/src/thread_manager.rs
+1 -1
View File
@@ -1024,7 +1024,7 @@ impl ThreadManager {
}
fn agent_control_for_config(&self, config: &Config) -> AgentControl {
AgentControl::new(Arc::downgrade(&self.state), config.rollout_budget)
AgentControl::new(Arc::downgrade(&self.state), config.rollout_budget.clone())
}
#[cfg(test)]
codex-rs/core/tests/suite/rollout_budget.rs
+19 -17
View File
@@ -25,12 +25,14 @@ use std::time::Duration;
use test_case::test_case;
use tokio::time::timeout;
const ROLLOUT_BUDGET: RolloutBudgetConfig = RolloutBudgetConfig {
limit_tokens: 100,
reminder_interval_tokens: 25,
sampling_token_weight: 1.0,
prefill_token_weight: 1.0,
};
fn rollout_budget() -> RolloutBudgetConfig {
RolloutBudgetConfig {
limit_tokens: 100,
reminder_at_remaining_tokens: vec![75, 50, 25],
sampling_token_weight: 1.0,
prefill_token_weight: 1.0,
}
}
fn rollout_budget_texts(request: &ResponsesRequest) -> Vec<String> {
request
@@ -51,7 +53,7 @@ fn wire_request_contains(request: &wiremock::Request, text: &str) -> bool {
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn adds_weighted_initial_and_periodic_reminders() -> Result<()> {
async fn adds_weighted_initial_and_threshold_reminders() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
@@ -83,7 +85,7 @@ async fn adds_weighted_initial_and_periodic_reminders() -> Result<()> {
config.rollout_budget = Some(RolloutBudgetConfig {
sampling_token_weight: 2.0,
prefill_token_weight: 0.5,
..ROLLOUT_BUDGET
..rollout_budget()
});
})
.build(&server)
@@ -171,7 +173,7 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> {
.features
.enable(Feature::MultiAgentV2)
.expect("test config should allow multi-agent v2");
config.rollout_budget = Some(ROLLOUT_BUDGET);
config.rollout_budget = Some(rollout_budget());
})
.build(&server)
.await?;
@@ -218,8 +220,8 @@ async fn exhausted_budget_aborts_current_and_later_turns() -> Result<()> {
.with_config(|config| {
config.rollout_budget = Some(RolloutBudgetConfig {
limit_tokens: 30,
reminder_interval_tokens: 10,
..ROLLOUT_BUDGET
reminder_at_remaining_tokens: vec![20, 10],
..rollout_budget()
});
})
.build(&server)
@@ -293,8 +295,8 @@ async fn compaction_budget_exhaustion_aborts_without_error_or_retry(remote_v2: b
config.model_provider = model_provider;
config.rollout_budget = Some(RolloutBudgetConfig {
limit_tokens: 10,
reminder_interval_tokens: 5,
..ROLLOUT_BUDGET
reminder_at_remaining_tokens: vec![5],
..rollout_budget()
});
if remote_v2 {
config
@@ -354,8 +356,8 @@ async fn restates_the_current_remainder_after_compaction() -> Result<()> {
.with_config(move |config| {
config.model_provider = model_provider;
config.rollout_budget = Some(RolloutBudgetConfig {
reminder_interval_tokens: 50,
..ROLLOUT_BUDGET
reminder_at_remaining_tokens: vec![50],
..rollout_budget()
});
})
.build(&server)
@@ -409,8 +411,8 @@ async fn restates_the_current_remainder_after_rollback() -> Result<()> {
let test = test_codex()
.with_config(|config| {
config.rollout_budget = Some(RolloutBudgetConfig {
reminder_interval_tokens: 50,
..ROLLOUT_BUDGET
reminder_at_remaining_tokens: vec![50],
..rollout_budget()
});
})
.build(&server)
codex-rs/features/src/feature_configs.rs
+2 -2
View File
@@ -109,8 +109,8 @@ pub struct RolloutBudgetConfigToml {
#[schemars(range(min = 1))]
pub limit_tokens: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
#[schemars(range(min = 1))]
pub reminder_interval_tokens: Option<i64>,
/// Remaining weighted-token values that trigger reminders when crossed.
pub reminder_at_remaining_tokens: Option<Vec<i64>>,
#[serde(skip_serializing_if = "Option::is_none")]
#[schemars(range(min = 0.0))]
pub sampling_token_weight: Option<f64>,