diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 5421d1729..e09f041e3 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -17,11 +17,13 @@ use codex_extension_api::ToolFinishInput; use codex_extension_api::ToolLifecycleContributor; use codex_extension_api::ToolLifecycleFuture; use codex_extension_api::TurnAbortInput; +use codex_extension_api::TurnErrorInput; use codex_extension_api::TurnLifecycleContributor; use codex_extension_api::TurnStartInput; use codex_extension_api::TurnStopInput; use codex_otel::MetricsClient; use codex_protocol::ThreadId; +use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadGoalStatus; @@ -248,6 +250,22 @@ where } runtime.accounting_state().finish_turn(turn_id); } + + 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; + }; + + if let Err(err) = runtime + .usage_limit_active_goal_for_turn(input.turn_id) + .await + { + tracing::warn!("failed to usage-limit active goal after usage-limit error: {err}"); + } + } } #[async_trait] diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 7524d40fe..92f4313fc 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -17,6 +17,7 @@ use codex_extension_api::ToolCallSource; use codex_extension_api::ToolExecutor; use codex_extension_api::ToolFinishInput; use codex_extension_api::ToolPayload; +use codex_extension_api::TurnErrorInput; use codex_extension_api::TurnStartInput; use codex_extension_api::TurnStopInput; use codex_goal_extension::GoalRuntimeHandle; @@ -26,6 +27,7 @@ use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; +use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::SessionSource; @@ -398,7 +400,7 @@ async fn budget_limited_goal_keeps_accounting_after_later_tool_finish() -> anyho } #[tokio::test] -async fn usage_limit_active_goal_accounts_progress_and_clears_accounting() -> anyhow::Result<()> { +async fn turn_error_usage_limit_accounts_progress_and_clears_accounting() -> anyhow::Result<()> { let runtime = test_runtime().await?; let thread_id = test_thread_id()?; seed_thread_metadata(runtime.as_ref(), thread_id).await?; @@ -425,11 +427,18 @@ async fn usage_limit_active_goal_accounts_progress_and_clears_accounting() -> an ), ) .await; - harness - .runtime_handle() - .usage_limit_active_goal_for_turn("turn-1") - .await - .map_err(anyhow::Error::msg)?; + 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; + } let goal = runtime .thread_goals()