diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index ac8f9d049e..0ceda1007a 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2533,7 +2533,7 @@ async fn turn_error_lifecycle_exposes_error_and_stores() { fn on_turn_error<'a>( &'a self, input: codex_extension_api::TurnErrorInput<'a>, - ) -> codex_extension_api::ExtensionFuture<'a, ()> { + ) -> codex_extension_api::ExtensionFuture<'a, Option> { Box::pin(async move { self.records .lock() @@ -2553,6 +2553,7 @@ async fn turn_error_lifecycle_exposes_error_and_stores() { .get::() .is_some(), }); + None }) } } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index e231f6af86..3f35cdbd56 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -153,15 +153,18 @@ pub(crate) async fn run_turn( // new user message are recorded. Estimate pending incoming items (context // diffs/full reinjection + user input) and trigger compaction preemptively // when they would push the thread over the compaction threshold. - if let Err(err) = run_pre_sampling_compact(&sess, &turn_context, &mut client_session).await { + while let Err(err) = run_pre_sampling_compact(&sess, &turn_context, &mut client_session).await { if matches!(err, CodexErr::TurnAborted) { return Err(err); } - let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) + let retry_delay = sess + .emit_turn_error_lifecycle(turn_context.as_ref(), err.to_codex_protocol_error()) .await; error!("Failed to run pre-sampling compact"); - return Ok(None); + let Some(delay) = retry_delay else { + return Ok(None); + }; + tokio::time::sleep(delay).await; } // run_turn owns the step used to seed context and make the first sampling request. @@ -186,7 +189,9 @@ pub(crate) async fn run_turn( if run_pending_session_start_hooks(&sess, &turn_context).await { return Ok(None); } - let mut can_drain_pending_input = input.is_empty(); + let mut can_drain_pending_input = !input + .iter() + .any(|item| matches!(item, TurnInput::UserInput { .. })); if run_hooks_and_record_inputs(&sess, &turn_context, &input).await { return Ok(None); } @@ -218,7 +223,7 @@ pub(crate) async fn run_turn( // one instance across retries within this turn. // Pending input is drained into history before building the next model request. // However, we defer that drain until after sampling in two cases: - // 1. At the start of a turn, so the fresh turn input in `input` gets sampled first. + // 1. At the start of a user turn, so the submitted user input gets sampled first. // 2. After auto-compact, when model/tool continuation needs to resume before any steer. let mut next_step_context = Some(first_step_context); @@ -361,7 +366,7 @@ pub(crate) async fn run_turn( return Err(err); } let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) + sess.emit_turn_error_lifecycle(turn_context.as_ref(), error) .await; return Ok(None); } @@ -445,11 +450,18 @@ pub(crate) async fn run_turn( Err(e) => { info!("Turn error: {e:#}"); let error = e.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) + let retry_delay = sess + .emit_turn_error_lifecycle(turn_context.as_ref(), error) .await; sess.track_turn_codex_error(turn_context.as_ref(), &e); let event = EventMsg::Error(e.to_error_event(/*message_prefix*/ None)); sess.send_event(&turn_context, event).await; + if let Some(delay) = retry_delay { + // Keep this turn alive so its recorded input is not appended again. + tokio::time::sleep(delay).await; + can_drain_pending_input = true; + continue; + } // let the user continue the conversation break; } diff --git a/codex-rs/core/src/tasks/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs index 782cc8c7de..785051253d 100644 --- a/codex-rs/core/src/tasks/lifecycle.rs +++ b/codex-rs/core/src/tasks/lifecycle.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use codex_extension_api::ExtensionData; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::TokenUsage; @@ -76,9 +78,10 @@ impl Session { &self, turn_context: &TurnContext, error: CodexErrorInfo, - ) { + ) -> Option { + let mut retry_delay: Option = None; for contributor in self.services.extensions.turn_lifecycle_contributors() { - contributor + retry_delay = retry_delay.or(contributor .on_turn_error(codex_extension_api::TurnErrorInput { turn_id: turn_context.sub_id.as_str(), error: error.clone(), @@ -86,7 +89,8 @@ impl Session { thread_store: &self.services.thread_extension_data, turn_store: turn_context.extension_data.as_ref(), }) - .await; + .await); } + retry_delay } } diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index 214039f3d7..fdcd4c0565 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -4937,53 +4937,6 @@ async fn snapshot_request_shape_pre_turn_compaction_context_window_exceeded() { ); } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn automatic_idle_turn_does_not_retry_failed_pre_turn_compaction() { - skip_if_no_network!(); - - let server = start_mock_server().await; - mount_sse_sequence( - &server, - vec![ - sse(vec![ - ev_assistant_message("m1", FIRST_REPLY), - ev_completed_with_tokens("r1", /*total_tokens*/ 500), - ]), - sse_failed( - "compact-overloaded", - "server_is_overloaded", - "Selected model is at capacity. Please try a different model.", - ), - ], - ) - .await; - - let mut model_provider = non_openai_model_provider(&server); - model_provider.stream_max_retries = Some(0); - let test = test_codex() - .with_config(move |config| { - config.model_provider = model_provider; - set_test_compact_prompt(config); - config.model_auto_compact_token_limit = Some(200); - }) - .build(&server) - .await - .expect("build codex"); - - test.submit_turn("USER_ONE").await.expect("submit turn"); - - test.codex - .try_start_turn_if_idle(vec![responses::user_message_item("GOAL_CONTINUATION")]) - .await - .expect("start automatic idle turn"); - tokio::time::timeout( - std::time::Duration::from_secs(5), - wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))), - ) - .await - .expect("failed pre-turn compaction should end the automatic turn"); -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn snapshot_request_shape_manual_compact_without_previous_user_messages() { skip_if_no_network!(); diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index 9c9852a276..b6c4d60d6a 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -3,6 +3,8 @@ use std::sync::Arc; use codex_core::CodexThread; use codex_core::config::CurrentTimeReminderConfig; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::TurnLifecycleContributor; use codex_features::Feature; use codex_protocol::AgentPath; use codex_protocol::items::SleepItem; @@ -27,6 +29,7 @@ use core_test_support::responses::ev_output_text_delta; use core_test_support::responses::ev_reasoning_item; use core_test_support::responses::ev_reasoning_item_added; use core_test_support::responses::ev_response_created; +use core_test_support::skip_if_no_network; use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::StreamingSseServer; use core_test_support::streaming_sse::start_streaming_sse_server; @@ -40,6 +43,17 @@ use serde_json::from_slice; use serde_json::json; use tokio::sync::oneshot; +struct RetryTurnErrors; + +impl TurnLifecycleContributor for RetryTurnErrors { + fn on_turn_error<'a>( + &'a self, + _input: codex_extension_api::TurnErrorInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, Option> { + Box::pin(async { Some(std::time::Duration::ZERO) }) + } +} + fn ev_message_item_done(id: &str, text: &str) -> Value { serde_json::json!({ "type": "response.output_item.done", @@ -287,6 +301,57 @@ async fn wait_for_sleep_item_completed(codex: &CodexThread, call_id: &str, durat ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn automatic_turn_retries_in_place_with_queued_mail_in_first_request() { + skip_if_no_network!(); + + const CONTINUATION: &str = "GOAL_CONTINUATION"; + const MAIL: &str = "CHILD_UPDATE"; + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse_failed( + "overloaded", + "server_is_overloaded", + "Selected model is at capacity. Please try a different model.", + ), + responses::sse_completed("response-2"), + ], + ) + .await; + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.turn_lifecycle_contributor(Arc::new(RetryTurnErrors)); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| config.model_provider.stream_max_retries = Some(0)) + .build(&server) + .await + .expect("build Codex test session"); + + submit_queue_only_agent_mail(&test.codex, MAIL).await; + test.codex + .try_start_turn_if_idle(vec![responses::user_message_item(CONTINUATION)]) + .await + .expect("start automatic turn"); + wait_for_turn_complete(&test.codex).await; + + let requests = response_mock.requests(); + let [first, second] = requests.as_slice() else { + panic!("expected two requests, got {}", requests.len()); + }; + assert_eq!(1, first.inputs_of_type("agent_message").len()); + assert!(first.body_contains_text(MAIL)); + assert_eq!( + 1, + second + .message_input_texts("user") + .iter() + .filter(|text| text.as_str() == CONTINUATION) + .count() + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn steer_interrupts_wait_agent_and_is_sent_in_follow_up_request() { const WAIT_CALL_ID: &str = "wait-call"; diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 96f1ee0b48..3fb30e56ff 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -1,6 +1,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::time::Duration; use codex_context_fragments::ContextualUserFragment; use codex_protocol::items::TurnItem; @@ -184,11 +185,16 @@ pub trait TurnLifecycleContributor: Send + Sync { }) } - /// Called when the host observes an error for a running turn. - fn on_turn_error<'a>(&'a self, input: TurnErrorInput<'a>) -> ExtensionFuture<'a, ()> { + /// Called when the host observes an error for a running turn. Returning a + /// delay asks the host to retry the failed operation without ending the turn. + fn on_turn_error<'a>( + &'a self, + input: TurnErrorInput<'a>, + ) -> ExtensionFuture<'a, Option> { Box::pin(async move { let _self = self; let _input = input; + None }) } } diff --git a/codex-rs/ext/goal/Cargo.toml b/codex-rs/ext/goal/Cargo.toml index 52556a6c87..4d75ba2f26 100644 --- a/codex-rs/ext/goal/Cargo.toml +++ b/codex-rs/ext/goal/Cargo.toml @@ -24,7 +24,7 @@ codex-tools = { workspace = true } codex-utils-template = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } -tokio = { workspace = true, features = ["rt", "sync", "time"] } +tokio = { workspace = true, features = ["sync"] } tracing = { workspace = true } [dev-dependencies] diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index e93011208d..e0ebd0c43a 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -1,5 +1,6 @@ use std::sync::Arc; use std::sync::Weak; +use std::time::Duration; use codex_analytics::AnalyticsEventsClient; use codex_core::ThreadManager; @@ -45,6 +46,10 @@ 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. Keep the retry cadence deliberately low. +const DEFERRED_GOAL_RETRY_DELAY: Duration = Duration::from_secs(5 * 60); + #[derive(Clone, Debug)] pub struct GoalExtensionConfig { pub enabled: bool, @@ -296,23 +301,21 @@ where }) } - fn on_turn_error<'a>(&'a self, input: TurnErrorInput<'a>) -> ExtensionFuture<'a, ()> { + fn on_turn_error<'a>( + &'a self, + input: TurnErrorInput<'a>, + ) -> ExtensionFuture<'a, Option> { Box::pin(async move { - let Some(runtime) = goal_runtime_handle(input.thread_store) else { - return; - }; + let runtime = goal_runtime_handle(input.thread_store)?; + if input.error == CodexErrorInfo::ServerOverloaded + && runtime + .accounting_state() + .turn_is_current_active_goal(input.turn_id) + { + return Some(DEFERRED_GOAL_RETRY_DELAY); + } 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 - if runtime - .accounting_state() - .turn_is_current_active_goal(input.turn_id) => - { - runtime.defer_retry(); - return; - } CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit, // The turn has ended because the error was non-retryable or its // retries were exhausted. Block the goal to prevent automatic @@ -329,6 +332,7 @@ where "failed to stop active goal after turn error: {err}" ); } + None }) } } diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 1df0568b58..2f2202cbb8 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -2,7 +2,6 @@ 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; @@ -21,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 DEFERRED_GOAL_RETRY_DELAY: Duration = Duration::from_secs(5 * 60); - #[derive(Clone)] pub struct GoalRuntimeHandle { inner: Arc, @@ -50,7 +45,6 @@ struct GoalRuntimeInner { thread_manager: Weak, accounting_state: Arc, enabled: AtomicBool, - retry_pending: AtomicBool, tools_available_for_thread: bool, goal_state_lock: Semaphore, } @@ -103,7 +97,6 @@ impl GoalRuntimeHandle { thread_manager, accounting_state, enabled: AtomicBool::new(config.enabled), - retry_pending: AtomicBool::new(false), tools_available_for_thread: config.tools_available_for_thread, goal_state_lock: Semaphore::new(/*permits*/ 1), }), @@ -122,33 +115,6 @@ impl GoalRuntimeHandle { self.is_enabled() && self.inner.tools_available_for_thread } - // Schedule at most one deferred retry for this live runtime. Idle - // continuation stays suppressed until the timer fires, and unloading the - // runtime cancels the retry by invalidating the weak reference. - pub(crate) fn defer_retry(&self) { - if self.inner.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(DEFERRED_GOAL_RETRY_DELAY).await; - let Some(inner) = runtime.upgrade() else { - return; - }; - inner.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 } @@ -391,9 +357,6 @@ impl GoalRuntimeHandle { } pub(crate) async fn continue_if_idle(&self) -> Result<(), String> { - if self.inner.retry_pending.load(Ordering::Relaxed) { - return Ok(()); - } if !self.tools_visible() { self.inner.accounting_state.clear_active_goal(); return Ok(()); diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index cc429d3cdb..f28bb88817 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -576,7 +576,7 @@ async fn turn_error_usage_limit_accounts_progress_and_clears_accounting() -> any #[tokio::test] async fn turn_error_blocks_goal() -> anyhow::Result<()> { assert_eq!( - codex_state::ThreadGoalStatus::Blocked, + (codex_state::ThreadGoalStatus::Blocked, false), goal_status_after_turn_error(CodexErrorInfo::Other).await? ); Ok(()) @@ -585,7 +585,7 @@ async fn turn_error_blocks_goal() -> anyhow::Result<()> { #[tokio::test] async fn server_overloaded_error_keeps_goal_active() -> anyhow::Result<()> { assert_eq!( - codex_state::ThreadGoalStatus::Active, + (codex_state::ThreadGoalStatus::Active, true), goal_status_after_turn_error(CodexErrorInfo::ServerOverloaded).await? ); Ok(()) @@ -593,7 +593,7 @@ async fn server_overloaded_error_keeps_goal_active() -> anyhow::Result<()> { async fn goal_status_after_turn_error( error: CodexErrorInfo, -) -> anyhow::Result { +) -> anyhow::Result<(codex_state::ThreadGoalStatus, bool)> { let runtime = test_runtime().await?; let thread_id = test_thread_id()?; seed_thread_metadata(runtime.as_ref(), thread_id).await?; @@ -609,14 +609,14 @@ async fn goal_status_after_turn_error( )) .await?; - harness.notify_turn_error("turn-1", error).await; + let retry_delay = harness.notify_turn_error("turn-1", error).await; let goal = runtime .thread_goals() .get_thread_goal(thread_id) .await? .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; - Ok(goal.status) + Ok((goal.status, retry_delay.is_some())) } #[tokio::test] @@ -1321,10 +1321,15 @@ impl GoalExtensionHarness { } } - async fn notify_turn_error(&self, turn_id: &str, error: CodexErrorInfo) { + async fn notify_turn_error( + &self, + turn_id: &str, + error: CodexErrorInfo, + ) -> Option { let turn_store = ExtensionData::new(turn_id); + let mut retry_delay = None; for contributor in self.registry.turn_lifecycle_contributors() { - contributor + retry_delay = contributor .on_turn_error(TurnErrorInput { turn_id, error: error.clone(), @@ -1334,6 +1339,7 @@ impl GoalExtensionHarness { }) .await; } + retry_delay } fn runtime_handle(&self) -> Arc {