codex_wrapper/
rollout_budget.rs

1//! Native per-execution rollout-budget configuration.
2//!
3//! This is deliberately separate from [`crate::TokenBudget`]. That type sees
4//! usage only after a complete CLI turn and can refuse the next turn; this
5//! one asks Codex itself to stop an in-progress `exec` at a response boundary.
6//!
7//! The native meter is not portable total-token usage. Codex 0.145 and 0.146
8//! always compute `output_tokens * sampling_token_weight` plus non-cached
9//! input tokens times `prefill_token_weight`. Starting with 0.147, Codex
10//! prefers a provider-reported `codex_rollout_budget_units` value when one is
11//! available and otherwise uses that weighted fallback. Cached input is not
12//! included in the fallback. Provider-reported units may have different,
13//! opaque semantics, so hosts must not treat a CLI upgrade as an unchanged
14//! portable meter. One response can cross the limit before Codex observes it.
15
16use crate::error::{Error, Result};
17
18/// A validated Codex-native rollout budget for one CLI execution.
19///
20/// Use [`RolloutBudgetConfig::builder`] and pass the result to
21/// [`crate::ExecCommand::rollout_budget`] or
22/// [`crate::ExecResumeCommand::rollout_budget`]. The same config shape is
23/// accepted by both opening and resumed `exec` commands.
24#[derive(Debug, Clone, PartialEq)]
25pub struct RolloutBudgetConfig {
26    limit_tokens: u64,
27    reminder_at_remaining_tokens: Vec<u64>,
28    sampling_token_weight: f64,
29    prefill_token_weight: f64,
30}
31
32impl RolloutBudgetConfig {
33    /// Start a builder with the native rollout-budget-unit limit.
34    #[must_use]
35    pub fn builder(limit_tokens: u64) -> RolloutBudgetConfigBuilder {
36        RolloutBudgetConfigBuilder {
37            limit_tokens,
38            reminder_at_remaining_tokens: Vec::new(),
39            sampling_token_weight: 1.0,
40            prefill_token_weight: 1.0,
41        }
42    }
43
44    /// The configured native rollout-budget-unit limit.
45    #[must_use]
46    pub fn limit_tokens(&self) -> u64 {
47        self.limit_tokens
48    }
49
50    /// Remaining rollout-budget-unit thresholds that make Codex restate the budget.
51    #[must_use]
52    pub fn reminder_at_remaining_tokens(&self) -> &[u64] {
53        &self.reminder_at_remaining_tokens
54    }
55
56    /// Weight applied to generated output tokens when provider units are absent.
57    #[must_use]
58    pub fn sampling_token_weight(&self) -> f64 {
59        self.sampling_token_weight
60    }
61
62    /// Weight applied to non-cached input tokens when provider units are absent.
63    #[must_use]
64    pub fn prefill_token_weight(&self) -> f64 {
65        self.prefill_token_weight
66    }
67
68    pub(crate) fn config_override(&self) -> String {
69        let reminders = self
70            .reminder_at_remaining_tokens
71            .iter()
72            .map(u64::to_string)
73            .collect::<Vec<_>>()
74            .join(",");
75        format!(
76            "features.rollout_budget={{enabled=true,limit_tokens={},reminder_at_remaining_tokens=[{}],sampling_token_weight={},prefill_token_weight={}}}",
77            self.limit_tokens, reminders, self.sampling_token_weight, self.prefill_token_weight,
78        )
79    }
80
81    pub(crate) fn is_config_override(value: &str) -> bool {
82        value.starts_with("features.rollout_budget={enabled=true,limit_tokens=")
83    }
84}
85
86/// Builder for [`RolloutBudgetConfig`].
87#[derive(Debug, Clone)]
88pub struct RolloutBudgetConfigBuilder {
89    limit_tokens: u64,
90    reminder_at_remaining_tokens: Vec<u64>,
91    sampling_token_weight: f64,
92    prefill_token_weight: f64,
93}
94
95impl RolloutBudgetConfigBuilder {
96    /// Replace the remaining-unit thresholds that make Codex restate the budget.
97    ///
98    /// An empty list is valid and disables threshold reminders. Codex still
99    /// includes the initial remaining-budget message and enforces the limit.
100    #[must_use]
101    pub fn reminder_at_remaining_tokens(
102        mut self,
103        thresholds: impl IntoIterator<Item = u64>,
104    ) -> Self {
105        self.reminder_at_remaining_tokens = thresholds.into_iter().collect();
106        self
107    }
108
109    /// Set the weight for generated output tokens when provider units are absent.
110    #[must_use]
111    pub fn sampling_token_weight(mut self, weight: f64) -> Self {
112        self.sampling_token_weight = weight;
113        self
114    }
115
116    /// Set the weight for non-cached input tokens when provider units are absent.
117    #[must_use]
118    pub fn prefill_token_weight(mut self, weight: f64) -> Self {
119        self.prefill_token_weight = weight;
120        self
121    }
122
123    /// Validate and build the native rollout-budget config.
124    pub fn build(self) -> Result<RolloutBudgetConfig> {
125        if self.limit_tokens == 0 || self.limit_tokens > i64::MAX as u64 {
126            return Err(invalid("limit_tokens must be in 1..=i64::MAX"));
127        }
128        if self
129            .reminder_at_remaining_tokens
130            .iter()
131            .any(|&threshold| threshold == 0 || threshold >= self.limit_tokens)
132        {
133            return Err(invalid(
134                "reminder thresholds must be positive and below limit_tokens",
135            ));
136        }
137        for (field, weight) in [
138            ("sampling_token_weight", self.sampling_token_weight),
139            ("prefill_token_weight", self.prefill_token_weight),
140        ] {
141            if !weight.is_finite() || weight < 0.0 {
142                return Err(invalid(format!("{field} must be finite and non-negative")));
143            }
144        }
145        Ok(RolloutBudgetConfig {
146            limit_tokens: self.limit_tokens,
147            reminder_at_remaining_tokens: self.reminder_at_remaining_tokens,
148            sampling_token_weight: self.sampling_token_weight,
149            prefill_token_weight: self.prefill_token_weight,
150        })
151    }
152}
153
154fn invalid(message: impl Into<String>) -> Error {
155    Error::InvalidRolloutBudget {
156        message: message.into(),
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    #[test]
165    fn valid_config_serializes_as_one_strict_cli_override() {
166        let budget = RolloutBudgetConfig::builder(100_000)
167            .reminder_at_remaining_tokens([50_000, 10_000])
168            .sampling_token_weight(1.5)
169            .prefill_token_weight(0.25)
170            .build()
171            .expect("valid budget");
172
173        assert_eq!(budget.limit_tokens(), 100_000);
174        assert_eq!(budget.reminder_at_remaining_tokens(), [50_000, 10_000]);
175        assert_eq!(
176            budget.config_override(),
177            "features.rollout_budget={enabled=true,limit_tokens=100000,reminder_at_remaining_tokens=[50000,10000],sampling_token_weight=1.5,prefill_token_weight=0.25}"
178        );
179    }
180
181    #[test]
182    fn invalid_limits_thresholds_and_weights_fail_before_launch() {
183        for result in [
184            RolloutBudgetConfig::builder(0).build(),
185            RolloutBudgetConfig::builder(i64::MAX as u64 + 1).build(),
186            RolloutBudgetConfig::builder(100)
187                .reminder_at_remaining_tokens([0])
188                .build(),
189            RolloutBudgetConfig::builder(100)
190                .reminder_at_remaining_tokens([100])
191                .build(),
192            RolloutBudgetConfig::builder(100)
193                .sampling_token_weight(f64::NAN)
194                .build(),
195            RolloutBudgetConfig::builder(100)
196                .prefill_token_weight(-1.0)
197                .build(),
198        ] {
199            assert!(matches!(result, Err(Error::InvalidRolloutBudget { .. })));
200        }
201    }
202}
codex_wrapper/rollout_budget.rs

codex_wrapper/
rollout_budget.rs
