From 2c59806fe04484d312e943aa7c047fe20569686c Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 20 Apr 2026 14:31:22 +0100 Subject: [PATCH] feat: add metric to track the number of turns with memory usage (#18662) Add a metric `codex.turn.memory` to know if a turn used memories or not. This is not part of the other turn metrics as a label to limit cardinality --- codex-rs/core/src/memories/citations.rs | 16 ++--- codex-rs/core/src/memories/citations_tests.rs | 17 +++-- codex-rs/core/src/session/mod.rs | 8 +++ codex-rs/core/src/state/turn.rs | 1 + codex-rs/core/src/stream_events_utils.rs | 23 +++++-- codex-rs/core/src/tasks/mod.rs | 32 +++++++++ codex-rs/core/src/tasks/mod_tests.rs | 66 +++++++++++++++++-- codex-rs/otel/src/metrics/names.rs | 1 + 8 files changed, 138 insertions(+), 26 deletions(-) diff --git a/codex-rs/core/src/memories/citations.rs b/codex-rs/core/src/memories/citations.rs index d8642880f..ca32b6e6d 100644 --- a/codex-rs/core/src/memories/citations.rs +++ b/codex-rs/core/src/memories/citations.rs @@ -42,16 +42,12 @@ pub fn parse_memory_citation(citations: Vec) -> Option { } } -pub fn get_thread_id_from_citations(citations: Vec) -> Vec { - let mut result = Vec::new(); - if let Some(memory_citation) = parse_memory_citation(citations) { - for rollout_id in memory_citation.rollout_ids { - if let Ok(thread_id) = ThreadId::try_from(rollout_id.as_str()) { - result.push(thread_id); - } - } - } - result +pub fn thread_ids_from_memory_citation(memory_citation: &MemoryCitation) -> Vec { + memory_citation + .rollout_ids + .iter() + .filter_map(|id| ThreadId::try_from(id.as_str()).ok()) + .collect() } fn parse_memory_citation_entry(line: &str) -> Option { diff --git a/codex-rs/core/src/memories/citations_tests.rs b/codex-rs/core/src/memories/citations_tests.rs index 49d4a6743..1513c8d11 100644 --- a/codex-rs/core/src/memories/citations_tests.rs +++ b/codex-rs/core/src/memories/citations_tests.rs @@ -1,10 +1,10 @@ -use super::get_thread_id_from_citations; use super::parse_memory_citation; +use super::thread_ids_from_memory_citation; use codex_protocol::ThreadId; use pretty_assertions::assert_eq; #[test] -fn get_thread_id_from_citations_extracts_thread_ids() { +fn parse_memory_citation_supports_legacy_thread_ids() { let first = ThreadId::new(); let second = ThreadId::new(); @@ -12,18 +12,25 @@ fn get_thread_id_from_citations_extracts_thread_ids() { "\n\nMEMORY.md:1-2|note=[x]\n\n\n{first}\nnot-a-uuid\n{second}\n\n" )]; - assert_eq!(get_thread_id_from_citations(citations), vec![first, second]); + let parsed = parse_memory_citation(citations).expect("memory citation should parse"); + + assert_eq!( + thread_ids_from_memory_citation(&parsed), + vec![first, second] + ); } #[test] -fn get_thread_id_from_citations_supports_legacy_rollout_ids() { +fn parse_memory_citation_supports_rollout_ids() { let thread_id = ThreadId::new(); let citations = vec![format!( "\n\n{thread_id}\n\n" )]; - assert_eq!(get_thread_id_from_citations(citations), vec![thread_id]); + let parsed = parse_memory_citation(citations).expect("memory citation should parse"); + + assert_eq!(thread_ids_from_memory_citation(&parsed), vec![thread_id]); } #[test] diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 6406bbda3..1b3815c7f 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2793,6 +2793,14 @@ impl Session { .set_mailbox_delivery_phase(MailboxDeliveryPhase::CurrentTurn); } + pub(crate) async fn record_memory_citation_for_turn(&self, sub_id: &str) { + let turn_state = self.turn_state_for_sub_id(sub_id).await; + let Some(turn_state) = turn_state else { + return; + }; + turn_state.lock().await.has_memory_citation = true; + } + async fn turn_state_for_sub_id( &self, sub_id: &str, diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index 18a68600b..f1b7bec6b 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -105,6 +105,7 @@ pub(crate) struct TurnState { mailbox_delivery_phase: MailboxDeliveryPhase, granted_permissions: Option, pub(crate) tool_calls: u64, + pub(crate) has_memory_citation: bool, pub(crate) token_usage_at_turn_start: TokenUsage, } diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index a3534fcc3..44c70b931 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -9,8 +9,8 @@ use codex_utils_stream_parser::strip_citations; use tokio_util::sync::CancellationToken; use crate::function_tool::FunctionCallError; -use crate::memories::citations::get_thread_id_from_citations; use crate::memories::citations::parse_memory_citation; +use crate::memories::citations::thread_ids_from_memory_citation; use crate::parse_turn_item; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -137,7 +137,12 @@ pub(crate) async fn record_completed_response_item( .await; } mark_thread_memory_mode_polluted_if_external_context(sess, turn_context, item).await; - record_stage1_output_usage_for_completed_item(turn_context, item).await; + let has_memory_citation = + record_stage1_output_usage_and_detect_memory_citation(turn_context, item).await; + if has_memory_citation { + sess.record_memory_citation_for_turn(&turn_context.sub_id) + .await; + } } fn response_item_may_include_external_context(item: &ResponseItem) -> bool { @@ -167,23 +172,27 @@ pub(crate) async fn mark_thread_memory_mode_polluted_if_external_context( .await; } -async fn record_stage1_output_usage_for_completed_item( +async fn record_stage1_output_usage_and_detect_memory_citation( turn_context: &TurnContext, item: &ResponseItem, -) { +) -> bool { let Some(raw_text) = raw_assistant_output_text_from_item(item) else { - return; + return false; }; let (_, citations) = strip_citations(&raw_text); - let thread_ids = get_thread_id_from_citations(citations); + let Some(memory_citation) = parse_memory_citation(citations) else { + return false; + }; + let thread_ids = thread_ids_from_memory_citation(&memory_citation); if thread_ids.is_empty() { - return; + return true; } if let Some(db) = state_db::get_state_db(turn_context.config.as_ref()).await { let _ = db.record_stage1_output_usage(&thread_ids).await; } + true } /// Handle a completed output item from the model stream, recording it and diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 2064a5a5d..e7030d1d9 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -35,6 +35,7 @@ use codex_login::AuthManager; use codex_models_manager::manager::ModelsManager; use codex_otel::SessionTelemetry; use codex_otel::TURN_E2E_DURATION_METRIC; +use codex_otel::TURN_MEMORY_METRIC; use codex_otel::TURN_NETWORK_PROXY_METRIC; use codex_otel::TURN_TOKEN_USAGE_METRIC; use codex_otel::TURN_TOOL_CALL_METRIC; @@ -96,6 +97,29 @@ fn emit_turn_network_proxy_metric( ); } +fn emit_turn_memory_metric( + session_telemetry: &SessionTelemetry, + feature_enabled: bool, + config_enabled: bool, + has_citations: bool, +) { + let read_allowed = feature_enabled && config_enabled; + session_telemetry.counter( + TURN_MEMORY_METRIC, + /*inc*/ 1, + &[ + ("read_allowed", bool_tag(read_allowed)), + ("feature_enabled", bool_tag(feature_enabled)), + ("config_use_memories", bool_tag(config_enabled)), + ("has_citations", bool_tag(has_citations)), + ], + ); +} + +fn bool_tag(value: bool) -> &'static str { + if value { "true" } else { "false" } +} + /// Thin wrapper that exposes the parts of [`Session`] task runners need. #[derive(Clone)] pub(crate) struct SessionTaskContext { @@ -409,6 +433,7 @@ impl Session { let mut pending_input = Vec::::new(); let mut should_clear_active_turn = false; let mut token_usage_at_turn_start = None; + let mut turn_had_memory_citation = false; let mut turn_tool_calls = 0_u64; let turn_state = { let mut active = self.active_turn.lock().await; @@ -428,6 +453,7 @@ impl Session { if let Some(turn_state) = turn_state { let mut ts = turn_state.lock().await; pending_input = ts.take_pending_input(); + turn_had_memory_citation = ts.has_memory_citation; turn_tool_calls = ts.tool_calls; token_usage_at_turn_start = Some(ts.token_usage_at_turn_start.clone()); } @@ -531,6 +557,12 @@ impl Session { &[("token_type", "reasoning_output"), tmp_mem], ); } + emit_turn_memory_metric( + &self.services.session_telemetry, + turn_context.features.enabled(Feature::MemoryTool), + turn_context.config.memories.use_memories, + turn_had_memory_citation, + ); let (completed_at, duration_ms) = turn_context .turn_timing_state .completed_at_and_duration_ms() diff --git a/codex-rs/core/src/tasks/mod_tests.rs b/codex-rs/core/src/tasks/mod_tests.rs index 22bf0a99c..805544f76 100644 --- a/codex-rs/core/src/tasks/mod_tests.rs +++ b/codex-rs/core/src/tasks/mod_tests.rs @@ -1,7 +1,9 @@ +use super::emit_turn_memory_metric; use super::emit_turn_network_proxy_metric; use codex_otel::MetricsClient; use codex_otel::MetricsConfig; use codex_otel::SessionTelemetry; +use codex_otel::TURN_MEMORY_METRIC; use codex_otel::TURN_NETWORK_PROXY_METRIC; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; @@ -55,8 +57,8 @@ fn attributes_to_map<'a>( .collect() } -fn metric_point(resource_metrics: &ResourceMetrics) -> (BTreeMap, u64) { - let metric = find_metric(resource_metrics, TURN_NETWORK_PROXY_METRIC); +fn metric_point(resource_metrics: &ResourceMetrics, name: &str) -> (BTreeMap, u64) { + let metric = find_metric(resource_metrics, name); match metric.data() { AggregatedMetrics::U64(data) => match data { MetricData::Sum(sum) => { @@ -84,7 +86,7 @@ fn emit_turn_network_proxy_metric_records_active_turn() { let snapshot = session_telemetry .snapshot_metrics() .expect("runtime metrics snapshot"); - let (attrs, value) = metric_point(&snapshot); + let (attrs, value) = metric_point(&snapshot, TURN_NETWORK_PROXY_METRIC); assert_eq!(value, 1); assert_eq!( @@ -109,7 +111,7 @@ fn emit_turn_network_proxy_metric_records_inactive_turn() { let snapshot = session_telemetry .snapshot_metrics() .expect("runtime metrics snapshot"); - let (attrs, value) = metric_point(&snapshot); + let (attrs, value) = metric_point(&snapshot, TURN_NETWORK_PROXY_METRIC); assert_eq!(value, 1); assert_eq!( @@ -120,3 +122,59 @@ fn emit_turn_network_proxy_metric_records_inactive_turn() { ]) ); } + +#[test] +fn emit_turn_memory_metric_records_read_allowed_with_citations() { + let session_telemetry = test_session_telemetry(); + + emit_turn_memory_metric( + &session_telemetry, + /*feature_enabled*/ true, + /*config_enabled*/ true, + /*has_citations*/ true, + ); + + let snapshot = session_telemetry + .snapshot_metrics() + .expect("runtime metrics snapshot"); + let (attrs, value) = metric_point(&snapshot, TURN_MEMORY_METRIC); + + assert_eq!(value, 1); + assert_eq!( + attrs, + BTreeMap::from([ + ("config_use_memories".to_string(), "true".to_string()), + ("feature_enabled".to_string(), "true".to_string()), + ("has_citations".to_string(), "true".to_string()), + ("read_allowed".to_string(), "true".to_string()), + ]) + ); +} + +#[test] +fn emit_turn_memory_metric_records_config_disabled_without_citations() { + let session_telemetry = test_session_telemetry(); + + emit_turn_memory_metric( + &session_telemetry, + /*feature_enabled*/ true, + /*config_enabled*/ false, + /*has_citations*/ false, + ); + + let snapshot = session_telemetry + .snapshot_metrics() + .expect("runtime metrics snapshot"); + let (attrs, value) = metric_point(&snapshot, TURN_MEMORY_METRIC); + + assert_eq!(value, 1); + assert_eq!( + attrs, + BTreeMap::from([ + ("config_use_memories".to_string(), "false".to_string()), + ("feature_enabled".to_string(), "true".to_string()), + ("has_citations".to_string(), "false".to_string()), + ("read_allowed".to_string(), "false".to_string()), + ]) + ); +} diff --git a/codex-rs/otel/src/metrics/names.rs b/codex-rs/otel/src/metrics/names.rs index 9dc718d2d..2d564bac2 100644 --- a/codex-rs/otel/src/metrics/names.rs +++ b/codex-rs/otel/src/metrics/names.rs @@ -24,6 +24,7 @@ pub const TURN_E2E_DURATION_METRIC: &str = "codex.turn.e2e_duration_ms"; pub const TURN_TTFT_DURATION_METRIC: &str = "codex.turn.ttft.duration_ms"; pub const TURN_TTFM_DURATION_METRIC: &str = "codex.turn.ttfm.duration_ms"; pub const TURN_NETWORK_PROXY_METRIC: &str = "codex.turn.network_proxy"; +pub const TURN_MEMORY_METRIC: &str = "codex.turn.memory"; pub const TURN_TOOL_CALL_METRIC: &str = "codex.turn.tool.call"; pub const TURN_TOKEN_USAGE_METRIC: &str = "codex.turn.token_usage"; pub const PROFILE_USAGE_METRIC: &str = "codex.profile.usage";