codex: preserve goal capacity backoff across resumes (#31176)

codex-rs/app-server/tests/suite/v2/thread_resume.rs

@@ -1761,7 +1761,7 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul
     )
     .await??;
     tokio::time::sleep(Duration::from_millis(100)).await;
-    wait_for_responses_request_count(&server, 3).await?;
+    wait_for_responses_request_count(&server, /*expected_count*/ 3).await?;
 
     Ok(())
 }

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

@@ -4,6 +4,7 @@ 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;
@@ -19,6 +20,11 @@ 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 {
@@ -86,6 +92,9 @@ 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 {
@@ -240,6 +249,7 @@ 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,
@@ -278,6 +288,7 @@ 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);
 
@@ -307,6 +318,47 @@ 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();
@@ -320,4 +372,10 @@ 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,6 +156,12 @@ 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!(
@@ -305,8 +311,12 @@ where
             let reason = match input.error {
                 // Capacity retries do not consume the user's token budget, so
                 // leave the goal active for a delayed idle continuation.
-                CodexErrorInfo::ServerOverloaded => {
-                    runtime.defer_capacity_retry();
+                CodexErrorInfo::ServerOverloaded
+                    if runtime
+                        .accounting_state()
+                        .turn_is_current_active_goal(input.turn_id) =>
+                {
+                    self.goal_service.defer_capacity_retry(runtime.thread_id());
                     return;
                 }
                 CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit,

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

@@ -1,9 +1,7 @@
 use std::sync::Arc;
 use std::sync::Weak;
 use std::sync::atomic::AtomicBool;
-use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
-use std::time::Duration;
 
 use codex_core::ThreadManager;
 use codex_protocol::ThreadId;
@@ -22,10 +20,6 @@ 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>,
@@ -51,7 +45,6 @@ struct GoalRuntimeInner {
     thread_manager: Weak<ThreadManager>,
     accounting_state: Arc<GoalAccountingState>,
     enabled: AtomicBool,
-    capacity_retry_generation: AtomicU64,
     tools_available_for_thread: bool,
     goal_state_lock: Semaphore,
 }
@@ -104,7 +97,6 @@ impl GoalRuntimeHandle {
                 thread_manager,
                 accounting_state,
                 enabled: AtomicBool::new(config.enabled),
-                capacity_retry_generation: AtomicU64::new(0),
                 tools_available_for_thread: config.tools_available_for_thread,
                 goal_state_lock: Semaphore::new(/*permits*/ 1),
             }),
@@ -123,31 +115,6 @@ impl GoalRuntimeHandle {
         self.is_enabled() && self.inner.tools_available_for_thread
     }
 
-    pub(crate) fn defer_capacity_retry(&self) {
-        let generation = self
-            .inner
-            .capacity_retry_generation
-            .fetch_add(1, Ordering::Relaxed)
-            .saturating_add(1);
-
-        let runtime = self.clone();
-        drop(tokio::spawn(async move {
-            tokio::time::sleep(SERVER_OVERLOADED_GOAL_RETRY_DELAY).await;
-            if runtime
-                .inner
-                .capacity_retry_generation
-                .compare_exchange(generation, 0, Ordering::Relaxed, Ordering::Relaxed)
-                .is_ok()
-                && 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
     }
@@ -390,9 +357,6 @@ impl GoalRuntimeHandle {
     }
 
     pub(crate) async fn continue_if_idle(&self) -> Result<(), String> {
-        if self.inner.capacity_retry_generation.load(Ordering::Relaxed) != 0 {
-            return Ok(());
-        }
         if !self.tools_visible() {
             self.inner.accounting_state.clear_active_goal();
             return Ok(());