[codex] add rollout token budget configuration (varlength 1/N) (#28746)

## What

This PR defines the structured configuration contract for shared rollout
token budgets (across ALL agent threads under 1 rollout).

```toml
[features.rollout_budget]
enabled = true
limit_tokens = 100000
reminder_interval_tokens = 10000
sampling_token_weight = 1.0
prefill_token_weight = 0.1
```

The reminder interval defaults to 10% of the rollout limit. Sampling and
prefill weights default to `1.0`.

## Scope

This PR only defines and validates configuration. It does not track
usage, inject reminders, or stop a rollout. Accounting and reminders are
implemented in the stacked follow-up #28494.

The existing `token_budget` feature remains unchanged. `rollout_budget`
has its own feature key and configuration type.

## Tests

The config test verifies that the structured fields resolve into
`RolloutBudgetConfig` and do not enable the existing `token_budget`
feature.

Local checks:

- `just write-config-schema`
- `just test -p codex-core load_config_resolves_rollout_budget`
- `cargo check -p codex-thread-manager-sample`
- `git diff --check`

The full workspace test suite was not run locally.
This commit is contained in:
rka-oai
2026-06-18 04:29:47 -07:00
committed by GitHub
Unverified
parent c73296a0f0
commit ecc4c30e28
8 changed files with 264 additions and 0 deletions
+61
View File
@@ -461,6 +461,67 @@ direct_only_tool_namespaces = ["mcp__history", "mcp__notes"]
Ok(())
}
#[tokio::test]
async fn load_config_resolves_rollout_budget() -> std::io::Result<()> {
let codex_home = tempdir()?;
let config_toml: ConfigToml = toml::from_str(
r#"
[features.rollout_budget]
enabled = true
limit_tokens = 100000
reminder_interval_tokens = 10000
sampling_token_weight = 1.0
prefill_token_weight = 0.1
"#,
)
.expect("TOML deserialization should succeed");
let config = Config::load_from_base_config_with_overrides(
config_toml,
ConfigOverrides::default(),
codex_home.abs(),
)
.await?;
assert!(config.features.enabled(Feature::RolloutBudget));
assert!(!config.features.enabled(Feature::TokenBudget));
assert_eq!(
config.rollout_budget,
Some(RolloutBudgetConfig {
limit_tokens: 100_000,
reminder_interval_tokens: 10_000,
sampling_token_weight: 1.0,
prefill_token_weight: 0.1,
})
);
Ok(())
}
#[tokio::test]
async fn load_config_rejects_enabled_rollout_budget_without_limit() -> std::io::Result<()> {
for config_toml in [
"[features]\nrollout_budget = true\n",
"[features.rollout_budget]\nenabled = true\n",
] {
let codex_home = tempdir()?;
let config_toml: ConfigToml =
toml::from_str(config_toml).expect("TOML deserialization should succeed");
let err = Config::load_from_base_config_with_overrides(
config_toml,
ConfigOverrides::default(),
codex_home.abs(),
)
.await
.expect_err("enabled rollout budget without limit_tokens should be rejected");
assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
assert_eq!(
err.to_string(),
"features.rollout_budget.limit_tokens is required when rollout_budget is enabled"
);
}
Ok(())
}
#[test]
fn rejects_provider_auth_with_env_key() {
let err = toml::from_str::<ConfigToml>(
+72
View File
@@ -1020,6 +1020,9 @@ pub struct Config {
/// Settings specific to the task-path-based multi-agent tool surface.
pub multi_agent_v2: MultiAgentV2Config,
/// Shared token budget for the root thread and its sub-agents.
pub rollout_budget: Option<RolloutBudgetConfig>,
/// Centralized feature flags; source of truth for feature gating.
pub features: ManagedFeatures,
@@ -1064,6 +1067,14 @@ pub struct CodeModeConfig {
pub direct_only_tool_namespaces: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize)]
pub struct RolloutBudgetConfig {
pub limit_tokens: i64,
pub reminder_interval_tokens: i64,
pub sampling_token_weight: f64,
pub prefill_token_weight: f64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct MultiAgentV2Config {
pub max_concurrent_threads_per_session: usize,
@@ -2469,6 +2480,65 @@ fn resolve_multi_agent_v2_config(config_toml: &ConfigToml) -> MultiAgentV2Config
}
}
fn resolve_rollout_budget_config(
config_toml: &ConfigToml,
features: &ManagedFeatures,
) -> std::io::Result<Option<RolloutBudgetConfig>> {
if !features.enabled(Feature::RolloutBudget) {
return Ok(None);
}
let missing_limit_error = || {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.rollout_budget.limit_tokens is required when rollout_budget is enabled",
)
};
let Some(FeatureToml::Config(config)) = config_toml
.features
.as_ref()
.and_then(|features| features.rollout_budget.as_ref())
else {
return Err(missing_limit_error());
};
let Some(limit_tokens) = config.limit_tokens else {
return Err(missing_limit_error());
};
if limit_tokens <= 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"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 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.rollout_budget.reminder_interval_tokens must be positive",
));
}
let sampling_token_weight = config.sampling_token_weight.unwrap_or(1.0);
let prefill_token_weight = config.prefill_token_weight.unwrap_or(1.0);
for (field, weight) in [
("sampling_token_weight", sampling_token_weight),
("prefill_token_weight", prefill_token_weight),
] {
if !weight.is_finite() || weight < 0.0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("features.rollout_budget.{field} must be finite and non-negative"),
));
}
}
Ok(Some(RolloutBudgetConfig {
limit_tokens,
reminder_interval_tokens,
sampling_token_weight,
prefill_token_weight,
}))
}
fn resolve_terminal_resize_reflow_config(config_toml: &ConfigToml) -> TerminalResizeReflowConfig {
let Some(tui) = config_toml.tui.as_ref() else {
return TerminalResizeReflowConfig::default();
@@ -3146,6 +3216,7 @@ impl Config {
resolve_experimental_request_user_input_enabled(&cfg);
let code_mode = resolve_code_mode_config(&cfg);
let multi_agent_v2 = resolve_multi_agent_v2_config(&cfg);
let rollout_budget = resolve_rollout_budget_config(&cfg, &features)?;
let terminal_resize_reflow = resolve_terminal_resize_reflow_config(&cfg);
let agent_roles =
@@ -3686,6 +3757,7 @@ impl Config {
background_terminal_max_timeout,
ghost_snapshot,
multi_agent_v2,
rollout_budget,
features,
suppress_unstable_features_warning: cfg
.suppress_unstable_features_warning
+30
View File
@@ -6,6 +6,7 @@ use codex_features::Feature;
use codex_features::FeatureToml;
use codex_features::FeaturesToml;
use codex_features::MultiAgentV2ConfigToml;
use codex_features::RolloutBudgetConfigToml;
use codex_protocol::ThreadId;
use crate::config::Config;
@@ -148,6 +149,12 @@ fn save_config_resolved_fields(
resolved_config_to_toml(&config.multi_agent_v2, "features.multi_agent_v2")?;
multi_agent_v2.enabled = Some(config.features.enabled(Feature::MultiAgentV2));
features.multi_agent_v2 = Some(FeatureToml::Config(multi_agent_v2));
if let Some(rollout_budget) = config.rollout_budget.as_ref() {
let mut rollout_budget: RolloutBudgetConfigToml =
resolved_config_to_toml(rollout_budget, "features.rollout_budget")?;
rollout_budget.enabled = Some(config.features.enabled(Feature::RolloutBudget));
features.rollout_budget = Some(FeatureToml::Config(rollout_budget));
}
lock_config.memories = Some(resolved_config_to_toml::<MemoriesToml>(
&config.memories,
"memories",
@@ -209,6 +216,18 @@ mod tests {
#[tokio::test]
async fn lock_contains_prompts_and_materializes_features() {
let mut sc = crate::session::tests::make_session_configuration_for_tests().await;
let mut config = (*sc.original_config_do_not_use).clone();
config.rollout_budget = Some(crate::config::RolloutBudgetConfig {
limit_tokens: 100_000,
reminder_interval_tokens: 10_000,
sampling_token_weight: 1.0,
prefill_token_weight: 0.25,
});
config
.features
.enable(Feature::RolloutBudget)
.expect("rollout_budget should be enableable in tests");
sc.original_config_do_not_use = Arc::new(config);
sc.base_instructions = "resolved instructions".to_string();
sc.developer_instructions = Some("resolved developer instructions".to_string());
sc.compact_prompt = Some("resolved compact prompt".to_string());
@@ -273,6 +292,17 @@ mod tests {
})
));
assert_eq!(
features.rollout_budget,
Some(FeatureToml::Config(RolloutBudgetConfigToml {
enabled: Some(true),
limit_tokens: Some(100_000),
reminder_interval_tokens: Some(10_000),
sampling_token_weight: Some(1.0),
prefill_token_weight: Some(0.25),
}))
);
assert_eq!(lockfile.version, crate::config_lock::CONFIG_LOCK_VERSION);
}