codex: scope goal capacity backoff to live runtime (#31176)

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

@@ -4,7 +4,6 @@ use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::PoisonError;
 use std::sync::Weak;
-use std::time::Duration;
 
 use codex_protocol::ThreadId;
 use codex_protocol::protocol::EventMsg;
@@ -20,11 +19,6 @@ use crate::tool::fill_empty_thread_preview_if_possible;
 use crate::tool::protocol_goal_from_state;
 use crate::tool::state_status_from_protocol;
 use crate::tool::validate_goal_budget;
-use tokio::time::Instant;
-
-// Capacity failures do not consume user tokens, but retrying immediately can
-// create a tight loop of failed turns. Keep the retry cadence deliberately low.
-const SERVER_OVERLOADED_GOAL_RETRY_DELAY: Duration = Duration::from_secs(5 * 60);
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum GoalServiceError {
@@ -92,9 +86,6 @@ impl GoalSetOutcome {
 #[derive(Debug, Default)]
 pub struct GoalService {
     runtimes: Mutex<HashMap<String, Weak<GoalRuntimeHandle>>>,
-    // Keep this above the per-thread runtime lifecycle so unloading and
-    // resuming a thread cannot bypass its pending backoff.
-    capacity_retry_deadlines: Mutex<HashMap<String, Instant>>,
 }
 
 impl GoalService {
@@ -249,7 +240,6 @@ impl GoalService {
         if objective.is_some() {
             fill_empty_thread_preview_if_possible(state_db, thread_id, &goal).await;
         }
-        self.clear_capacity_retry(thread_id);
         Ok(GoalSetOutcome {
             goal: protocol_goal_from_state(goal.clone()),
             state_goal: goal,
@@ -288,7 +278,6 @@ impl GoalService {
                 GoalServiceError::Internal(format!("failed to clear thread goal: {err}"))
             })?;
         let cleared = cleared_goal.is_some();
-        self.clear_capacity_retry(thread_id);
         drop(goal_state_permit);
         drop(runtime);
 
@@ -318,47 +307,6 @@ impl GoalService {
         }
     }
 
-    pub(crate) fn defer_capacity_retry(self: &Arc<Self>, thread_id: ThreadId) {
-        let deadline = Instant::now() + SERVER_OVERLOADED_GOAL_RETRY_DELAY;
-        self.capacity_retry_deadlines()
-            .insert(thread_id.to_string(), deadline);
-
-        let service = Arc::downgrade(self);
-        drop(tokio::spawn(async move {
-            tokio::time::sleep_until(deadline).await;
-            let Some(service) = service.upgrade() else {
-                return;
-            };
-            let key = thread_id.to_string();
-            let retry_is_current = {
-                let mut deadlines = service.capacity_retry_deadlines();
-                let retry_is_current = deadlines.get(&key) == Some(&deadline);
-                if retry_is_current {
-                    deadlines.remove(&key);
-                }
-                retry_is_current
-            };
-            if retry_is_current
-                && let Some(runtime) = service.runtime_for_thread(thread_id)
-                && let Err(err) = runtime.continue_if_idle().await
-            {
-                tracing::warn!(
-                    "failed to continue active goal after capacity retry delay for {thread_id}: {err}"
-                );
-            }
-        }));
-    }
-
-    pub(crate) fn capacity_retry_pending(&self, thread_id: ThreadId) -> bool {
-        self.capacity_retry_deadlines()
-            .contains_key(&thread_id.to_string())
-    }
-
-    fn clear_capacity_retry(&self, thread_id: ThreadId) {
-        self.capacity_retry_deadlines()
-            .remove(&thread_id.to_string());
-    }
-
     fn runtime_for_thread(&self, thread_id: ThreadId) -> Option<Arc<GoalRuntimeHandle>> {
         let key = thread_id.to_string();
         let mut runtimes = self.runtimes();
@@ -372,10 +320,4 @@ impl GoalService {
     fn runtimes(&self) -> std::sync::MutexGuard<'_, HashMap<String, Weak<GoalRuntimeHandle>>> {
         self.runtimes.lock().unwrap_or_else(PoisonError::into_inner)
     }
-
-    fn capacity_retry_deadlines(&self) -> std::sync::MutexGuard<'_, HashMap<String, Instant>> {
-        self.capacity_retry_deadlines
-            .lock()
-            .unwrap_or_else(PoisonError::into_inner)
-    }
 }

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

@@ -156,12 +156,6 @@ where
             let Some(runtime) = goal_runtime_handle(input.thread_store) else {
                 return;
             };
-            if self
-                .goal_service
-                .capacity_retry_pending(runtime.thread_id())
-            {
-                return;
-            }
 
             if let Err(err) = runtime.continue_if_idle().await {
                 tracing::warn!(
@@ -316,7 +310,7 @@ where
                         .accounting_state()
                         .turn_is_current_active_goal(input.turn_id) =>
                 {
-                    self.goal_service.defer_capacity_retry(runtime.thread_id());
+                    runtime.defer_capacity_retry();
                     return;
                 }
                 CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit,

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

@@ -2,6 +2,7 @@ use std::sync::Arc;
 use std::sync::Weak;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
+use std::time::Duration;
 
 use codex_core::ThreadManager;
 use codex_protocol::ThreadId;
@@ -20,6 +21,10 @@ use crate::tool::protocol_goal_from_state;
 use tokio::sync::Semaphore;
 use tokio::sync::SemaphorePermit;
 
+// Capacity failures do not consume user tokens, but retrying immediately can
+// create a tight loop of failed turns. Keep the retry cadence deliberately low.
+const SERVER_OVERLOADED_GOAL_RETRY_DELAY: Duration = Duration::from_secs(5 * 60);
+
 #[derive(Clone)]
 pub struct GoalRuntimeHandle {
     inner: Arc<GoalRuntimeInner>,
@@ -45,6 +50,7 @@ struct GoalRuntimeInner {
     thread_manager: Weak<ThreadManager>,
     accounting_state: Arc<GoalAccountingState>,
     enabled: AtomicBool,
+    capacity_retry_pending: AtomicBool,
     tools_available_for_thread: bool,
     goal_state_lock: Semaphore,
 }
@@ -97,6 +103,7 @@ impl GoalRuntimeHandle {
                 thread_manager,
                 accounting_state,
                 enabled: AtomicBool::new(config.enabled),
+                capacity_retry_pending: AtomicBool::new(false),
                 tools_available_for_thread: config.tools_available_for_thread,
                 goal_state_lock: Semaphore::new(/*permits*/ 1),
             }),
@@ -115,6 +122,34 @@ impl GoalRuntimeHandle {
         self.is_enabled() && self.inner.tools_available_for_thread
     }
 
+    pub(crate) fn defer_capacity_retry(&self) {
+        if self
+            .inner
+            .capacity_retry_pending
+            .swap(true, Ordering::Relaxed)
+        {
+            return;
+        }
+
+        // Do not keep an unloaded thread runtime alive just for this timer.
+        // A newly resumed runtime intentionally starts without this backoff.
+        let runtime = Arc::downgrade(&self.inner);
+        drop(tokio::spawn(async move {
+            tokio::time::sleep(SERVER_OVERLOADED_GOAL_RETRY_DELAY).await;
+            let Some(inner) = runtime.upgrade() else {
+                return;
+            };
+            inner.capacity_retry_pending.store(false, Ordering::Relaxed);
+            let runtime = GoalRuntimeHandle { inner };
+            if let Err(err) = runtime.continue_if_idle().await {
+                tracing::warn!(
+                    "failed to continue active goal after capacity retry delay for {}: {err}",
+                    runtime.thread_id()
+                );
+            }
+        }));
+    }
+
     pub(crate) fn thread_id(&self) -> ThreadId {
         self.inner.thread_id
     }
@@ -357,6 +392,9 @@ impl GoalRuntimeHandle {
     }
 
     pub(crate) async fn continue_if_idle(&self) -> Result<(), String> {
+        if self.inner.capacity_retry_pending.load(Ordering::Relaxed) {
+            return Ok(());
+        }
         if !self.tools_visible() {
             self.inner.accounting_state.clear_active_goal();
             return Ok(());