Add goal extension telemetry parity

codex-rs/ext/goal/Cargo.toml

@@ -17,6 +17,7 @@ workspace = true
 async-trait = { workspace = true }
 codex-core = { workspace = true }
 codex-extension-api = { workspace = true }
+codex-otel = { workspace = true }
 codex-protocol = { workspace = true }
 codex-state = { workspace = true }
 codex-tools = { workspace = true }

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

@@ -8,6 +8,7 @@ use codex_extension_api::ExtensionData;
 use codex_extension_api::ExtensionEventSink;
 use codex_extension_api::ExtensionRegistryBuilder;
 use codex_extension_api::ThreadLifecycleContributor;
+use codex_extension_api::ThreadResumeInput;
 use codex_extension_api::ThreadStartInput;
 use codex_extension_api::TokenUsageContributor;
 use codex_extension_api::ToolCallOutcome;
@@ -19,13 +20,15 @@ use codex_extension_api::TurnAbortInput;
 use codex_extension_api::TurnLifecycleContributor;
 use codex_extension_api::TurnStartInput;
 use codex_extension_api::TurnStopInput;
+use codex_otel::MetricsClient;
 use codex_protocol::ThreadId;
 use codex_protocol::protocol::ThreadGoalStatus;
 use codex_protocol::protocol::TokenUsageInfo;
 
 use crate::accounting::BudgetLimitedGoalDisposition;
 use crate::accounting::GoalAccountingState;
 use crate::events::GoalEventEmitter;
+use crate::metrics::GoalMetrics;
 use crate::runtime::GoalRuntimeHandle;
 use crate::spec::UPDATE_GOAL_TOOL_NAME;
 use crate::steering::budget_limit_steering_item;
@@ -46,6 +49,7 @@ impl GoalExtensionConfig {
 pub struct GoalExtension<C> {
     state_dbs: Arc<codex_state::StateRuntime>,
     event_emitter: GoalEventEmitter,
+    metrics: GoalMetrics,
     thread_manager: Weak<ThreadManager>,
     goals_enabled: Arc<dyn Fn(&C) -> bool + Send + Sync>,
 }
@@ -60,12 +64,14 @@ impl<C> GoalExtension<C> {
     pub(crate) fn new_with_host_capabilities(
         state_dbs: Arc<codex_state::StateRuntime>,
         event_sink: Arc<dyn ExtensionEventSink>,
+        metrics_client: Option<MetricsClient>,
         thread_manager: Weak<ThreadManager>,
         goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
     ) -> Self {
         Self {
             state_dbs,
             event_emitter: GoalEventEmitter::new(event_sink),
+            metrics: GoalMetrics::new(metrics_client),
             thread_manager,
             goals_enabled: Arc::new(goals_enabled),
         }
@@ -93,13 +99,48 @@ where
                 thread_id,
                 Arc::clone(&self.state_dbs),
                 self.event_emitter.clone(),
+                self.metrics.clone(),
                 self.thread_manager.clone(),
                 accounting_state,
                 enabled,
             )
         });
         runtime.set_enabled(enabled);
     }
+
+    async fn on_thread_resume(&self, input: ThreadResumeInput<'_>) {
+        let Some(runtime) = goal_runtime_handle(input.thread_store) else {
+            return;
+        };
+        if !runtime.is_enabled() {
+            return;
+        }
+
+        let goal = match self
+            .state_dbs
+            .thread_goals()
+            .get_thread_goal(runtime.thread_id())
+            .await
+        {
+            Ok(goal) => goal,
+            Err(err) => {
+                tracing::warn!(
+                    "failed to restore goal runtime after thread resume for {}: {err}",
+                    runtime.thread_id()
+                );
+                return;
+            }
+        };
+        match goal {
+            Some(goal) if goal.status == codex_state::ThreadGoalStatus::Active => {
+                runtime
+                    .accounting_state()
+                    .mark_idle_goal_active(goal.goal_id);
+                self.metrics.record_resumed();
+            }
+            Some(_) | None => runtime.accounting_state().clear_active_goal(),
+        }
+    }
 }
 
 impl<C> ConfigContributor<C> for GoalExtension<C>
@@ -320,18 +361,21 @@ where
                 Arc::clone(&self.state_dbs),
                 runtime.accounting_state(),
                 self.event_emitter.clone(),
