diff --git a/codex-rs/config/src/schema.rs b/codex-rs/config/src/schema.rs index c641f1703..6d3f901bf 100644 --- a/codex-rs/config/src/schema.rs +++ b/codex-rs/config/src/schema.rs @@ -44,6 +44,15 @@ pub fn features_schema(schema_gen: &mut SchemaGenerator) -> Schema { ); continue; } + if feature.id == codex_features::Feature::RolloutBudget { + validation.properties.insert( + feature.key.to_string(), + schema_gen.subschema_for::>(), + ); + continue; + } if feature.id == codex_features::Feature::AppsMcpPathOverride { validation.properties.insert( feature.key.to_string(), diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 148901a70..852f6de39 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -611,6 +611,9 @@ "responses_websockets_v2": { "type": "boolean" }, + "rollout_budget": { + "$ref": "#/definitions/FeatureToml_for_RolloutBudgetConfigToml" + }, "runtime_metrics": { "type": "boolean" }, @@ -908,6 +911,16 @@ } ] }, + "FeatureToml_for_RolloutBudgetConfigToml": { + "anyOf": [ + { + "type": "boolean" + }, + { + "$ref": "#/definitions/RolloutBudgetConfigToml" + } + ] + }, "FeedbackConfigToml": { "additionalProperties": false, "properties": { @@ -2572,6 +2585,35 @@ } ] }, + "RolloutBudgetConfigToml": { + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean" + }, + "limit_tokens": { + "format": "int64", + "minimum": 1.0, + "type": "integer" + }, + "prefill_token_weight": { + "format": "double", + "minimum": 0.0, + "type": "number" + }, + "reminder_interval_tokens": { + "format": "int64", + "minimum": 1.0, + "type": "integer" + }, + "sampling_token_weight": { + "format": "double", + "minimum": 0.0, + "type": "number" + } + }, + "type": "object" + }, "SandboxMode": { "enum": [ "read-only", @@ -4772,6 +4814,9 @@ "responses_websockets_v2": { "type": "boolean" }, + "rollout_budget": { + "$ref": "#/definitions/FeatureToml_for_RolloutBudgetConfigToml" + }, "runtime_metrics": { "type": "boolean" }, diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index dd6169fbb..2244e3e1b 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -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::( diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index a54137b58..042761336 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -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, + /// 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, } +#[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> { + 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 diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index ec2056581..b5f1689ad 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/codex-rs/core/src/session/config_lock.rs @@ -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::( &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); } diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index 0acb9b7d3..ad4b6da63 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -73,6 +73,35 @@ impl FeatureConfig for MultiAgentV2ConfigToml { } } +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct RolloutBudgetConfigToml { + #[serde(skip_serializing_if = "Option::is_none")] + pub enabled: Option, + #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(range(min = 1))] + pub limit_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(range(min = 1))] + pub reminder_interval_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(range(min = 0.0))] + pub sampling_token_weight: Option, + #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(range(min = 0.0))] + pub prefill_token_weight: Option, +} + +impl FeatureConfig for RolloutBudgetConfigToml { + fn enabled(&self) -> Option { + self.enabled + } + + fn set_enabled(&mut self, enabled: bool) { + self.enabled = Some(enabled); + } +} + #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] #[serde(deny_unknown_fields)] pub(crate) struct RemovedAppsMcpPathOverrideConfigToml { diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index b887237a1..731ea20a1 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -23,6 +23,7 @@ pub use feature_configs::NetworkProxyDomainPermissionToml; pub use feature_configs::NetworkProxyModeToml; pub use feature_configs::NetworkProxyUnixSocketPermissionToml; use feature_configs::RemovedAppsMcpPathOverrideConfigToml; +pub use feature_configs::RolloutBudgetConfigToml; use legacy::LegacyFeatureToggles; pub use legacy::legacy_feature_keys; @@ -203,6 +204,8 @@ pub enum Feature { Goals, /// Add current context-window metadata to model-visible context. TokenBudget, + /// Track and report a shared token budget across a session's agent threads. + RolloutBudget, /// Expose an input-interruptible sleep tool. SleepTool, /// Route MCP tool approval prompts through the MCP elicitation request path. @@ -618,6 +621,8 @@ pub struct FeaturesToml { pub code_mode: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub multi_agent_v2: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rollout_budget: Option>, #[serde(default, rename = "apps_mcp_path_override", skip_serializing)] #[schemars(skip)] removed_apps_mcp_path_override: Option>, @@ -650,6 +655,9 @@ impl FeaturesToml { if let Some(enabled) = self.multi_agent_v2.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::MultiAgentV2.key().to_string(), enabled); } + if let Some(enabled) = self.rollout_budget.as_ref().and_then(FeatureToml::enabled) { + entries.insert(Feature::RolloutBudget.key().to_string(), enabled); + } if let Some(enabled) = self.network_proxy.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::NetworkProxy.key().to_string(), enabled); } @@ -661,6 +669,7 @@ impl FeaturesToml { let Self { code_mode, multi_agent_v2, + rollout_budget, removed_apps_mcp_path_override: _, network_proxy, entries, @@ -674,6 +683,8 @@ impl FeaturesToml { materialize_resolved_feature_enabled(code_mode, enabled); } else if spec.id == Feature::MultiAgentV2 { materialize_resolved_feature_enabled(multi_agent_v2, enabled); + } else if spec.id == Feature::RolloutBudget { + materialize_resolved_feature_enabled(rollout_budget, enabled); } else if spec.id == Feature::NetworkProxy { materialize_resolved_feature_enabled(network_proxy, enabled); } else { @@ -1167,6 +1178,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::RolloutBudget, + key: "rollout_budget", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::SleepTool, key: "sleep_tool", diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index e2d803ce7..ca11ff907 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -277,6 +277,7 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R background_terminal_max_timeout: 300_000, ghost_snapshot: GhostSnapshotConfig::default(), multi_agent_v2: MultiAgentV2Config::default(), + rollout_budget: None, features: Default::default(), suppress_unstable_features_warning: false, active_project: ProjectConfig { trust_level: None },