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

@@ -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<i64>),
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct GoalSetRequest<'a> {
+    pub thread_id: ThreadId,
+    pub objective: GoalObjectiveUpdate<'a>,
+    pub status: Option<ThreadGoalStatus>,
+    pub token_budget: GoalTokenBudgetUpdate,
+}
+
+#[derive(Clone, Debug)]
+pub struct GoalSetOutcome {
+    pub goal: ThreadGoal,
+    state_goal: codex_state::ThreadGoal,
+    previous_goal: Option<PreviousGoalSnapshot>,
+}
+
+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<HashMap<String, Weak<GoalRuntimeHandle>>>,
+}
+
+impl GoalService {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub async fn get_thread_goal(
+        &self,
+        state_db: &codex_state::StateRuntime,
+        thread_id: ThreadId,
+    ) -> Result<Option<ThreadGoal>, 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<GoalSetOutcome, GoalServiceError> {
+        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<bool, GoalServiceError> {
+        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<GoalRuntimeHandle>) {
+        self.runtimes()
+            .insert(runtime.thread_id().to_string(), Arc::downgrade(runtime));
+    }
+
+    pub(crate) fn unregister_runtime(&self, runtime: &Arc<GoalRuntimeHandle>) {
+        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<Arc<GoalRuntimeHandle>> {
+        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<String, Weak<GoalRuntimeHandle>>> {
+        self.runtimes.lock().unwrap_or_else(PoisonError::into_inner)
+    }
+}

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<C> {
     event_emitter: GoalEventEmitter,
     metrics: GoalMetrics,
     thread_manager: Weak<ThreadManager>,
+    goal_service: Arc<GoalService>,
     goals_enabled: Arc<dyn Fn(&C) -> bool + Send + Sync>,
 }
 
@@ -72,6 +75,7 @@ impl<C> GoalExtension<C> {
         event_sink: Arc<dyn ExtensionEventSink>,
         metrics_client: Option<MetricsClient>,
         thread_manager: Weak<ThreadManager>,
+        goal_service: Arc<GoalService>,
         goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
     ) -> Self {
         Self {
@@ -79,6 +83,7 @@ impl<C> GoalExtension<C> {
             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<C> ConfigContributor<C> for GoalExtension<C>
@@ -408,6 +420,7 @@ pub fn install_with_backend<C>(
     state_dbs: Arc<codex_state::StateRuntime>,
     metrics_client: Option<MetricsClient>,
     thread_manager: Weak<ThreadManager>,
+    goal_service: Arc<GoalService>,
     goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
 ) where
     C: Send + Sync + 'static,
@@ -417,6 +430,7 @@ pub fn install_with_backend<C>(
         registry.event_sink(),
         metrics_client,
         thread_manager,
+        Arc::clone(&goal_service),
         goals_enabled,
     ));
     registry.thread_lifecycle_contributor(extension.clone());

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;

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| {
+        let objective_changed = previous_goal.as_ref().is_some_and(|previous_goal| {
-            previous_goal.goal_id != goal.goal_id
+            !replaced_existing_goal && previous_goal.objective != goal.objective
-                || previous_goal.status != codex_state::ThreadGoalStatus::Active
-                || 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() {

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<i64>) -> Result<(), String> {
+pub(crate) fn validate_goal_budget(value: Option<i64>) -> 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,

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
+    let outcome = harness
-        .runtime_handle()
+        .goal_service
-        .prepare_external_goal_mutation()
+        .set_thread_goal(
-        .await
+            runtime.as_ref(),
-        .map_err(anyhow::Error::msg)?;
+            GoalSetRequest {
-
+                thread_id,
-    let previous_goal = runtime
+                objective: GoalObjectiveUpdate::Set("new objective"),
-        .thread_goals()
+                status: Some(ThreadGoalStatus::Active),
-        .get_thread_goal(thread_id)
+                token_budget: GoalTokenBudgetUpdate::Keep,
-        .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()),
             },
         )
-        .await?
+        .await?;
-        .ok_or_else(|| anyhow::anyhow!("goal update should succeed"))?;
+    outcome.apply_runtime_effects(&harness.goal_service).await;
-    harness
-        .runtime_handle()
-        .apply_external_goal_set(
-            updated_goal,
-            Some(PreviousGoalSnapshot::from(&previous_goal)),
-        )
-        .await
-        .map_err(anyhow::Error::msg)?;
 
     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::<CapturedGoalEvent>::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<codex_state::StateRuntime>,
     thread_id: ThreadId,
@@ -937,11 +1005,13 @@ async fn installed_tools_with_start(
     persistent_thread_state_available: bool,
 ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
     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<GoalService>,
     sink: Arc<RecordingEventSink>,
 }
 
@@ -984,11 +1055,13 @@ impl GoalExtensionHarness {
     ) -> anyhow::Result<Self> {
         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);