feat: account active goal progress in the goal extension (#23696)

## Why

The goal extension can create and surface goals, but the live
turn-accounting path still stopped short of persisting active-goal
progress. That leaves token and wall-clock usage, plus
`ThreadGoalUpdated` events, out of sync with the extension boundary once
work actually advances or a goal transitions out of active state.

## What changed

- Teach `GoalAccountingState` to track the current turn, active goal,
token deltas, and wall-clock progress snapshots against the persisted
goal id.
- Flush active-goal accounting from tool-finish, turn-stop, and
turn-abort lifecycle hooks, and emit `ThreadGoalUpdated` events when
persisted progress changes.
- Route `create_goal` and `update_goal` through the same accounting
state so new goals start from the right baseline, final progress is
flushed before status changes, and `update_goal` can mark a goal
`blocked` as well as `complete`.
- Keep budget-limited goals accruing through the end of the turn while
clearing local active-goal state once a turn or explicit update is
finished.
- Expand backend and lifecycle coverage around store ids, baseline
reset, tool-finish accounting, budget-limited carry-through, and
blocked-goal updates.

## Testing

- Added focused backend coverage in
`codex-rs/ext/goal/tests/goal_extension_backend.rs` for baseline reset,
tool-finish accounting, budget-limited turns, and blocked-goal updates.
- Extended `codex-rs/core/src/session/tests.rs` to assert that lifecycle
inputs expose the expected session, thread, and turn store ids.

codex-rs/ext/goal/src/tool.rs

@@ -15,6 +15,8 @@ use codex_protocol::protocol::validate_thread_goal_objective;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::accounting::BudgetLimitedGoalDisposition;
+use crate::accounting::GoalAccountingState;
 use crate::events::GoalEventEmitter;
 use crate::spec::CREATE_GOAL_TOOL_NAME;
 use crate::spec::GET_GOAL_TOOL_NAME;
@@ -28,6 +30,7 @@ pub(crate) struct GoalToolExecutor {
     kind: GoalToolKind,
     thread_id: ThreadId,
     state_db: Arc<codex_state::StateRuntime>,
+    accounting_state: Arc<GoalAccountingState>,
     event_emitter: GoalEventEmitter,
 }
 
@@ -69,12 +72,14 @@ impl GoalToolExecutor {
     pub(crate) fn get(
         thread_id: ThreadId,
         state_db: Arc<codex_state::StateRuntime>,
+        accounting_state: Arc<GoalAccountingState>,
         event_emitter: GoalEventEmitter,
     ) -> Self {
         Self {
             kind: GoalToolKind::Get,
             thread_id,
             state_db,
+            accounting_state,
             event_emitter,
         }
     }
@@ -82,12 +87,14 @@ impl GoalToolExecutor {
     pub(crate) fn create(
         thread_id: ThreadId,
         state_db: Arc<codex_state::StateRuntime>,
+        accounting_state: Arc<GoalAccountingState>,
         event_emitter: GoalEventEmitter,
     ) -> Self {
         Self {
             kind: GoalToolKind::Create,
             thread_id,
             state_db,
+            accounting_state,
             event_emitter,
         }
     }
@@ -95,12 +102,14 @@ impl GoalToolExecutor {
     pub(crate) fn update(
         thread_id: ThreadId,
         state_db: Arc<codex_state::StateRuntime>,
+        accounting_state: Arc<GoalAccountingState>,
         event_emitter: GoalEventEmitter,
     ) -> Self {
         Self {
             kind: GoalToolKind::Update,
             thread_id,
             state_db,
+            accounting_state,
             event_emitter,
         }
     }
@@ -179,8 +188,11 @@ impl GoalToolExecutor {
                 )
             })?;
         fill_empty_thread_preview_if_possible(self.state_db.as_ref(), self.thread_id, &goal).await;
+        let turn_id = self
+            .accounting_state
+            .mark_current_turn_goal_active(goal.goal_id.clone());
         let goal = protocol_goal_from_state(goal);
-        self.emit_goal_updated_from_tool_call(&invocation, goal.clone());
+        self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone());
         goal_response(Some(goal), CompletionBudgetReport::Omit)
     }
 
