From bc1cc8d3b02dc18c97fa6917010b97584f1d04de Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Sun, 5 Jul 2026 10:08:17 -0700 Subject: [PATCH] codex: scope goal capacity backoff to live runtime (#31176) --- codex-rs/ext/goal/src/api.rs | 58 ------------------------------ codex-rs/ext/goal/src/extension.rs | 8 +---- codex-rs/ext/goal/src/runtime.rs | 38 ++++++++++++++++++++ 3 files changed, 39 insertions(+), 65 deletions(-) diff --git a/codex-rs/ext/goal/src/api.rs b/codex-rs/ext/goal/src/api.rs index 871718b9f1..807f481d0d 100644 --- a/codex-rs/ext/goal/src/api.rs +++ b/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>>, - // Keep this above the per-thread runtime lifecycle so unloading and - // resuming a thread cannot bypass its pending backoff. - capacity_retry_deadlines: Mutex>, } 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, 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> { let key = thread_id.to_string(); let mut runtimes = self.runtimes(); @@ -372,10 +320,4 @@ impl GoalService { fn runtimes(&self) -> std::sync::MutexGuard<'_, HashMap>> { self.runtimes.lock().unwrap_or_else(PoisonError::into_inner) } - - fn capacity_retry_deadlines(&self) -> std::sync::MutexGuard<'_, HashMap> { - self.capacity_retry_deadlines - .lock() - .unwrap_or_else(PoisonError::into_inner) - } } diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 5df73145f2..2ba3c71fd7 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/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, diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 2f2202cbb8..0383547ee7 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/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, @@ -45,6 +50,7 @@ struct GoalRuntimeInner { thread_manager: Weak, accounting_state: Arc, 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(());