agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add goal extension telemetry parity (#24615)

Browse Source 
## 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.
This commit is contained in:
jif-oai
2026-05-26 19:48:32 +02:00
committed by GitHub
Unverified
parent db9cb04fb6
commit 08504e86fb
8 changed files with 237 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -2968,6 +2968,7 @@ dependencies = [
"chrono",
"codex-core",
"codex-extension-api",
"codex-otel",
"codex-protocol",
"codex-state",
"codex-tools",
codex-rs/ext/goal/Cargo.toml
+1
View File
@@ -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
+46
View File
@@ -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<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,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<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
+1
View File
@@ -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
+84
View File
@@ -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
+49
View File
@@ -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)
}))
}
}
codex-rs/ext/goal/src/tool.rs
+41 -1
View File
@@ -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<codex_state::StateRuntime>,
accounting_state: Arc<GoalAccountingState>,
event_emitter: GoalEventEmitter,
metrics: GoalMetrics,
}
#[derive(Clone, Copy)]
@@ -74,6 +76,7 @@ impl GoalToolExecutor {
state_db: Arc<codex_state::StateRuntime>,
accounting_state: Arc<GoalAccountingState>,
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<codex_state::StateRuntime>,
accounting_state: Arc<GoalAccountingState>,
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<codex_state::StateRuntime>,
accounting_state: Arc<GoalAccountingState>,
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<Option<codex_state::ThreadGoalStatus>, 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<T>(arguments: &str) -> Result<T, FunctionCallError>
codex-rs/ext/goal/tests/goal_extension_backend.rs
+14 -2
View File
@@ -594,7 +594,13 @@ async fn installed_tools(
thread_id: ThreadId,
) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
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<Self> {
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());