Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions codex-rs/core/config.schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
4 changes: 2 additions & 2 deletions codex-rs/core/src/config/config_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
"#,
Expand All @@ -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,
})
Expand Down
26 changes: 18 additions & 8 deletions codex-rs/core/src/config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1097,10 +1097,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,
}
Expand Down Expand Up @@ -2596,13 +2596,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);
Expand All @@ -2620,7 +2630,7 @@ fn resolve_rollout_budget_config(
}
Ok(Some(RolloutBudgetConfig {
limit_tokens,
reminder_interval_tokens,
reminder_at_remaining_tokens,
sampling_token_weight,
prefill_token_weight,
}))
Expand Down
14 changes: 9 additions & 5 deletions codex-rs/core/src/rollout_budget.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
})
}
Expand Down
4 changes: 2 additions & 2 deletions codex-rs/core/src/session/config_lock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
});
Expand Down Expand Up @@ -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),
}))
Expand Down
2 changes: 1 addition & 1 deletion codex-rs/core/src/thread_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
36 changes: 19 additions & 17 deletions codex-rs/core/tests/suite/rollout_budget.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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?;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions codex-rs/features/src/feature_configs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>>,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve rollout budget config compatibility

When users have an existing rollout-budget config or exported config-lock file, changing the TOML shape to only accept reminder_at_remaining_tokens while RolloutBudgetConfigToml still denies unknown fields makes prior reminder_interval_tokens configs fail to deserialize before migration; configs that only set limit_tokens also now fail because the resolver requires this new field instead of deriving the previous default. Please keep a deprecated alias/translation or fallback so existing rollout-budget sessions/configs continue to start.

AGENTS.md reference: AGENTS.md:L103-L111

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not an issue

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Cap rollout budget threshold lists

When a managed or user config supplies a large/dense reminder_at_remaining_tokens array, each newly crossed index can append another <rollout_budget> developer-context fragment over the life of the thread, but this new list has no length cap or deduplication before being stored and scanned on every request. Please enforce a small maximum (and ideally normalize duplicates) so the number of injected rollout-budget fragments remains bounded.

AGENTS.md reference: AGENTS.md:L98-L98

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not an issue

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3 Badge Constrain threshold values in the schema

The resolver rejects zero or negative reminder_at_remaining_tokens, but the generated schema for this new array now advertises each item as an unconstrained integer. Users relying on config.schema.json can therefore get editor/CI validation success for configs that Codex later rejects at startup, so the item schema should carry the same positive-value constraint as the runtime check.

Useful? React with 👍 / 👎.

#[serde(skip_serializing_if = "Option::is_none")]
#[schemars(range(min = 0.0))]
pub sampling_token_weight: Option<f64>,
Expand Down
Loading