From 08504e86fbbb3951796d9eb2f97d9c67b1a41b00 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 26 May 2026 19:48:32 +0200 Subject: [PATCH] Add goal extension telemetry parity (#24615) ## Why `core/src/goals.rs` already emits OTEL metrics for goal creation, resume, terminal transitions, token counts, and duration. As `/goal` moves into `ext/goal`, the extension needs to preserve that telemetry contract instead of only emitting app-visible `ThreadGoalUpdated` events. This keeps the existing `codex.goal.*` metric surface intact while goal lifecycle ownership shifts toward the extension. ## What changed - Added an extension-local `GoalMetrics` helper that records the existing `codex.goal.*` counters and histograms through `codex-otel`. - Threaded an optional `MetricsClient` through `install_with_backend`, `GoalExtension`, `GoalRuntimeHandle`, and `GoalToolExecutor`. - Emitted created, resumed, and terminal goal metrics from the extension paths that create goals, restore active goals on thread resume, account budget limits, complete or block goals, and handle external goal mutations. - Updated existing goal extension test setup callsites to pass `None` for metrics when instrumentation is not under test. ## Verification Not run locally. --- codex-rs/Cargo.lock | 1 + codex-rs/ext/goal/Cargo.toml | 1 + codex-rs/ext/goal/src/extension.rs | 46 ++++++++++ codex-rs/ext/goal/src/lib.rs | 1 + codex-rs/ext/goal/src/metrics.rs | 84 +++++++++++++++++++ codex-rs/ext/goal/src/runtime.rs | 49 +++++++++++ codex-rs/ext/goal/src/tool.rs | 42 +++++++++- .../ext/goal/tests/goal_extension_backend.rs | 16 +++- 8 files changed, 237 insertions(+), 3 deletions(-) create mode 100644 codex-rs/ext/goal/src/metrics.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f95eb7a29..78acd3478 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2968,6 +2968,7 @@ dependencies = [ "chrono", "codex-core", "codex-extension-api", + "codex-otel", "codex-protocol", "codex-state", "codex-tools", diff --git a/codex-rs/ext/goal/Cargo.toml b/codex-rs/ext/goal/Cargo.toml index 80dbfa74b..3bf0eca8d 100644 --- a/codex-rs/ext/goal/Cargo.toml +++ b/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 } diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 82609285b..aa0f658da 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/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,6 +20,7 @@ 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; @@ -26,6 +28,7 @@ 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 { state_dbs: Arc, event_emitter: GoalEventEmitter, + metrics: GoalMetrics, thread_manager: Weak, goals_enabled: Arc bool + Send + Sync>, } @@ -60,12 +64,14 @@ impl GoalExtension { pub(crate) fn new_with_host_capabilities( state_dbs: Arc, event_sink: Arc, + metrics_client: Option, thread_manager: Weak, 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,6 +99,7 @@ where thread_id, Arc::clone(&self.state_dbs), self.event_emitter.clone(), + self.metrics.clone(), self.thread_manager.clone(), accounting_state, enabled, @@ -100,6 +107,40 @@ where }); 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 ConfigContributor for GoalExtension @@ -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( registry: &mut ExtensionRegistryBuilder, state_dbs: Arc, + metrics_client: Option, thread_manager: Weak, goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, ) where @@ -348,6 +393,7 @@ pub fn install_with_backend( let extension = Arc::new(GoalExtension::new_with_host_capabilities( state_dbs, registry.event_sink(), + metrics_client, thread_manager, goals_enabled, )); diff --git a/codex-rs/ext/goal/src/lib.rs b/codex-rs/ext/goal/src/lib.rs index c779f462e..545d9e699 100644 --- a/codex-rs/ext/goal/src/lib.rs +++ b/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; diff --git a/codex-rs/ext/goal/src/metrics.rs b/codex-rs/ext/goal/src/metrics.rs new file mode 100644 index 000000000..0bb227ef6 --- /dev/null +++ b/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, +} + +impl GoalMetrics { + pub(crate) fn new(metrics_client: Option) -> 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, + 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, + 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, + ); + } +} diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 074f78f7e..689d0578c 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/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, event_emitter: GoalEventEmitter, + metrics: GoalMetrics, thread_manager: Weak, accounting_state: Arc, enabled: AtomicBool, @@ -61,6 +63,7 @@ impl GoalRuntimeHandle { thread_id: ThreadId, state_dbs: Arc, event_emitter: GoalEventEmitter, + metrics: GoalMetrics, thread_manager: Weak, accounting_state: Arc, 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, 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) + })) + } } diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index 39142cfe2..90ea0eb44 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -18,6 +18,7 @@ use serde::Serialize; use crate::accounting::BudgetLimitedGoalDisposition; use crate::accounting::GoalAccountingState; use crate::events::GoalEventEmitter; +use crate::metrics::GoalMetrics; use crate::spec::CREATE_GOAL_TOOL_NAME; use crate::spec::GET_GOAL_TOOL_NAME; use crate::spec::UPDATE_GOAL_TOOL_NAME; @@ -32,6 +33,7 @@ pub(crate) struct GoalToolExecutor { state_db: Arc, accounting_state: Arc, event_emitter: GoalEventEmitter, + metrics: GoalMetrics, } #[derive(Clone, Copy)] @@ -74,6 +76,7 @@ impl GoalToolExecutor { state_db: Arc, accounting_state: Arc, event_emitter: GoalEventEmitter, + metrics: GoalMetrics, ) -> Self { Self { kind: GoalToolKind::Get, @@ -81,6 +84,7 @@ impl GoalToolExecutor { state_db, accounting_state, event_emitter, + metrics, } } @@ -89,6 +93,7 @@ impl GoalToolExecutor { state_db: Arc, accounting_state: Arc, event_emitter: GoalEventEmitter, + metrics: GoalMetrics, ) -> Self { Self { kind: GoalToolKind::Create, @@ -96,6 +101,7 @@ impl GoalToolExecutor { state_db, accounting_state, event_emitter, + metrics, } } @@ -104,6 +110,7 @@ impl GoalToolExecutor { state_db: Arc, accounting_state: Arc, event_emitter: GoalEventEmitter, + metrics: GoalMetrics, ) -> Self { Self { kind: GoalToolKind::Update, @@ -111,6 +118,7 @@ impl GoalToolExecutor { state_db, accounting_state, event_emitter, + metrics, } } } @@ -191,6 +199,7 @@ impl GoalToolExecutor { let turn_id = self .accounting_state .mark_current_turn_goal_active(goal.goal_id.clone()); + self.metrics.record_created(); let goal = protocol_goal_from_state(goal); self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone()); goal_response(Some(goal), CompletionBudgetReport::Omit) @@ -224,6 +233,9 @@ impl GoalToolExecutor { BudgetLimitedGoalDisposition::ClearActive, ) .await?; + let previous_status = self + .current_goal_status_for_metrics(/*expected_goal_id*/ None) + .await?; let goal = self .state_db .thread_goals() @@ -240,12 +252,14 @@ impl GoalToolExecutor { .map_err(|err| { FunctionCallError::RespondToModel(format!("failed to update goal: {err}")) })? - .map(protocol_goal_from_state) .ok_or_else(|| { FunctionCallError::RespondToModel( "cannot update goal because this thread has no goal".to_string(), ) })?; + self.metrics + .record_terminal_if_status_changed(previous_status, &goal); + let goal = protocol_goal_from_state(goal); let turn_id = self.accounting_state.clear_current_turn_goal(); self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone()); goal_response( @@ -280,6 +294,9 @@ impl GoalToolExecutor { let Some(snapshot) = self.accounting_state.progress_snapshot(turn_id.as_str()) else { return Ok(None); }; + let previous_status = self + .current_goal_status_for_metrics(Some(snapshot.expected_goal_id.as_str())) + .await?; let outcome = self .state_db .thread_goals() @@ -296,6 +313,8 @@ impl GoalToolExecutor { })?; Ok(match outcome { codex_state::GoalAccountingOutcome::Updated(goal) => { + self.metrics + .record_terminal_if_status_changed(previous_status, &goal); self.accounting_state.mark_progress_accounted_for_status( turn_id.as_str(), &snapshot, @@ -313,6 +332,27 @@ impl GoalToolExecutor { codex_state::GoalAccountingOutcome::Unchanged(_) => None, }) } + + async fn current_goal_status_for_metrics( + &self, + expected_goal_id: Option<&str>, + ) -> Result, FunctionCallError> { + let goal = self + .state_db + .thread_goals() + .get_thread_goal(self.thread_id) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to read goal metrics status: {err}" + )) + })?; + Ok(goal.and_then(|goal| { + expected_goal_id + .is_none_or(|expected_goal_id| goal.goal_id == expected_goal_id) + .then_some(goal.status) + })) + } } fn parse_arguments(arguments: &str) -> Result diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 85253345f..0a3c76898 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -594,7 +594,13 @@ async fn installed_tools( thread_id: ThreadId, ) -> Vec>> { let mut builder = ExtensionRegistryBuilder::<()>::new(); - install_with_backend(&mut builder, runtime, Weak::new(), |_| true); + install_with_backend( + &mut builder, + runtime, + /*metrics_client*/ None, + Weak::new(), + |_| true, + ); let registry = builder.build(); let session_store = ExtensionData::new("session-1"); let thread_store = ExtensionData::new(thread_id.to_string()); @@ -629,7 +635,13 @@ impl GoalExtensionHarness { ) -> anyhow::Result { let sink = Arc::new(RecordingEventSink::default()); let mut builder = ExtensionRegistryBuilder::<()>::with_event_sink(sink.clone()); - install_with_backend(&mut builder, runtime, Weak::new(), |_| true); + install_with_backend( + &mut builder, + runtime, + /*metrics_client*/ None, + Weak::new(), + |_| true, + ); let registry = builder.build(); let session_store = ExtensionData::new("session-1"); let thread_store = ExtensionData::new(thread_id.to_string());