Add configurable goal token budget limits (#37878)

## What changed

- Add `goals.max_goal_token_budget` as a positive-integer configuration setting.
- Use the configured maximum as the default budget for new goals and when `tokenBudget` is reset to `null`.
- Reject goal creation and updates whose token budget exceeds the configured maximum, including requests through goal tools and `thread/goal/set`.
- Respect managed configuration precedence and per-thread configuration overrides.

## Testing

- Cover configuration parsing and managed overrides.
- Cover defaulting, resetting, and rejecting oversized budgets through the goal service, goal tools, and app-server API.

GitOrigin-RevId: f8d7e6418cdc237d454c8cf47bb32ba0d44a60cf

codex-rs/ext/goal/src/api.rs

@@ -54,6 +54,7 @@ pub struct GoalSetRequest<'a> {
     pub objective: GoalObjectiveUpdate<'a>,
     pub status: Option<ThreadGoalStatus>,
     pub token_budget: GoalTokenBudgetUpdate,
+    pub max_goal_token_budget: Option<i64>,
 }
 
 #[derive(Clone, Debug)]
@@ -150,6 +151,7 @@ impl GoalService {
             objective,
             status,
             token_budget,
+            max_goal_token_budget,
         } = request;
         let status = status.map(state_status_from_protocol);
         let objective = match objective {
@@ -158,14 +160,16 @@ impl GoalService {
         };
         let token_budget = match token_budget {
             GoalTokenBudgetUpdate::Keep => None,
-            GoalTokenBudgetUpdate::Set(token_budget) => Some(token_budget),
+            GoalTokenBudgetUpdate::Set(token_budget) => {
+                Some(token_budget.or(max_goal_token_budget))
+            }
         };
 
         if let Some(objective) = objective {
             validate_thread_goal_objective(objective).map_err(GoalServiceError::InvalidRequest)?;
         }
         if objective.is_some() || token_budget.is_some() {
-            validate_goal_budget(token_budget.flatten())
+            validate_goal_budget(token_budget.flatten(), max_goal_token_budget)
                 .map_err(GoalServiceError::InvalidRequest)?;
         }
 
@@ -225,7 +229,7 @@ impl GoalService {
                         thread_id,
                         objective,
                         status.unwrap_or(codex_state::ThreadGoalStatus::Active),
-                        token_budget.flatten(),
+                        token_budget.flatten().or(max_goal_token_budget),
                     )
                     .await
                     .map_err(|err| {

codex-rs/ext/goal/src/extension.rs

@@ -48,12 +48,7 @@ use crate::tool::GoalToolExecutor;
 #[derive(Clone, Debug)]
 pub struct GoalExtensionConfig {
     pub enabled: bool,
-}
-
-impl GoalExtensionConfig {
-    fn from_enabled(enabled: bool) -> Self {
-        Self { enabled }
-    }
+    pub max_goal_token_budget: Option<i64>,
 }
 
 #[derive(Clone)]
@@ -64,7 +59,7 @@ pub struct GoalExtension<C> {
     metrics: GoalMetrics,
     thread_manager: Weak<ThreadManager>,
     goal_service: Arc<GoalService>,
-    goals_enabled: Arc<dyn Fn(&C) -> bool + Send + Sync>,
+    goal_config: Arc<dyn Fn(&C) -> GoalExtensionConfig + Send + Sync>,
 }
 
 impl<C> std::fmt::Debug for GoalExtension<C> {
@@ -81,7 +76,7 @@ impl<C> GoalExtension<C> {
         metrics_client: Option<MetricsClient>,
         thread_manager: Weak<ThreadManager>,
         goal_service: Arc<GoalService>,
-        goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
+        goal_config: impl Fn(&C) -> GoalExtensionConfig + Send + Sync + 'static,
     ) -> Self {
         Self {
             state_dbs,
@@ -90,7 +85,7 @@ impl<C> GoalExtension<C> {
             metrics: GoalMetrics::new(metrics_client),
             thread_manager,
             goal_service,
-            goals_enabled: Arc::new(goals_enabled),
+            goal_config: Arc::new(goal_config),
         }
     }
 }
@@ -101,15 +96,14 @@ where
 {
     fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> {
         Box::pin(async move {
-            let enabled = (self.goals_enabled)(input.config);
+            let config = (self.goal_config)(input.config);
+            let enabled = config.enabled;
             let tools_available_for_thread = input.persistent_thread_state_available
                 && !matches!(
                     input.session_source,
                     SessionSource::SubAgent(SubAgentSource::Review)
                 );
-            input
-                .thread_store
-                .insert(GoalExtensionConfig::from_enabled(enabled));
+            input.thread_store.insert(config);
             let accounting_state = input
                 .thread_store
                 .get_or_init::<GoalAccountingState>(GoalAccountingState::default);
@@ -186,8 +180,9 @@ where
         _previous_config: &C,
         new_config: &C,
     ) {
-        let enabled = (self.goals_enabled)(new_config);
-        thread_store.insert(GoalExtensionConfig::from_enabled(enabled));
+        let config = (self.goal_config)(new_config);
+        let enabled = config.enabled;
+        thread_store.insert(config);
         if let Some(runtime) = goal_runtime_handle(thread_store) {
             runtime.set_enabled(enabled);
         }
@@ -427,6 +422,9 @@ where
         if !runtime.tools_visible() {
             return Vec::new();
         }
+        let max_goal_token_budget = thread_store
+            .get::<GoalExtensionConfig>()
+            .and_then(|config| config.max_goal_token_budget);
 
         vec![
             Arc::new(GoalToolExecutor::get(
@@ -444,6 +442,7 @@ where
                 self.analytics.clone(),
                 self.event_emitter.clone(),
                 self.metrics.clone(),
+                max_goal_token_budget,
             )),
             Arc::new(GoalToolExecutor::update(
                 runtime.thread_id(),
@@ -464,7 +463,7 @@ pub fn install_with_backend<C>(
     metrics_client: Option<MetricsClient>,
     thread_manager: Weak<ThreadManager>,
     goal_service: Arc<GoalService>,
-    goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
+    goal_config: impl Fn(&C) -> GoalExtensionConfig + Send + Sync + 'static,
 ) where
     C: Send + Sync + 'static,
 {
@@ -475,7 +474,7 @@ pub fn install_with_backend<C>(
         metrics_client,
         thread_manager,
         Arc::clone(&goal_service),
-        goals_enabled,
+        goal_config,
     ));
     registry.thread_lifecycle_contributor(extension.clone());
     registry.config_contributor(extension.clone());

codex-rs/ext/goal/src/tool.rs

@@ -36,6 +36,7 @@ pub(crate) struct GoalToolExecutor {
     analytics: GoalAnalytics,
     event_emitter: GoalEventEmitter,
     metrics: GoalMetrics,
+    max_goal_token_budget: Option<i64>,
 }
 
 #[derive(Clone, Copy)]
@@ -89,6 +90,7 @@ impl GoalToolExecutor {
             analytics,
             event_emitter,
             metrics,
+            max_goal_token_budget: None,
         }
     }
 
@@ -99,6 +101,7 @@ impl GoalToolExecutor {
         analytics: GoalAnalytics,
         event_emitter: GoalEventEmitter,
         metrics: GoalMetrics,
+        max_goal_token_budget: Option<i64>,
     ) -> Self {
         Self {
             kind: GoalToolKind::Create,
@@ -108,6 +111,7 @@ impl GoalToolExecutor {
             analytics,
             event_emitter,
             metrics,
+            max_goal_token_budget,
         }
     }
 
@@ -127,6 +131,7 @@ impl GoalToolExecutor {
             analytics,
             event_emitter,
             metrics,
+            max_goal_token_budget: None,
         }
     }
 }
@@ -185,7 +190,9 @@ impl GoalToolExecutor {
         request.objective = request.objective.trim().to_string();
         validate_thread_goal_objective(&request.objective)
             .map_err(FunctionCallError::RespondToModel)?;
-        validate_goal_budget(request.token_budget).map_err(FunctionCallError::RespondToModel)?;
+        request.token_budget = request.token_budget.or(self.max_goal_token_budget);
+        validate_goal_budget(request.token_budget, self.max_goal_token_budget)
+            .map_err(FunctionCallError::RespondToModel)?;
 
         let goal = self
             .state_db
@@ -397,12 +404,23 @@ where
         .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))
 }
 
-pub(crate) fn validate_goal_budget(value: Option<i64>) -> Result<(), String> {
+pub(crate) fn validate_goal_budget(
+    value: Option<i64>,
+    max_goal_token_budget: Option<i64>,
+) -> Result<(), String> {
     if let Some(value) = value
         && value <= 0
     {
         return Err("goal budgets must be positive when provided".to_string());
     }
+    if let Some(value) = value
+        && let Some(max_goal_token_budget) = max_goal_token_budget
+        && value > max_goal_token_budget
+    {
+        return Err(format!(
+            "goal token budget {value} exceeds the maximum allowed goal token budget of {max_goal_token_budget}"
+        ));
+    }
     Ok(())
 }