From eba392762c2fcdcffa9dc4a6a826f27a88c611d5 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 26 May 2026 17:55:14 +0100 Subject: [PATCH 1/4] feat: goal metrics --- 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 | 39 +++++++++ .../ext/goal/tests/goal_extension_backend.rs | 4 +- 8 files changed, 223 insertions(+), 2 deletions(-) create mode 100644 codex-rs/ext/goal/src/metrics.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f95eb7a2993c..78acd3478dd8 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 80dbfa74b578..3bf0eca8d51c 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 82609285b7a1..aa0f658dae5c 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 c779f462e029..545d9e699a20 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 000000000000..0bb227ef6185 --- /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 074f78f7ecc1..689d0578cf1f 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 39142cfe28d3..d483fdf1c0cd 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,7 @@ impl GoalToolExecutor { BudgetLimitedGoalDisposition::ClearActive, ) .await?; + let previous_status = self.current_goal_status_for_metrics(None).await?; let goal = self .state_db .thread_goals() @@ -246,6 +256,9 @@ impl GoalToolExecutor { "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 +293,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 +312,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 +331,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 85253345f6d0..8f06f8a93e85 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -594,7 +594,7 @@ 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, 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 +629,7 @@ 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, None, Weak::new(), |_| true); let registry = builder.build(); let session_store = ExtensionData::new("session-1"); let thread_store = ExtensionData::new(thread_id.to_string()); From 88905e9cfc99106ba760193f46aecee376bf649d Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 26 May 2026 18:16:36 +0100 Subject: [PATCH 2/4] fix --- codex-rs/ext/goal/src/tool.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index d483fdf1c0cd..8de6a916c36c 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -250,7 +250,6 @@ 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(), From 2d5ce9284a0820965747ab0ee30faf870baa0342 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 26 May 2026 18:32:13 +0100 Subject: [PATCH 3/4] Fix goal backend argument lint --- .../ext/goal/tests/goal_extension_backend.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 8f06f8a93e85..0a3c76898542 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, None, 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, None, 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()); From 78f7743654e43d7fc626609ab95c987286a95fe2 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 26 May 2026 18:39:24 +0100 Subject: [PATCH 4/4] Fix goal status argument lint --- codex-rs/ext/goal/src/tool.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index 8de6a916c36c..90ea0eb44be9 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -233,7 +233,9 @@ impl GoalToolExecutor { BudgetLimitedGoalDisposition::ClearActive, ) .await?; - let previous_status = self.current_goal_status_for_metrics(None).await?; + let previous_status = self + .current_goal_status_for_metrics(/*expected_goal_id*/ None) + .await?; let goal = self .state_db .thread_goals()