+                self.metrics.clone(),
             )),
             Arc::new(GoalToolExecutor::create(
                 runtime.thread_id(),
                 Arc::clone(&self.state_dbs),
                 runtime.accounting_state(),
                 self.event_emitter.clone(),
+                self.metrics.clone(),
             )),
             Arc::new(GoalToolExecutor::update(
                 runtime.thread_id(),
                 Arc::clone(&self.state_dbs),
                 runtime.accounting_state(),
                 self.event_emitter.clone(),
+                self.metrics.clone(),
             )),
         ]
     }
@@ -340,6 +384,7 @@ where
 pub fn install_with_backend<C>(
     registry: &mut ExtensionRegistryBuilder<C>,
     state_dbs: Arc<codex_state::StateRuntime>,
+    metrics_client: Option<MetricsClient>,
     thread_manager: Weak<ThreadManager>,
     goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
 ) where
@@ -348,6 +393,7 @@ pub fn install_with_backend<C>(
     let extension = Arc::new(GoalExtension::new_with_host_capabilities(
         state_dbs,
         registry.event_sink(),
+        metrics_client,
         thread_manager,
         goals_enabled,
     ));

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

@@ -7,6 +7,7 @@
 mod accounting;
 mod events;
 mod extension;
+mod metrics;
 mod runtime;
 mod spec;
 mod steering;

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

@@ -0,0 +1,84 @@
+use codex_otel::GOAL_BLOCKED_METRIC;
+use codex_otel::GOAL_BUDGET_LIMITED_METRIC;
+use codex_otel::GOAL_COMPLETED_METRIC;
+use codex_otel::GOAL_CREATED_METRIC;
+use codex_otel::GOAL_DURATION_SECONDS_METRIC;
+use codex_otel::GOAL_RESUMED_METRIC;
+use codex_otel::GOAL_TOKEN_COUNT_METRIC;
+use codex_otel::GOAL_USAGE_LIMITED_METRIC;
+use codex_otel::MetricsClient;
+
+#[derive(Clone, Default)]
+pub(crate) struct GoalMetrics {
+    metrics_client: Option<MetricsClient>,
+}
+
+impl GoalMetrics {
+    pub(crate) fn new(metrics_client: Option<MetricsClient>) -> Self {
+        Self { metrics_client }
+    }
+
+    pub(crate) fn record_created(&self) {
+        let Some(metrics_client) = self.metrics_client.as_ref() else {
+            return;
+        };
+        let _ = metrics_client.counter(GOAL_CREATED_METRIC, /*inc*/ 1, &[]);
+    }
+
+    pub(crate) fn record_resumed(&self) {
+        let Some(metrics_client) = self.metrics_client.as_ref() else {
+            return;
+        };
+        let _ = metrics_client.counter(GOAL_RESUMED_METRIC, /*inc*/ 1, &[]);
+    }
+
+    pub(crate) fn record_resumed_if_status_changed(
+        &self,
+        previous_status: Option<codex_state::ThreadGoalStatus>,
+        goal_status: codex_state::ThreadGoalStatus,
+    ) {
+        if goal_status == codex_state::ThreadGoalStatus::Active
+            && matches!(
+                previous_status,
+                Some(
+                    codex_state::ThreadGoalStatus::Paused
+                        | codex_state::ThreadGoalStatus::Blocked
+                        | codex_state::ThreadGoalStatus::UsageLimited
+                )
+            )
+        {
+            self.record_resumed();
+        }
+    }
+
+    pub(crate) fn record_terminal_if_status_changed(
+        &self,
+        previous_status: Option<codex_state::ThreadGoalStatus>,
+        goal: &codex_state::ThreadGoal,
+    ) {
+        if previous_status == Some(goal.status) {
+            return;
+        }
+
+        let counter = match goal.status {
+            codex_state::ThreadGoalStatus::Blocked => GOAL_BLOCKED_METRIC,
+            codex_state::ThreadGoalStatus::UsageLimited => GOAL_USAGE_LIMITED_METRIC,
+            codex_state::ThreadGoalStatus::BudgetLimited => GOAL_BUDGET_LIMITED_METRIC,
+            codex_state::ThreadGoalStatus::Complete => GOAL_COMPLETED_METRIC,
+            codex_state::ThreadGoalStatus::Active | codex_state::ThreadGoalStatus::Paused => {
+                return;
+            }
+        };
+        let Some(metrics_client) = self.metrics_client.as_ref() else {
+            return;
+        };
+        let status_tag = [("status", goal.status.as_str())];
+        let _ = metrics_client.counter(counter, /*inc*/ 1, &[]);
+        let _ = metrics_client.histogram(GOAL_TOKEN_COUNT_METRIC, goal.tokens_used, &status_tag);
+        let _ = metrics_client.histogram(
+            GOAL_DURATION_SECONDS_METRIC,
+            goal.time_used_seconds,
+            &status_tag,
+        );
+    }
+}

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

@@ -11,6 +11,7 @@ use codex_protocol::protocol::ThreadGoal;
 use crate::accounting::BudgetLimitedGoalDisposition;
 use crate::accounting::GoalAccountingState;
 use crate::events::GoalEventEmitter;
+use crate::metrics::GoalMetrics;
 use crate::steering::objective_updated_steering_item;
 use crate::tool::protocol_goal_from_state;
 
@@ -23,6 +24,7 @@ struct GoalRuntimeInner {
     thread_id: ThreadId,
     state_dbs: Arc<codex_state::StateRuntime>,
     event_emitter: GoalEventEmitter,
+    metrics: GoalMetrics,
     thread_manager: Weak<ThreadManager>,
     accounting_state: Arc<GoalAccountingState>,
     enabled: AtomicBool,
@@ -61,6 +63,7 @@ impl GoalRuntimeHandle {
         thread_id: ThreadId,
         state_dbs: Arc<codex_state::StateRuntime>,
         event_emitter: GoalEventEmitter,
+        metrics: GoalMetrics,
         thread_manager: Weak<ThreadManager>,
         accounting_state: Arc<GoalAccountingState>,
         enabled: bool,
@@ -70,6 +73,7 @@ impl GoalRuntimeHandle {
                 thread_id,
                 state_dbs,
                 event_emitter,
+                metrics,
                 thread_manager,
                 accounting_state,
                 enabled: AtomicBool::new(enabled),
@@ -127,6 +131,21 @@ impl GoalRuntimeHandle {
             return Ok(());
         }
 
+        let replaced_existing_goal = previous_goal
+            .as_ref()
+            .is_some_and(|previous_goal| previous_goal.goal_id != goal.goal_id);
+        if previous_goal.is_none() || replaced_existing_goal {
+            self.inner.metrics.record_created();
+        }
+        let previous_status = previous_goal
+            .as_ref()
+            .and_then(|previous_goal| (!replaced_existing_goal).then_some(previous_goal.status));
+        self.inner
+            .metrics
+            .record_resumed_if_status_changed(previous_status, goal.status);
+        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
@@ -202,6 +221,9 @@ impl GoalRuntimeHandle {
         let Some(snapshot) = accounting.progress_snapshot(turn_id) else {
             return Ok(None);
         };
+        let previous_status = self
+            .current_goal_status_for_metrics(Some(snapshot.expected_goal_id.as_str()))
+            .await?;
         let outcome = self
             .inner
             .state_dbs
@@ -218,6 +240,9 @@ impl GoalRuntimeHandle {
         Ok(match outcome {
             codex_state::GoalAccountingOutcome::Updated(goal) => {
                 let goal_id = goal.goal_id.clone();
+                self.inner
+                    .metrics
+                    .record_terminal_if_status_changed(previous_status, &goal);
                 accounting.mark_progress_accounted_for_status(
                     turn_id,
                     &snapshot,
@@ -246,6 +271,9 @@ impl GoalRuntimeHandle {
         let Some(snapshot) = accounting.idle_progress_snapshot() else {
             return Ok(None);
         };
+        let previous_status = self
+            .current_goal_status_for_metrics(Some(snapshot.expected_goal_id.as_str()))
+            .await?;
         let outcome = self
             .inner
             .state_dbs
@@ -262,6 +290,9 @@ impl GoalRuntimeHandle {
         Ok(match outcome {
             codex_state::GoalAccountingOutcome::Updated(goal) => {
                 let goal_id = goal.goal_id.clone();
+                self.inner
+                    .metrics
+                    .record_terminal_if_status_changed(previous_status, &goal);
                 accounting.mark_idle_progress_accounted_for_status(
                     &snapshot,
                     goal.status,
@@ -281,4 +312,22 @@ impl GoalRuntimeHandle {
             }
         })
     }
+
+    async fn current_goal_status_for_metrics(
+        &self,
+        expected_goal_id: Option<&str>,
+    ) -> Result<Option<codex_state::ThreadGoalStatus>, String> {
+        let goal = self
+            .inner
+            .state_dbs
+            .thread_goals()
+            .get_thread_goal(self.thread_id())
+            .await
+            .map_err(|err| err.to_string())?;
+        Ok(goal.and_then(|goal| {
+            expected_goal_id
+                .is_none_or(|expected_goal_id| goal.goal_id == expected_goal_id)
+                .then_some(goal.status)
+        }))
+    }
 }