From 7782d6733b553da2f56ae0e82dc904db8cd8bf2c Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Wed, 27 May 2026 00:34:13 -0700 Subject: [PATCH] Simplify goal extension refactor --- codex-rs/app-server/src/request_processors.rs | 1 - .../thread_goal_processor.rs | 97 +++++++++---------- codex-rs/core/src/tasks/mod.rs | 2 +- codex-rs/core/src/tools/router.rs | 1 - .../ext/extension-api/src/contributors.rs | 10 +- codex-rs/ext/extension-api/src/lib.rs | 1 - codex-rs/ext/goal/src/extension.rs | 3 +- codex-rs/ext/goal/src/lib.rs | 1 - codex-rs/ext/goal/src/runtime.rs | 19 +--- .../ext/goal/tests/goal_extension_backend.rs | 8 +- 10 files changed, 53 insertions(+), 90 deletions(-) diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 58e2a95847..719e956741 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -334,7 +334,6 @@ use codex_feedback::FeedbackUploadOptions; use codex_git_utils::git_diff_to_remote; use codex_git_utils::resolve_root_git_project_for_trust; use codex_goal_extension::GoalRuntimeHandle; -use codex_goal_extension::PreviousGoalSnapshot; use codex_login::AuthManager; use codex_login::CLIENT_ID; use codex_login::CodexAuth; diff --git a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs index 0d66f26937..121019525b 100644 --- a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs @@ -145,15 +145,14 @@ impl ThreadGoalRequestProcessor { prepare_external_goal_mutation(running_thread.as_deref()).await; let should_set_thread_preview = objective.is_some(); - let (goal, previous_goal) = (if let Some(objective) = objective { - let existing_goal = state_db - .thread_goals() - .get_thread_goal(thread_id) - .await - .map_err(|err| invalid_request(err.to_string()))?; - if let Some(goal) = existing_goal.as_ref() { - let previous_goal = PreviousGoalSnapshot::from(goal); - state_db + let existing_goal = state_db + .thread_goals() + .get_thread_goal(thread_id) + .await + .map_err(|err| invalid_request(err.to_string()))?; + let (goal, previous_goal) = match (objective, existing_goal) { + (Some(objective), Some(existing_goal)) => { + let goal = state_db .thread_goals() .update_thread_goal( thread_id, @@ -161,20 +160,20 @@ impl ThreadGoalRequestProcessor { objective: Some(objective.to_string()), status, token_budget: params.token_budget, - expected_goal_id: Some(goal.goal_id.clone()), + expected_goal_id: Some(existing_goal.goal_id.clone()), }, ) .await - .and_then(|goal| { - goal.ok_or_else(|| { - anyhow::anyhow!( - "cannot update goal for thread {thread_id}: no goal exists" - ) - }) - }) - .map(|goal| (goal, Some(previous_goal))) - } else { - state_db + .map_err(|err| invalid_request(err.to_string()))? + .ok_or_else(|| { + invalid_request(format!( + "cannot update goal for thread {thread_id}: no goal exists" + )) + })?; + (goal, Some(existing_goal)) + } + (Some(objective), None) => { + let goal = state_db .thread_goals() .replace_thread_goal( thread_id, @@ -183,40 +182,36 @@ impl ThreadGoalRequestProcessor { params.token_budget.flatten(), ) .await - .map(|goal| (goal, None)) + .map_err(|err| invalid_request(err.to_string()))?; + (goal, None) } - } else { - let existing_goal = state_db - .thread_goals() - .get_thread_goal(thread_id) - .await - .map_err(|err| invalid_request(err.to_string()))?; - let Some(existing_goal) = existing_goal else { + (None, Some(existing_goal)) => { + let goal = state_db + .thread_goals() + .update_thread_goal( + thread_id, + codex_state::GoalUpdate { + objective: None, + status, + token_budget: params.token_budget, + expected_goal_id: None, + }, + ) + .await + .map_err(|err| invalid_request(err.to_string()))? + .ok_or_else(|| { + invalid_request(format!( + "cannot update goal for thread {thread_id}: no goal exists" + )) + })?; + (goal, Some(existing_goal)) + } + (None, None) => { return Err(invalid_request(format!( "cannot update goal for thread {thread_id}: no goal exists" ))); - }; - let previous_goal = PreviousGoalSnapshot::from(&existing_goal); - state_db - .thread_goals() - .update_thread_goal( - thread_id, - codex_state::GoalUpdate { - objective: None, - status, - token_budget: params.token_budget, - expected_goal_id: None, - }, - ) - .await - .and_then(|goal| { - goal.ok_or_else(|| { - anyhow::anyhow!("cannot update goal for thread {thread_id}: no goal exists") - }) - }) - .map(|goal| (goal, Some(previous_goal))) - }) - .map_err(|err| invalid_request(err.to_string()))?; + } + }; if should_set_thread_preview && let Err(err) = state_db .set_thread_preview_if_empty(thread_id, goal.objective.as_str()) @@ -453,7 +448,7 @@ async fn prepare_external_goal_mutation(thread: Option<&CodexThread>) { async fn apply_external_goal_set( thread: Option<&CodexThread>, goal: codex_state::ThreadGoal, - previous_goal: Option, + previous_goal: Option, ) { let Some(thread) = thread else { return; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index ea1dd2354c..b29ded1e19 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -512,7 +512,7 @@ impl Session { }) .await { - requested_items = Some(request.items); + requested_items = Some(request); break; } } diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 04f5ebbc37..dd55c622d5 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -230,7 +230,6 @@ pub(crate) fn extension_tool_executors( contributor.tools_for_turn(codex_extension_api::ToolContributionInput { session_store: &session.services.session_extension_data, thread_store: &session.services.thread_extension_data, - turn_store: Some(turn_context.extension_data.as_ref()), session_source: &turn_context.session_source, persistent_thread: !turn_context.config.ephemeral && session.state_db().is_some(), }) diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index ce215f3ed7..ee15b82bc3 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -141,19 +141,12 @@ pub struct ToolContributionInput<'a> { pub session_store: &'a ExtensionData, /// Store scoped to this thread runtime. pub thread_store: &'a ExtensionData, - /// Store scoped to this turn runtime, if the host is collecting for a turn. - pub turn_store: Option<&'a ExtensionData>, /// Source for the current turn's session. pub session_source: &'a SessionSource, /// Whether the current thread has persistent state available. pub persistent_thread: bool, } -/// Extension request to start a new idle turn. -pub struct IdleTurnRequest { - pub items: Vec, -} - /// Context supplied when the host is idle and extensions may request work. pub struct IdleTurnInput<'a> { /// Effective collaboration mode for the next default turn. @@ -164,7 +157,8 @@ pub struct IdleTurnInput<'a> { pub thread_store: &'a ExtensionData, } -pub type IdleTurnFuture<'a> = Pin> + Send + 'a>>; +pub type IdleTurnFuture<'a> = + Pin>> + Send + 'a>>; /// Extension contribution that can request host-owned work while a thread is idle. pub trait IdleTurnContributor: Send + Sync { diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index b6bb84f47f..f92aca46b6 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -28,7 +28,6 @@ pub use contributors::ContextContributor; pub use contributors::IdleTurnContributor; pub use contributors::IdleTurnFuture; pub use contributors::IdleTurnInput; -pub use contributors::IdleTurnRequest; pub use contributors::PromptFragment; pub use contributors::PromptSlot; pub use contributors::ThreadLifecycleContributor; diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index a678343c4b..4eb09fcee1 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -10,7 +10,6 @@ use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::IdleTurnContributor; use codex_extension_api::IdleTurnFuture; use codex_extension_api::IdleTurnInput; -use codex_extension_api::IdleTurnRequest; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadResumeInput; use codex_extension_api::ThreadStartInput; @@ -429,7 +428,7 @@ where .idle_continuation_items(input.collaboration_mode.mode) .await { - Ok(Some(items)) => Some(IdleTurnRequest { items }), + Ok(Some(items)) => Some(items), Ok(None) => None, Err(err) => { tracing::warn!("failed to request idle goal continuation: {err}"); diff --git a/codex-rs/ext/goal/src/lib.rs b/codex-rs/ext/goal/src/lib.rs index 7056744baa..08a7a65da7 100644 --- a/codex-rs/ext/goal/src/lib.rs +++ b/codex-rs/ext/goal/src/lib.rs @@ -17,7 +17,6 @@ pub use extension::GoalExtension; pub use extension::GoalExtensionConfig; pub use extension::install_with_backend; pub use runtime::GoalRuntimeHandle; -pub use runtime::PreviousGoalSnapshot; pub use spec::CREATE_GOAL_TOOL_NAME; pub use spec::GET_GOAL_TOOL_NAME; pub use spec::UPDATE_GOAL_TOOL_NAME; diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 61ed289c83..12a347d771 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -37,23 +37,6 @@ pub(crate) struct AccountedGoalProgress { pub(crate) goal_id: String, } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct PreviousGoalSnapshot { - pub goal_id: String, - pub status: codex_state::ThreadGoalStatus, - pub objective: String, -} - -impl From<&codex_state::ThreadGoal> for PreviousGoalSnapshot { - fn from(goal: &codex_state::ThreadGoal) -> Self { - Self { - goal_id: goal.goal_id.clone(), - status: goal.status, - objective: goal.objective.clone(), - } - } -} - impl std::fmt::Debug for GoalRuntimeHandle { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("GoalRuntimeHandle").finish_non_exhaustive() @@ -127,7 +110,7 @@ impl GoalRuntimeHandle { pub async fn apply_external_goal_set( &self, goal: codex_state::ThreadGoal, - previous_goal: Option, + previous_goal: Option, ) -> Result<(), String> { if !self.is_enabled() { 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 f528ff23b7..1dfc8e3c41 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -21,7 +21,6 @@ use codex_extension_api::TurnErrorInput; use codex_extension_api::TurnStartInput; use codex_extension_api::TurnStopInput; use codex_goal_extension::GoalRuntimeHandle; -use codex_goal_extension::PreviousGoalSnapshot; use codex_goal_extension::install_with_backend; use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationMode; @@ -661,10 +660,7 @@ async fn external_goal_set_active_resets_baseline_without_live_thread() -> anyho .ok_or_else(|| anyhow::anyhow!("goal update should succeed"))?; harness .runtime_handle() - .apply_external_goal_set( - updated_goal, - Some(PreviousGoalSnapshot::from(&previous_goal)), - ) + .apply_external_goal_set(updated_goal, Some(previous_goal)) .await .map_err(anyhow::Error::msg)?; @@ -924,7 +920,7 @@ impl GoalExtensionHarness { }) .await { - return Some(request.items); + return Some(request); } } None