@@ -189,15 +201,29 @@ impl GoalToolExecutor {
         invocation: ToolCall,
     ) -> Result<Box<dyn ToolOutput>, FunctionCallError> {
         let args: UpdateGoalArgs = parse_arguments(invocation.function_arguments()?)?;
-        if args.status != ThreadGoalStatus::Complete {
+        if !matches!(
+            args.status,
+            ThreadGoalStatus::Complete | ThreadGoalStatus::Blocked
+        ) {
             return Err(FunctionCallError::RespondToModel(
-                "update_goal can only mark the existing goal complete; pause, resume, and budget-limited status changes are controlled by the user or system"
+                "update_goal can only mark the existing goal complete or blocked; pause, resume, budget-limited, and usage-limited status changes are controlled by the user or system"
                     .to_string(),
             ));
         }
 
-        // TODO: update_goal needs a host callback before completion to flush
-        // final active-turn accounting with budget steering suppressed.
+        self.account_active_goal_progress(
+            match args.status {
+                ThreadGoalStatus::Complete => codex_state::GoalAccountingMode::ActiveOrComplete,
+                ThreadGoalStatus::Blocked => codex_state::GoalAccountingMode::ActiveOrStopped,
+                ThreadGoalStatus::Active
+                | ThreadGoalStatus::Paused
+                | ThreadGoalStatus::UsageLimited
+                | ThreadGoalStatus::BudgetLimited => unreachable!("status validated above"),
+            },
+            invocation.call_id.as_str(),
+            BudgetLimitedGoalDisposition::ClearActive,
+        )
+        .await?;
         let goal = self
             .state_db
             .thread_goals()
@@ -205,14 +231,14 @@ impl GoalToolExecutor {
                 self.thread_id,
                 codex_state::GoalUpdate {
                     objective: None,
-                    status: Some(codex_state::ThreadGoalStatus::Complete),
+                    status: Some(state_status_from_protocol(args.status)),
                     token_budget: None,
                     expected_goal_id: None,
                 },
             )
             .await
             .map_err(|err| {
-                FunctionCallError::RespondToModel(format!("failed to complete goal: {err}"))
+                FunctionCallError::RespondToModel(format!("failed to update goal: {err}"))
             })?
             .map(protocol_goal_from_state)
             .ok_or_else(|| {
@@ -220,19 +246,72 @@ impl GoalToolExecutor {
                     "cannot update goal because this thread has no goal".to_string(),
                 )
             })?;
-        self.emit_goal_updated_from_tool_call(&invocation, goal.clone());
-        goal_response(Some(goal), CompletionBudgetReport::Include)
+        let turn_id = self.accounting_state.clear_current_turn_goal();
+        self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone());
+        goal_response(
+            Some(goal),
+            if args.status == ThreadGoalStatus::Complete {
+                CompletionBudgetReport::Include
+            } else {
+                CompletionBudgetReport::Omit
+            },
+        )
     }
 
-    fn emit_goal_updated_from_tool_call(&self, invocation: &ToolCall, goal: ThreadGoal) {
-        // TODO: ToolCall should expose the current turn submission id so goal
-        // tool events can set ThreadGoalUpdatedEvent.turn_id exactly as core
-        // does today. Until then, correlate the event with the tool call id.
-        self.event_emitter.thread_goal_updated(
-            invocation.call_id.clone(),
-            /*turn_id*/ None,
-            goal,
-        );
+    fn emit_goal_updated_from_tool_call(
+        &self,
+        invocation: &ToolCall,
+        turn_id: Option<String>,
+        goal: ThreadGoal,
+    ) {
+        self.event_emitter
+            .thread_goal_updated(invocation.call_id.clone(), turn_id, goal);
+    }
+
+    async fn account_active_goal_progress(
+        &self,
+        mode: codex_state::GoalAccountingMode,
+        event_id: &str,
+        budget_limited_goal_disposition: BudgetLimitedGoalDisposition,
+    ) -> Result<Option<ThreadGoal>, FunctionCallError> {
+        let Some(turn_id) = self.accounting_state.current_turn_id() else {
+            return Ok(None);
+        };
+        let Some(snapshot) = self.accounting_state.progress_snapshot(turn_id.as_str()) else {
+            return Ok(None);
+        };
+        let outcome = self
+            .state_db
+            .thread_goals()
+            .account_thread_goal_usage(
+                self.thread_id,
+                snapshot.time_delta_seconds,
+                snapshot.token_delta,
+                mode,
+                Some(snapshot.expected_goal_id.as_str()),
+            )
+            .await
+            .map_err(|err| {
+                FunctionCallError::RespondToModel(format!("failed to account goal progress: {err}"))
+            })?;
+        Ok(match outcome {
+            codex_state::GoalAccountingOutcome::Updated(goal) => {
+                self.accounting_state.mark_progress_accounted_for_status(
+                    turn_id.as_str(),
+                    &snapshot,
+                    goal.status,
+                    budget_limited_goal_disposition,
+                );
+                let goal = protocol_goal_from_state(goal);
+                self.event_emitter.thread_goal_updated(
+                    event_id.to_string(),
+                    Some(turn_id),
+                    goal.clone(),
+                );
+                Some(goal)
+            }
+            codex_state::GoalAccountingOutcome::Unchanged(_) => None,
+        })
     }
 }
 
@@ -298,7 +377,7 @@ async fn fill_empty_thread_preview_if_possible(
     }
 }
 
-fn protocol_goal_from_state(goal: codex_state::ThreadGoal) -> ThreadGoal {
+pub(crate) fn protocol_goal_from_state(goal: codex_state::ThreadGoal) -> ThreadGoal {
     ThreadGoal {
         thread_id: goal.thread_id,
         objective: goal.objective,
@@ -322,6 +401,17 @@ fn protocol_status_from_state(status: codex_state::ThreadGoalStatus) -> ThreadGo
     }
 }
 
+fn state_status_from_protocol(status: ThreadGoalStatus) -> codex_state::ThreadGoalStatus {
+    match status {
+        ThreadGoalStatus::Active => codex_state::ThreadGoalStatus::Active,
+        ThreadGoalStatus::Paused => codex_state::ThreadGoalStatus::Paused,
+        ThreadGoalStatus::Blocked => codex_state::ThreadGoalStatus::Blocked,
+        ThreadGoalStatus::UsageLimited => codex_state::ThreadGoalStatus::UsageLimited,
+        ThreadGoalStatus::BudgetLimited => codex_state::ThreadGoalStatus::BudgetLimited,
+        ThreadGoalStatus::Complete => codex_state::ThreadGoalStatus::Complete,
+    }
+}
+
 fn completion_budget_report(goal: &ThreadGoal) -> Option<String> {
     if goal.token_budget.is_none() && goal.time_used_seconds <= 0 {
         None