diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index 519ec9cb79..e843826472 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/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(()) } diff --git a/codex-rs/ext/goal/src/api.rs b/codex-rs/ext/goal/src/api.rs index 807f481d0d..871718b9f1 100644 --- a/codex-rs/ext/goal/src/api.rs +++ b/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>>, + // 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 { @@ -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, 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(); @@ -320,4 +372,10 @@ 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 a4da74132c..5df73145f2 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/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, diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index c9f21d2db6..2f2202cbb8 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/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, @@ -51,7 +45,6 @@ struct GoalRuntimeInner { thread_manager: Weak, accounting_state: Arc, 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(());