diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 563c40b50..78ae0b5b4 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -36,6 +36,7 @@ use crate::accounting::GoalAccountingState; use crate::api::GoalService; use crate::events::GoalEventEmitter; use crate::metrics::GoalMetrics; +use crate::runtime::ActiveGoalStopReason; use crate::runtime::GoalRuntimeConfig; use crate::runtime::GoalRuntimeHandle; use crate::spec::UPDATE_GOAL_TOOL_NAME; @@ -278,18 +279,26 @@ where } async fn on_turn_error(&self, input: TurnErrorInput<'_>) { - if input.error != CodexErrorInfo::UsageLimitExceeded { - return; - } let Some(runtime) = goal_runtime_handle(input.thread_store) else { return; }; + let reason = match input.error { + 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 + // continuation from looping and consuming tokens, as can happen + // with compaction errors. + _ => ActiveGoalStopReason::TurnError, + }; if let Err(err) = runtime - .usage_limit_active_goal_for_turn(input.turn_id) + .stop_active_goal_for_turn(input.turn_id, reason) .await { - tracing::warn!("failed to usage-limit active goal after usage-limit error: {err}"); + tracing::warn!( + error = ?input.error, + "failed to stop active goal after turn error: {err}" + ); } } } diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 9818282ca..2641dfb94 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -28,6 +28,11 @@ pub(crate) struct GoalRuntimeConfig { pub(crate) tools_available_for_thread: bool, } +pub(crate) enum ActiveGoalStopReason { + TurnError, + UsageLimit, +} + struct GoalRuntimeInner { thread_id: ThreadId, state_dbs: Arc, @@ -216,10 +221,23 @@ impl GoalRuntimeHandle { } pub async fn usage_limit_active_goal_for_turn(&self, turn_id: &str) -> Result<(), String> { + self.stop_active_goal_for_turn(turn_id, ActiveGoalStopReason::UsageLimit) + .await + } + + /// Accounts the ending turn and stops its active goal after a terminal error. + pub(crate) async fn stop_active_goal_for_turn( + &self, + turn_id: &str, + reason: ActiveGoalStopReason, + ) -> Result<(), String> { if !self.is_enabled() { return Ok(()); } + // Hold this through accounting and the status update so external goal + // mutations and idle continuation cannot interleave between them. + let _goal_state_permit = self.goal_state_permit().await?; if !self .inner .accounting_state @@ -228,23 +246,54 @@ impl GoalRuntimeHandle { return Ok(()); } - let progress_event_id = format!("{turn_id}:usage-limit-progress"); + let (event_name, status) = match reason { + ActiveGoalStopReason::TurnError => { + ("turn-error", codex_state::ThreadGoalStatus::Blocked) + } + ActiveGoalStopReason::UsageLimit => { + ("usage-limit", codex_state::ThreadGoalStatus::UsageLimited) + } + }; self.account_active_goal_progress( turn_id, - progress_event_id.as_str(), + &format!("{turn_id}:{event_name}-progress"), codex_state::GoalAccountingMode::ActiveOnly, BudgetLimitedGoalDisposition::ClearActive, ) .await?; - let previous_status = self - .current_goal_status_for_metrics(/*expected_goal_id*/ None) - .await?; + let Some(active_goal) = self + .inner + .state_dbs + .thread_goals() + .get_thread_goal(self.thread_id()) + .await + .map_err(|err| err.to_string())? + else { + self.inner.accounting_state.clear_active_goal(); + return Ok(()); + }; + let can_stop = active_goal.status == codex_state::ThreadGoalStatus::Active + || (active_goal.status == codex_state::ThreadGoalStatus::BudgetLimited + && status == codex_state::ThreadGoalStatus::UsageLimited); + if !can_stop { + self.inner.accounting_state.clear_active_goal(); + return Ok(()); + } + let previous_status = Some(active_goal.status); let Some(goal) = self .inner .state_dbs .thread_goals() - .usage_limit_active_thread_goal(self.thread_id()) + .update_thread_goal( + self.thread_id(), + codex_state::GoalUpdate { + objective: None, + status: Some(status), + token_budget: None, + expected_goal_id: Some(active_goal.goal_id), + }, + ) .await .map_err(|err| err.to_string())? else { @@ -256,7 +305,7 @@ impl GoalRuntimeHandle { self.inner.accounting_state.clear_active_goal(); let goal = protocol_goal_from_state(goal); self.inner.event_emitter.thread_goal_updated( - format!("{turn_id}:usage-limit"), + format!("{turn_id}:{event_name}"), Some(turn_id.to_string()), goal, ); diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 3c437a088..f89a08ebb 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -494,18 +494,9 @@ async fn turn_error_usage_limit_accounts_progress_and_clears_accounting() -> any ), ) .await; - let turn_store = ExtensionData::new("turn-1"); - for contributor in harness.registry.turn_lifecycle_contributors() { - contributor - .on_turn_error(TurnErrorInput { - turn_id: "turn-1", - error: CodexErrorInfo::UsageLimitExceeded, - session_store: &harness.session_store, - thread_store: &harness.thread_store, - turn_store: &turn_store, - }) - .await; - } + harness + .notify_turn_error("turn-1", CodexErrorInfo::UsageLimitExceeded) + .await; let goal = runtime .thread_goals() @@ -557,6 +548,36 @@ async fn turn_error_usage_limit_accounts_progress_and_clears_accounting() -> any Ok(()) } +#[tokio::test] +async fn turn_error_blocks_goal() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call( + "create_goal", + "call-create-goal", + json!({ "objective": "ship goal extension backend" }), + )) + .await?; + + harness + .notify_turn_error("turn-1", CodexErrorInfo::Other) + .await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!(codex_state::ThreadGoalStatus::Blocked, goal.status); + Ok(()) +} + #[tokio::test] async fn usage_limit_budget_limited_goal_accounts_remaining_progress() -> anyhow::Result<()> { let runtime = test_runtime().await?; @@ -1255,6 +1276,21 @@ impl GoalExtensionHarness { } } + async fn notify_turn_error(&self, turn_id: &str, error: CodexErrorInfo) { + let turn_store = ExtensionData::new(turn_id); + for contributor in self.registry.turn_lifecycle_contributors() { + contributor + .on_turn_error(TurnErrorInput { + turn_id, + error: error.clone(), + session_store: &self.session_store, + thread_store: &self.thread_store, + turn_store: &turn_store, + }) + .await; + } + } + fn runtime_handle(&self) -> Arc { self.thread_store .get::()