codex: address PR review feedback (#31176)

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

@@ -1538,6 +1538,11 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul
             responses::ev_response_created("goal-continuation"),
             responses::ev_completed_with_tokens("goal-continuation", /*total_tokens*/ 200),
         ]),
+        responses::sse_failed(
+            "capacity-error",
+            "server_is_overloaded",
+            "Selected model is at capacity. Please try a different model.",
+        ),
     ])
     .await;
     let codex_home = TempDir::new()?;
@@ -1718,6 +1723,46 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul
     let clear_again: ThreadGoalClearResponse = to_response(clear_again_resp)?;
     assert!(!clear_again.cleared);
 
+    let retry_goal_id = mcp
+        .send_raw_request(
+            "thread/goal/set",
+            Some(json!({
+                "threadId": thread.id,
+                "objective": "retry after capacity recovers",
+            })),
+        )
+        .await?;
+    let retry_goal_resp: JSONRPCResponse = timeout(
+        DEFAULT_READ_TIMEOUT,
+        mcp.read_stream_until_response_message(RequestId::Integer(retry_goal_id)),
+    )
+    .await??;
+    let _retry_goal: ThreadGoalSetResponse = to_response(retry_goal_resp)?;
+    timeout(
+        DEFAULT_READ_TIMEOUT,
+        mcp.read_stream_until_notification_message("error"),
+    )
+    .await??;
+    timeout(
+        DEFAULT_READ_TIMEOUT,
+        mcp.read_stream_until_notification_message("turn/completed"),
+    )
+    .await??;
+
+    let resume_id = mcp
+        .send_thread_resume_request(ThreadResumeParams {
+            thread_id: thread.id,
+            ..Default::default()
+        })
+        .await?;
+    let _resume_resp: JSONRPCResponse = timeout(
+        DEFAULT_READ_TIMEOUT,
+        mcp.read_stream_until_response_message(RequestId::Integer(resume_id)),
+    )
+    .await??;
+    tokio::time::sleep(Duration::from_millis(100)).await;
+    wait_for_responses_request_count(&server, 3).await?;
+
     Ok(())
 }
 

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

@@ -1,6 +1,5 @@
 use std::sync::Arc;
 use std::sync::Weak;
-use std::time::Duration;
 
 use codex_analytics::AnalyticsEventsClient;
 use codex_core::ThreadManager;
@@ -46,13 +45,6 @@ use crate::spec::UPDATE_GOAL_TOOL_NAME;
 use crate::steering::budget_limit_steering_item;
 use crate::tool::GoalToolExecutor;
 
-// 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);
-
-// A thread-scoped one-shot marker consumed by the next idle lifecycle callback.
-struct DeferredCapacityRetry;
-
 #[derive(Clone, Debug)]
 pub struct GoalExtensionConfig {
     pub enabled: bool,
@@ -165,26 +157,6 @@ where
                 return;
             };
 
-            if input
-                .thread_store
-                .remove::<DeferredCapacityRetry>()
-                .is_some()
-            {
-                // Detach the delay so terminal turn cleanup is not held open.
-                // continue_if_idle rechecks the goal, thread, and queued work,
-                // so a stale wakeup after pause, completion, or user input is harmless.
-                drop(tokio::spawn(async move {
-                    tokio::time::sleep(SERVER_OVERLOADED_GOAL_RETRY_DELAY).await;
-                    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()
-                        );
-                    }
-                }));
-                return;
-            }
-
             if let Err(err) = runtime.continue_if_idle().await {
                 tracing::warn!(
                     "failed to continue active goal for idle thread {}: {err}",
@@ -334,7 +306,7 @@ where
                 // Capacity retries do not consume the user's token budget, so
                 // leave the goal active for a delayed idle continuation.
                 CodexErrorInfo::ServerOverloaded => {
-                    let _ = input.thread_store.insert(DeferredCapacityRetry);
+                    runtime.defer_capacity_retry();
                     return;
                 }
                 CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit,

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

@@ -1,7 +1,9 @@
 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;
@@ -20,6 +22,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 +51,7 @@ 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,
 }
@@ -97,6 +104,7 @@ 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),
             }),
@@ -115,6 +123,31 @@ 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
     }
@@ -357,6 +390,9 @@ 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(());