diff --git a/codex-rs/ext/goal/src/api.rs b/codex-rs/ext/goal/src/api.rs new file mode 100644 index 0000000000..3420aa0940 --- /dev/null +++ b/codex-rs/ext/goal/src/api.rs @@ -0,0 +1,269 @@ +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 GoalApiError { + InvalidRequest(String), + Internal(String), +} + +impl fmt::Display for GoalApiError { + 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 GoalApiError {} + +#[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_api: &GoalApi) { + if let Some(runtime) = goal_api.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 GoalApi { + runtimes: Mutex>>, +} + +impl GoalApi { + pub fn new() -> Self { + Self::default() + } + + pub async fn get_thread_goal( + &self, + state_db: &codex_state::StateRuntime, + thread_id: ThreadId, + ) -> Result, GoalApiError> { + state_db + .thread_goals() + .get_thread_goal(thread_id) + .await + .map(|goal| goal.map(protocol_goal_from_state)) + .map_err(|err| GoalApiError::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(GoalApiError::InvalidRequest)?; + } + if objective.is_some() || token_budget.is_some() { + validate_goal_budget(token_budget.flatten()).map_err(GoalApiError::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| { + GoalApiError::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| { + GoalApiError::Internal(format!("failed to update thread goal: {err}")) + })? + .ok_or_else(|| { + GoalApiError::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| { + GoalApiError::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| { + GoalApiError::Internal(format!("failed to read thread goal: {err}")) + })? + .ok_or_else(|| { + GoalApiError::InvalidRequest(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, + expected_goal_id: None, + }, + ) + .await + .map_err(|err| { + GoalApiError::Internal(format!("failed to update thread goal: {err}")) + })? + .ok_or_else(|| { + GoalApiError::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| GoalApiError::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)); + } + + fn runtime_for_thread(&self, thread_id: ThreadId) -> Option> { + let key = thread_id.to_string(); + let runtime = self.runtimes().get(&key).and_then(Weak::upgrade); + if runtime.is_none() { + self.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 c823889ce1..411dee18fc 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -30,6 +30,7 @@ use codex_protocol::protocol::TokenUsageInfo; use crate::accounting::BudgetLimitedGoalDisposition; use crate::accounting::GoalAccountingState; +use crate::api::GoalApi; use crate::events::GoalEventEmitter; use crate::metrics::GoalMetrics; use crate::runtime::GoalRuntimeConfig; @@ -55,6 +56,7 @@ pub struct GoalExtension { event_emitter: GoalEventEmitter, metrics: GoalMetrics, thread_manager: Weak, + goal_api: Arc, goals_enabled: Arc bool + Send + Sync>, } @@ -70,6 +72,7 @@ impl GoalExtension { event_sink: Arc, metrics_client: Option, thread_manager: Weak, + goal_api: Arc, goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, ) -> Self { Self { @@ -77,6 +80,7 @@ impl GoalExtension { event_emitter: GoalEventEmitter::new(event_sink), metrics: GoalMetrics::new(metrics_client), thread_manager, + goal_api, goals_enabled: Arc::new(goals_enabled), } } @@ -118,6 +122,7 @@ where ) }); runtime.set_enabled(enabled); + self.goal_api.register_runtime(&runtime); } async fn on_thread_resume(&self, input: ThreadResumeInput<'_>) { @@ -390,6 +395,7 @@ pub fn install_with_backend( state_dbs: Arc, metrics_client: Option, thread_manager: Weak, + goal_api: Arc, goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, ) where C: Send + Sync + 'static, @@ -399,6 +405,7 @@ pub fn install_with_backend( registry.event_sink(), metrics_client, thread_manager, + Arc::clone(&goal_api), 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 545d9e699a..a9ccd96a38 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::GoalApi; +pub use api::GoalApiError; +pub use api::GoalObjectiveUpdate; +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/tool.rs b/codex-rs/ext/goal/src/tool.rs index 90ea0eb44b..45ef3c54d0 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 7524d40fe1..4d5c1e3de7 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -19,8 +19,11 @@ use codex_extension_api::ToolFinishInput; use codex_extension_api::ToolPayload; use codex_extension_api::TurnStartInput; use codex_extension_api::TurnStopInput; +use codex_goal_extension::GoalApi; +use codex_goal_extension::GoalObjectiveUpdate; use codex_goal_extension::GoalRuntimeHandle; -use codex_goal_extension::PreviousGoalSnapshot; +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; @@ -769,7 +772,7 @@ 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_api_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?; @@ -806,38 +809,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_api + .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_api).await; harness .record_token_usage( @@ -908,6 +892,48 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul Ok(()) } +#[tokio::test] +async fn goal_api_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 = GoalApi::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, @@ -928,11 +954,13 @@ async fn installed_tools_with_start( persistent_thread_state_available: bool, ) -> Vec>> { let mut builder = ExtensionRegistryBuilder::<()>::new(); + let goal_api = Arc::new(GoalApi::new()); install_with_backend( &mut builder, runtime, /*metrics_client*/ None, Weak::new(), + goal_api, |_| true, ); let registry = builder.build(); @@ -965,6 +993,7 @@ struct GoalExtensionHarness { registry: codex_extension_api::ExtensionRegistry<()>, session_store: ExtensionData, thread_store: ExtensionData, + goal_api: Arc, sink: Arc, } @@ -975,11 +1004,13 @@ impl GoalExtensionHarness { ) -> anyhow::Result { let sink = Arc::new(RecordingEventSink::default()); let mut builder = ExtensionRegistryBuilder::<()>::with_event_sink(sink.clone()); + let goal_api = Arc::new(GoalApi::new()); install_with_backend( &mut builder, runtime, /*metrics_client*/ None, Weak::new(), + Arc::clone(&goal_api), |_| true, ); let registry = builder.build(); @@ -1001,6 +1032,7 @@ impl GoalExtensionHarness { registry, session_store, thread_store, + goal_api, sink, }) }