From 823381e867ae21babe8a41c1b070fc21b5ddd81d Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 26 May 2026 20:01:13 +0200 Subject: [PATCH] fix: restore goal accounting after thread resume (#24626) ## Why Goal idle accounting is supposed to survive a thread resume. Previously, the resume hook restored the active goal state inline from the extension lifecycle contributor, which left the runtime handle without a reusable restoration path and made the behavior hard to cover directly. When a thread with an active goal was resumed, goal accounting could lose track of the active idle goal instead of continuing to accrue elapsed time. ## What changed - Moved thread-resume restoration into `GoalRuntimeHandle::restore_after_resume()` so the runtime owns rehydrating active goal accounting from persisted thread goal state. - Kept disabled goal runtimes as a no-op and preserved the existing warning path when persisted goal state cannot be loaded. - Added a backend regression test that seeds an active goal, resumes the thread, waits briefly, and verifies elapsed idle time is reflected on the next external goal mutation. ## Testing - Not run locally; this metadata update only rewrote the PR title/body. --- codex-rs/ext/goal/src/extension.rs | 31 ++-------- codex-rs/ext/goal/src/runtime.rs | 24 ++++++++ .../ext/goal/tests/goal_extension_backend.rs | 59 +++++++++++++++++++ 3 files changed, 88 insertions(+), 26 deletions(-) diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index aa0f658dae..68068acbb4 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -112,33 +112,12 @@ where let Some(runtime) = goal_runtime_handle(input.thread_store) else { return; }; - if !runtime.is_enabled() { - return; - } - let goal = match self - .state_dbs - .thread_goals() - .get_thread_goal(runtime.thread_id()) - .await - { - Ok(goal) => goal, - Err(err) => { - tracing::warn!( - "failed to restore goal runtime after thread resume for {}: {err}", - runtime.thread_id() - ); - return; - } - }; - match goal { - Some(goal) if goal.status == codex_state::ThreadGoalStatus::Active => { - runtime - .accounting_state() - .mark_idle_goal_active(goal.goal_id); - self.metrics.record_resumed(); - } - Some(_) | None => runtime.accounting_state().clear_active_goal(), + if let Err(err) = runtime.restore_after_resume().await { + tracing::warn!( + "failed to restore goal runtime after thread resume for {}: {err}", + runtime.thread_id() + ); } } } diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 689d0578cf..0139cc2c52 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -192,6 +192,30 @@ impl GoalRuntimeHandle { Ok(()) } + pub async fn restore_after_resume(&self) -> Result<(), String> { + if !self.is_enabled() { + return Ok(()); + } + + let goal = self + .inner + .state_dbs + .thread_goals() + .get_thread_goal(self.thread_id()) + .await + .map_err(|err| err.to_string())?; + match goal { + Some(goal) if goal.status == codex_state::ThreadGoalStatus::Active => { + self.inner + .accounting_state + .mark_idle_goal_active(goal.goal_id); + self.inner.metrics.record_resumed(); + } + Some(_) | None => self.inner.accounting_state.clear_active_goal(), + } + Ok(()) + } + pub(crate) async fn inject_active_turn_steering(&self, item: ResponseInputItem) { let Some(thread_manager) = self.inner.thread_manager.upgrade() else { tracing::debug!("skipping goal steering because thread manager is unavailable"); diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 0a3c768985..0a36439b25 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -2,11 +2,13 @@ use std::sync::Arc; use std::sync::Mutex; use std::sync::PoisonError; use std::sync::Weak; +use std::time::Duration; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::FunctionCallError; +use codex_extension_api::ThreadResumeInput; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolCall; use codex_extension_api::ToolCallOutcome; @@ -589,6 +591,52 @@ async fn external_goal_set_active_resets_baseline_without_live_thread() -> anyho Ok(()) } +#[tokio::test] +async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + runtime + .thread_goals() + .replace_thread_goal( + thread_id, + "ship goal extension backend", + codex_state::ThreadGoalStatus::Active, + /*token_budget*/ None, + ) + .await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + + harness.resume_thread().await; + tokio::time::sleep(Duration::from_millis(1_100)).await; + harness + .runtime_handle() + .prepare_external_goal_mutation() + .await + .map_err(anyhow::Error::msg)?; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!(ThreadGoalStatus::Active, protocol_status(goal.status)); + assert!( + goal.time_used_seconds >= 1, + "resumed idle accounting should add elapsed wall-clock time" + ); + assert_eq!( + vec![CapturedGoalEvent { + event_id: format!("{thread_id}:external-goal-mutation"), + turn_id: None, + status: ThreadGoalStatus::Active, + tokens_used: 0, + }], + harness.sink.goal_events() + ); + Ok(()) +} + async fn installed_tools( runtime: Arc, thread_id: ThreadId, @@ -719,6 +767,17 @@ impl GoalExtensionHarness { } } + async fn resume_thread(&self) { + for contributor in self.registry.thread_lifecycle_contributors() { + contributor + .on_thread_resume(ThreadResumeInput { + session_store: &self.session_store, + thread_store: &self.thread_store, + }) + .await; + } + } + async fn notify_tool_finish(&self, turn_id: &str, call_id: &str, tool_name: &str) { let turn_store = ExtensionData::new(turn_id); let tool_name = codex_extension_api::ToolName::plain(tool_name);