diff --git a/codex-rs/core/src/goals.rs b/codex-rs/core/src/goals.rs index 6e0663655c..a4206730f9 100644 --- a/codex-rs/core/src/goals.rs +++ b/codex-rs/core/src/goals.rs @@ -30,7 +30,9 @@ use codex_prompts::continuation_prompt; use codex_prompts::objective_updated_prompt; use codex_protocol::ThreadId; use codex_protocol::config_types::ModeKind; +use codex_protocol::error::CodexErr; use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ThreadGoal; use codex_protocol::protocol::ThreadGoalStatus; @@ -129,13 +131,14 @@ pub(crate) enum GoalRuntimeEvent<'a> { turn_context: &'a TurnContext, turn_completed: bool, }, + TurnErrored { + turn_context: &'a TurnContext, + error: &'a CodexErr, + }, MaybeContinueIfIdle, TaskAborted { turn_context: Option<&'a TurnContext>, }, - UsageLimitReached { - turn_context: &'a TurnContext, - }, ExternalMutationStarting, ExternalSet { external_set: ExternalGoalSet, @@ -144,6 +147,12 @@ pub(crate) enum GoalRuntimeEvent<'a> { ThreadResumed, } +#[derive(Clone, Copy)] +enum TerminalGoalErrorStatus { + Paused, + UsageLimited, +} + pub(crate) struct GoalRuntimeState { pub(crate) state_db: Mutex>, pub(crate) budget_limit_reported_goal_id: Mutex>, @@ -359,6 +368,14 @@ impl Session { .await; Ok(()) }), + GoalRuntimeEvent::TurnErrored { + turn_context, + error, + } => Box::pin(async move { + self.handle_thread_goal_turn_error(turn_context, error) + .await?; + Ok(()) + }), GoalRuntimeEvent::MaybeContinueIfIdle => Box::pin(async move { self.maybe_continue_goal_if_idle_runtime().await; Ok(()) @@ -367,11 +384,6 @@ impl Session { self.handle_thread_goal_task_abort(turn_context).await; Ok(()) }), - GoalRuntimeEvent::UsageLimitReached { turn_context } => Box::pin(async move { - self.usage_limit_active_thread_goal_for_turn(turn_context) - .await?; - Ok(()) - }), GoalRuntimeEvent::ExternalMutationStarting => Box::pin(async move { if let Err(err) = self.account_thread_goal_before_external_mutation().await { tracing::warn!( @@ -1138,14 +1150,41 @@ impl Session { } } - async fn usage_limit_active_thread_goal_for_turn( + async fn handle_thread_goal_turn_error( &self, turn_context: &TurnContext, + error: &CodexErr, + ) -> anyhow::Result<()> { + if error.to_codex_protocol_error() == CodexErrorInfo::UsageLimitExceeded { + return self + .finish_active_thread_goal_after_terminal_error( + turn_context, + TerminalGoalErrorStatus::UsageLimited, + ) + .await; + } + + if matches!(error, CodexErr::TurnAborted | CodexErr::Interrupted) { + return Ok(()); + } + + // This is a terminal turn error. Leave request-local retry decisions to + // the turn; pause the goal so idle continuation does not replay it. + self.finish_active_thread_goal_after_terminal_error( + turn_context, + TerminalGoalErrorStatus::Paused, + ) + .await + } + + async fn finish_active_thread_goal_after_terminal_error( + &self, + turn_context: &TurnContext, + status: TerminalGoalErrorStatus, ) -> anyhow::Result<()> { if should_ignore_goal_for_mode(turn_context.collaboration_mode.mode) { return Ok(()); } - if !self.enabled(Feature::Goals) { return Ok(()); } @@ -1159,6 +1198,7 @@ impl Session { let Some(state_db) = self.state_db_for_thread_goals().await? else { return Ok(()); }; + self.account_thread_goal_progress( turn_context, BudgetLimitSteering::Suppressed, @@ -1168,11 +1208,21 @@ impl Session { let previous_status = self .current_goal_status_for_metrics(&state_db, /*expected_goal_id*/ None) .await?; - let Some(goal) = state_db - .thread_goals() - .usage_limit_active_thread_goal(self.thread_id) - .await? - else { + let goal = match status { + TerminalGoalErrorStatus::UsageLimited => { + state_db + .thread_goals() + .usage_limit_active_thread_goal(self.thread_id) + .await? + } + TerminalGoalErrorStatus::Paused => { + state_db + .thread_goals() + .pause_active_thread_goal(self.thread_id) + .await? + } + }; + let Some(goal) = goal else { return Ok(()); }; self.emit_goal_terminal_metrics_if_status_changed(previous_status, &goal); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index a2023ea378..f6aa1eaccd 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -9628,8 +9628,10 @@ async fn usage_limit_runtime_stops_active_goal_and_prevents_idle_continuation() .await; set_total_token_usage(&sess, post_goal_token_usage()).await; - sess.goal_runtime_apply(GoalRuntimeEvent::UsageLimitReached { + let error = CodexErr::QuotaExceeded; + sess.goal_runtime_apply(GoalRuntimeEvent::TurnErrored { turn_context: tc.as_ref(), + error: &error, }) .await?; @@ -9650,6 +9652,58 @@ async fn usage_limit_runtime_stops_active_goal_and_prevents_idle_continuation() Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn turn_error_pauses_active_goal_and_prevents_idle_continuation() -> anyhow::Result<()> { + let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; + sess.set_thread_goal( + tc.as_ref(), + SetGoalRequest { + objective: Some("Keep improving the benchmark".to_string()), + status: None, + token_budget: Some(Some(500)), + }, + ) + .await?; + sess.goal_runtime_apply(GoalRuntimeEvent::TurnStarted { + turn_context: tc.as_ref(), + token_usage: TokenUsage::default(), + }) + .await?; + sess.spawn_task( + Arc::clone(&tc), + Vec::new(), + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: false, + }, + ) + .await; + set_total_token_usage(&sess, post_goal_token_usage()).await; + + let error = CodexErr::InvalidRequest("bad compaction request".to_string()); + sess.goal_runtime_apply(GoalRuntimeEvent::TurnErrored { + turn_context: tc.as_ref(), + error: &error, + }) + .await?; + + let state_db = goal_test_state_db(sess.as_ref()).await?; + let goal = state_db + .thread_goals() + .get_thread_goal(sess.thread_id) + .await? + .expect("goal should remain persisted after turn error"); + assert_eq!(codex_state::ThreadGoalStatus::Paused, goal.status); + assert_eq!(70, goal.tokens_used); + + sess.abort_all_tasks(TurnAbortReason::Replaced).await; + sess.goal_runtime_apply(GoalRuntimeEvent::MaybeContinueIfIdle) + .await?; + assert!(sess.active_turn.lock().await.is_none()); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn external_goal_mutation_accounts_active_turn_before_status_change() -> anyhow::Result<()> { let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 996e59187e..fb9439e07f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -148,18 +148,7 @@ pub(crate) async fn run_turn( // 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 { - let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) - .await; - if error == CodexErrorInfo::UsageLimitExceeded - && let Err(err) = sess - .goal_runtime_apply(GoalRuntimeEvent::UsageLimitReached { - turn_context: turn_context.as_ref(), - }) - .await - { - warn!("failed to usage-limit active goal after usage-limit error: {err}"); - } + emit_turn_error_and_apply_goal_runtime(&sess, &turn_context, &err).await; error!("Failed to run pre-sampling compact"); return None; } @@ -301,20 +290,7 @@ pub(crate) async fn run_turn( ) .await { - let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) - .await; - if error == CodexErrorInfo::UsageLimitExceeded - && let Err(err) = sess - .goal_runtime_apply(GoalRuntimeEvent::UsageLimitReached { - turn_context: turn_context.as_ref(), - }) - .await - { - warn!( - "failed to usage-limit active goal after usage-limit error: {err}" - ); - } + emit_turn_error_and_apply_goal_runtime(&sess, &turn_context, &err).await; return None; } can_drain_pending_input = !model_needs_follow_up; @@ -387,6 +363,7 @@ pub(crate) async fn run_turn( let error = CodexErrorInfo::BadRequest; sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) .await; + apply_goal_runtime_turn_error(&sess, &turn_context, &codex_error).await; let event = EventMsg::Error(ErrorEvent { message: "Invalid image in your last message. Please remove it and try again." .to_string(), @@ -397,18 +374,7 @@ 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()) - .await; - if error == CodexErrorInfo::UsageLimitExceeded - && let Err(err) = sess - .goal_runtime_apply(GoalRuntimeEvent::UsageLimitReached { - turn_context: turn_context.as_ref(), - }) - .await - { - warn!("failed to usage-limit active goal after usage-limit error: {err}"); - } + emit_turn_error_and_apply_goal_runtime(&sess, &turn_context, &e).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; @@ -421,6 +387,32 @@ pub(crate) async fn run_turn( last_agent_message } +async fn emit_turn_error_and_apply_goal_runtime( + sess: &Arc, + turn_context: &Arc, + error: &CodexErr, +) { + sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.to_codex_protocol_error()) + .await; + apply_goal_runtime_turn_error(sess, turn_context, error).await; +} + +async fn apply_goal_runtime_turn_error( + sess: &Arc, + turn_context: &Arc, + error: &CodexErr, +) { + if let Err(err) = sess + .goal_runtime_apply(GoalRuntimeEvent::TurnErrored { + turn_context: turn_context.as_ref(), + error, + }) + .await + { + warn!("failed to apply goal runtime turn-error event: {err}"); + } +} + async fn run_hooks_and_record_inputs( sess: &Arc, turn_context: &Arc,