From f27bbbd49c0e05e763a96cc7fa677499de32b8d8 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 1 Jun 2026 11:32:13 +0200 Subject: [PATCH] Add goal extension GoalApi (#25096) ## Summary - add an extension-owned `GoalApi` for thread goal get/set/clear operations - register live goal runtimes with the API from the goal extension backend - cover the API and runtime-effect paths in goal extension tests ## Stack Follow-up app-server wiring PR: #25108 ## Validation - `just fmt` - `just fix -p codex-goal-extension` - `just test -p codex-goal-extension` --- codex-rs/ext/goal/src/api.rs | 286 ++++++++++++++++++ codex-rs/ext/goal/src/extension.rs | 14 + codex-rs/ext/goal/src/lib.rs | 7 + codex-rs/ext/goal/src/runtime.rs | 9 +- codex-rs/ext/goal/src/tool.rs | 8 +- .../ext/goal/tests/goal_extension_backend.rs | 149 +++++++-- 6 files changed, 433 insertions(+), 40 deletions(-) create mode 100644 codex-rs/ext/goal/src/api.rs diff --git a/codex-rs/ext/goal/src/api.rs b/codex-rs/ext/goal/src/api.rs new file mode 100644 index 000000000..0e23c295a --- /dev/null +++ b/codex-rs/ext/goal/src/api.rs @@ -0,0 +1,286 @@ +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::Weak; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::ThreadGoal; +use codex_protocol::protocol::ThreadGoalStatus; +use codex_protocol::protocol::validate_thread_goal_objective; + +use crate::runtime::GoalRuntimeHandle; +use crate::runtime::PreviousGoalSnapshot; +use crate::tool::fill_empty_thread_preview_if_possible; +use crate::tool::protocol_goal_from_state; +use crate::tool::state_status_from_protocol; +use crate::tool::validate_goal_budget; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum GoalServiceError { + InvalidRequest(String), + Internal(String), +} + +impl fmt::Display for GoalServiceError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidRequest(message) | Self::Internal(message) => f.write_str(message), + } + } +} + +impl std::error::Error for GoalServiceError {} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GoalObjectiveUpdate<'a> { + Keep, + Set(&'a str), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GoalTokenBudgetUpdate { + Keep, + Set(Option), +} + +#[derive(Clone, Copy, Debug)] +pub struct GoalSetRequest<'a> { + pub thread_id: ThreadId, + pub objective: GoalObjectiveUpdate<'a>, + pub status: Option, + pub token_budget: GoalTokenBudgetUpdate, +} + +#[derive(Clone, Debug)] +pub struct GoalSetOutcome { + pub goal: ThreadGoal, + state_goal: codex_state::ThreadGoal, + previous_goal: Option, +} + +impl GoalSetOutcome { + pub async fn apply_runtime_effects(&self, goal_service: &GoalService) { + if let Some(runtime) = goal_service.runtime_for_thread(self.goal.thread_id) + && let Err(err) = runtime + .apply_external_goal_set(self.state_goal.clone(), self.previous_goal.clone()) + .await + { + tracing::warn!("failed to apply external goal status runtime effects: {err}"); + } + } +} + +#[derive(Debug, Default)] +pub struct GoalService { + runtimes: Mutex>>, +} + +impl GoalService { + pub fn new() -> Self { + Self::default() + } + + pub async fn get_thread_goal( + &self, + state_db: &codex_state::StateRuntime, + thread_id: ThreadId, + ) -> Result, GoalServiceError> { + state_db + .thread_goals() + .get_thread_goal(thread_id) + .await + .map(|goal| goal.map(protocol_goal_from_state)) + .map_err(|err| GoalServiceError::Internal(format!("failed to read thread goal: {err}"))) + } + + pub async fn set_thread_goal( + &self, + state_db: &codex_state::StateRuntime, + request: GoalSetRequest<'_>, + ) -> Result { + let GoalSetRequest { + thread_id, + objective, + status, + token_budget, + } = request; + let status = status.map(state_status_from_protocol); + let objective = match objective { + GoalObjectiveUpdate::Keep => None, + GoalObjectiveUpdate::Set(objective) => Some(objective.trim()), + }; + let token_budget = match token_budget { + GoalTokenBudgetUpdate::Keep => None, + GoalTokenBudgetUpdate::Set(token_budget) => Some(token_budget), + }; + + if let Some(objective) = objective { + validate_thread_goal_objective(objective).map_err(GoalServiceError::InvalidRequest)?; + } + if objective.is_some() || token_budget.is_some() { + validate_goal_budget(token_budget.flatten()) + .map_err(GoalServiceError::InvalidRequest)?; + } + + if let Some(runtime) = self.runtime_for_thread(thread_id) + && let Err(err) = runtime.prepare_external_goal_mutation().await + { + tracing::warn!("failed to prepare external goal mutation: {err}"); + } + + 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| { + GoalServiceError::Internal(format!("failed to read thread goal: {err}")) + })?; + if let Some(existing_goal) = existing_goal.as_ref() { + let previous_goal = PreviousGoalSnapshot::from(existing_goal); + state_db + .thread_goals() + .update_thread_goal( + thread_id, + codex_state::GoalUpdate { + objective: Some(objective.to_string()), + status, + token_budget, + expected_goal_id: Some(existing_goal.goal_id.clone()), + }, + ) + .await + .map_err(|err| { + GoalServiceError::Internal(format!("failed to update thread goal: {err}")) + })? + .ok_or_else(|| { + GoalServiceError::InvalidRequest(format!( + "cannot update goal for thread {thread_id}: no goal exists" + )) + }) + .map(|goal| (goal, Some(previous_goal)))? + } else { + state_db + .thread_goals() + .replace_thread_goal( + thread_id, + objective, + status.unwrap_or(codex_state::ThreadGoalStatus::Active), + token_budget.flatten(), + ) + .await + .map_err(|err| { + GoalServiceError::Internal(format!("failed to replace thread goal: {err}")) + }) + .map(|goal| (goal, None))? + } + } else { + let existing_goal = state_db + .thread_goals() + .get_thread_goal(thread_id) + .await + .map_err(|err| { + GoalServiceError::Internal(format!("failed to read thread goal: {err}")) + })? + .ok_or_else(|| { + GoalServiceError::InvalidRequest(format!( + "cannot update goal for thread {thread_id}: no goal exists" + )) + })?; + let previous_goal = PreviousGoalSnapshot::from(&existing_goal); + let expected_goal_id = existing_goal.goal_id.clone(); + state_db + .thread_goals() + .update_thread_goal( + thread_id, + codex_state::GoalUpdate { + objective: None, + status, + token_budget, + expected_goal_id: Some(expected_goal_id), + }, + ) + .await + .map_err(|err| { + GoalServiceError::Internal(format!("failed to update thread goal: {err}")) + })? + .ok_or_else(|| { + GoalServiceError::InvalidRequest(format!( + "cannot update goal for thread {thread_id}: no goal exists" + )) + }) + .map(|goal| (goal, Some(previous_goal)))? + }; + + if objective.is_some() { + fill_empty_thread_preview_if_possible(state_db, thread_id, &goal).await; + } + Ok(GoalSetOutcome { + goal: protocol_goal_from_state(goal.clone()), + state_goal: goal, + previous_goal, + }) + } + + pub async fn clear_thread_goal( + &self, + state_db: &codex_state::StateRuntime, + thread_id: ThreadId, + ) -> Result { + if let Some(runtime) = self.runtime_for_thread(thread_id) + && let Err(err) = runtime.prepare_external_goal_mutation().await + { + tracing::warn!("failed to prepare external goal mutation: {err}"); + } + + let cleared = state_db + .thread_goals() + .delete_thread_goal(thread_id) + .await + .map_err(|err| { + GoalServiceError::Internal(format!("failed to clear thread goal: {err}")) + })?; + + if cleared + && let Some(runtime) = self.runtime_for_thread(thread_id) + && let Err(err) = runtime.apply_external_goal_clear().await + { + tracing::warn!("failed to apply external goal clear runtime effects: {err}"); + } + + Ok(cleared) + } + + pub(crate) fn register_runtime(&self, runtime: &Arc) { + self.runtimes() + .insert(runtime.thread_id().to_string(), Arc::downgrade(runtime)); + } + + pub(crate) fn unregister_runtime(&self, runtime: &Arc) { + let key = runtime.thread_id().to_string(); + let runtime = Arc::downgrade(runtime); + let mut runtimes = self.runtimes(); + if runtimes + .get(&key) + .is_some_and(|registered| registered.ptr_eq(&runtime)) + { + runtimes.remove(&key); + } + } + + fn runtime_for_thread(&self, thread_id: ThreadId) -> Option> { + let key = thread_id.to_string(); + let mut runtimes = self.runtimes(); + let runtime = runtimes.get(&key).and_then(Weak::upgrade); + if runtime.is_none() { + runtimes.remove(&key); + } + runtime + } + + fn runtimes(&self) -> std::sync::MutexGuard<'_, HashMap>> { + self.runtimes.lock().unwrap_or_else(PoisonError::into_inner) + } +} diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 6e1ebdafc..563c40b50 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -11,6 +11,7 @@ use codex_extension_api::ThreadIdleInput; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadResumeInput; use codex_extension_api::ThreadStartInput; +use codex_extension_api::ThreadStopInput; use codex_extension_api::TokenUsageContributor; use codex_extension_api::ToolCallOutcome; use codex_extension_api::ToolContributor; @@ -32,6 +33,7 @@ use codex_protocol::protocol::TokenUsageInfo; use crate::accounting::BudgetLimitedGoalDisposition; use crate::accounting::GoalAccountingState; +use crate::api::GoalService; use crate::events::GoalEventEmitter; use crate::metrics::GoalMetrics; use crate::runtime::GoalRuntimeConfig; @@ -57,6 +59,7 @@ pub struct GoalExtension { event_emitter: GoalEventEmitter, metrics: GoalMetrics, thread_manager: Weak, + goal_service: Arc, goals_enabled: Arc bool + Send + Sync>, } @@ -72,6 +75,7 @@ impl GoalExtension { event_sink: Arc, metrics_client: Option, thread_manager: Weak, + goal_service: Arc, goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, ) -> Self { Self { @@ -79,6 +83,7 @@ impl GoalExtension { event_emitter: GoalEventEmitter::new(event_sink), metrics: GoalMetrics::new(metrics_client), thread_manager, + goal_service, goals_enabled: Arc::new(goals_enabled), } } @@ -120,6 +125,7 @@ where ) }); runtime.set_enabled(enabled); + self.goal_service.register_runtime(&runtime); } async fn on_thread_resume(&self, input: ThreadResumeInput<'_>) { @@ -147,6 +153,12 @@ where ); } } + + async fn on_thread_stop(&self, input: ThreadStopInput<'_>) { + if let Some(runtime) = goal_runtime_handle(input.thread_store) { + self.goal_service.unregister_runtime(&runtime); + } + } } impl ConfigContributor for GoalExtension @@ -408,6 +420,7 @@ pub fn install_with_backend( state_dbs: Arc, metrics_client: Option, thread_manager: Weak, + goal_service: Arc, goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, ) where C: Send + Sync + 'static, @@ -417,6 +430,7 @@ pub fn install_with_backend( registry.event_sink(), metrics_client, thread_manager, + Arc::clone(&goal_service), goals_enabled, )); registry.thread_lifecycle_contributor(extension.clone()); diff --git a/codex-rs/ext/goal/src/lib.rs b/codex-rs/ext/goal/src/lib.rs index 545d9e699..8433e1f7d 100644 --- a/codex-rs/ext/goal/src/lib.rs +++ b/codex-rs/ext/goal/src/lib.rs @@ -5,6 +5,7 @@ //! accounting that can be represented with today's extension API. mod accounting; +mod api; mod events; mod extension; mod metrics; @@ -13,6 +14,12 @@ mod spec; mod steering; mod tool; +pub use api::GoalObjectiveUpdate; +pub use api::GoalService; +pub use api::GoalServiceError; +pub use api::GoalSetOutcome; +pub use api::GoalSetRequest; +pub use api::GoalTokenBudgetUpdate; pub use extension::GoalExtension; pub use extension::GoalExtensionConfig; pub use extension::install_with_backend; diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index fcb412495..8d79391c9 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -158,10 +158,8 @@ impl GoalRuntimeHandle { self.inner .metrics .record_terminal_if_status_changed(previous_status, &goal); - let should_steer_active_turn = previous_goal.as_ref().is_none_or(|previous_goal| { - previous_goal.goal_id != goal.goal_id - || previous_goal.status != codex_state::ThreadGoalStatus::Active - || previous_goal.objective != goal.objective + let objective_changed = previous_goal.as_ref().is_some_and(|previous_goal| { + !replaced_existing_goal && previous_goal.objective != goal.objective }); match goal.status { codex_state::ThreadGoalStatus::Active => { @@ -175,10 +173,11 @@ impl GoalRuntimeHandle { .accounting_state .mark_idle_goal_active(goal.goal_id.clone()); } - if should_steer_active_turn { + if objective_changed { let item = objective_updated_steering_item(&protocol_goal_from_state(goal)); self.inject_active_turn_steering(item).await; } + self.continue_if_idle().await?; } codex_state::ThreadGoalStatus::BudgetLimited => { if self.inner.accounting_state.current_turn_id().is_none() { diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index 90ea0eb44..45ef3c54d 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -363,7 +363,7 @@ where .map_err(|err| FunctionCallError::RespondToModel(err.to_string())) } -fn validate_goal_budget(value: Option) -> Result<(), String> { +pub(crate) fn validate_goal_budget(value: Option) -> Result<(), String> { if let Some(value) = value && value <= 0 { @@ -402,7 +402,7 @@ impl GoalToolResponse { } } -async fn fill_empty_thread_preview_if_possible( +pub(crate) async fn fill_empty_thread_preview_if_possible( state_db: &codex_state::StateRuntime, thread_id: ThreadId, goal: &codex_state::ThreadGoal, @@ -441,7 +441,9 @@ fn protocol_status_from_state(status: codex_state::ThreadGoalStatus) -> ThreadGo } } -fn state_status_from_protocol(status: ThreadGoalStatus) -> codex_state::ThreadGoalStatus { +pub(crate) fn state_status_from_protocol( + status: ThreadGoalStatus, +) -> codex_state::ThreadGoalStatus { match status { ThreadGoalStatus::Active => codex_state::ThreadGoalStatus::Active, ThreadGoalStatus::Paused => codex_state::ThreadGoalStatus::Paused, diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index f325c5f90..16b8cf261 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -11,6 +11,7 @@ use codex_extension_api::FunctionCallError; use codex_extension_api::NoopTurnItemEmitter; use codex_extension_api::ThreadResumeInput; use codex_extension_api::ThreadStartInput; +use codex_extension_api::ThreadStopInput; use codex_extension_api::ToolCall; use codex_extension_api::ToolCallOutcome; use codex_extension_api::ToolCallSource; @@ -20,8 +21,11 @@ use codex_extension_api::ToolPayload; use codex_extension_api::TurnErrorInput; use codex_extension_api::TurnStartInput; use codex_extension_api::TurnStopInput; +use codex_goal_extension::GoalObjectiveUpdate; use codex_goal_extension::GoalRuntimeHandle; -use codex_goal_extension::PreviousGoalSnapshot; +use codex_goal_extension::GoalService; +use codex_goal_extension::GoalSetRequest; +use codex_goal_extension::GoalTokenBudgetUpdate; use codex_goal_extension::install_with_backend; use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationMode; @@ -778,7 +782,8 @@ async fn external_goal_mutation_start_accounts_active_goal_progress() -> anyhow: } #[tokio::test] -async fn external_goal_set_active_resets_baseline_without_live_thread() -> anyhow::Result<()> { +async fn goal_service_external_set_active_resets_baseline_without_live_thread() -> anyhow::Result<()> +{ let runtime = test_runtime().await?; let thread_id = test_thread_id()?; seed_thread_metadata(runtime.as_ref(), thread_id).await?; @@ -815,38 +820,19 @@ async fn external_goal_set_active_resets_baseline_without_live_thread() -> anyho ), ) .await; - harness - .runtime_handle() - .prepare_external_goal_mutation() - .await - .map_err(anyhow::Error::msg)?; - - let previous_goal = runtime - .thread_goals() - .get_thread_goal(thread_id) - .await? - .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; - let updated_goal = runtime - .thread_goals() - .update_thread_goal( - thread_id, - codex_state::GoalUpdate { - objective: Some("new objective".to_string()), - status: Some(codex_state::ThreadGoalStatus::Active), - token_budget: None, - expected_goal_id: Some(previous_goal.goal_id.clone()), + let outcome = harness + .goal_service + .set_thread_goal( + runtime.as_ref(), + GoalSetRequest { + thread_id, + objective: GoalObjectiveUpdate::Set("new objective"), + status: Some(ThreadGoalStatus::Active), + token_budget: GoalTokenBudgetUpdate::Keep, }, ) - .await? - .ok_or_else(|| anyhow::anyhow!("goal update should succeed"))?; - harness - .runtime_handle() - .apply_external_goal_set( - updated_goal, - Some(PreviousGoalSnapshot::from(&previous_goal)), - ) - .await - .map_err(anyhow::Error::msg)?; + .await?; + outcome.apply_runtime_effects(&harness.goal_service).await; harness .record_token_usage( @@ -871,6 +857,46 @@ async fn external_goal_set_active_resets_baseline_without_live_thread() -> anyho Ok(()) } +#[tokio::test] +async fn thread_stop_unregisters_goal_runtime_from_service() -> 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(); + let create_tool = tool_by_name(&tools, "create_goal"); + create_tool + .handle(tool_call( + "create_goal", + "call-create-goal", + json!({ "objective": "ship goal extension backend" }), + )) + .await?; + harness.sink.clear(); + + harness + .record_token_usage( + "turn-1", + &token_usage( + /*input_tokens*/ 10, /*cached_input_tokens*/ 0, /*output_tokens*/ 0, + /*reasoning_output_tokens*/ 0, /*total_tokens*/ 10, + ), + ) + .await; + harness.stop_thread().await; + + assert!( + harness + .goal_service + .clear_thread_goal(runtime.as_ref(), thread_id) + .await? + ); + assert_eq!(Vec::::new(), harness.sink.goal_events()); + Ok(()) +} + #[tokio::test] async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Result<()> { let runtime = test_runtime().await?; @@ -917,6 +943,48 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul Ok(()) } +#[tokio::test] +async fn goal_service_sets_gets_and_clears_thread_goal() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let api = GoalService::new(); + + let set = api + .set_thread_goal( + runtime.as_ref(), + GoalSetRequest { + thread_id, + objective: GoalObjectiveUpdate::Set(" ship goal API ownership "), + status: None, + token_budget: GoalTokenBudgetUpdate::Set(Some(123)), + }, + ) + .await?; + let get = api + .get_thread_goal(runtime.as_ref(), thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + let metadata = runtime + .get_thread(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("seeded thread metadata should exist"))?; + + assert_eq!(set.goal, get); + assert_eq!("ship goal API ownership", get.objective); + assert_eq!(ThreadGoalStatus::Active, get.status); + assert_eq!(Some(123), get.token_budget); + assert_eq!(Some("ship goal API ownership"), metadata.preview.as_deref()); + + assert!(api.clear_thread_goal(runtime.as_ref(), thread_id).await?); + assert_eq!( + None, + api.get_thread_goal(runtime.as_ref(), thread_id).await? + ); + assert!(!api.clear_thread_goal(runtime.as_ref(), thread_id).await?); + Ok(()) +} + async fn installed_tools( runtime: Arc, thread_id: ThreadId, @@ -937,11 +1005,13 @@ async fn installed_tools_with_start( persistent_thread_state_available: bool, ) -> Vec>> { let mut builder = ExtensionRegistryBuilder::<()>::new(); + let goal_service = Arc::new(GoalService::new()); install_with_backend( &mut builder, runtime, /*metrics_client*/ None, Weak::new(), + goal_service, |_| true, ); let registry = builder.build(); @@ -974,6 +1044,7 @@ struct GoalExtensionHarness { registry: codex_extension_api::ExtensionRegistry<()>, session_store: ExtensionData, thread_store: ExtensionData, + goal_service: Arc, sink: Arc, } @@ -984,11 +1055,13 @@ impl GoalExtensionHarness { ) -> anyhow::Result { let sink = Arc::new(RecordingEventSink::default()); let mut builder = ExtensionRegistryBuilder::<()>::with_event_sink(sink.clone()); + let goal_service = Arc::new(GoalService::new()); install_with_backend( &mut builder, runtime, /*metrics_client*/ None, Weak::new(), + Arc::clone(&goal_service), |_| true, ); let registry = builder.build(); @@ -1010,6 +1083,7 @@ impl GoalExtensionHarness { registry, session_store, thread_store, + goal_service, sink, }) } @@ -1088,6 +1162,17 @@ impl GoalExtensionHarness { } } + async fn stop_thread(&self) { + for contributor in self.registry.thread_lifecycle_contributors() { + contributor + .on_thread_stop(ThreadStopInput { + 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);