From 8126af38799f8244c75e2a34dc13974b1a96f94c Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Mon, 4 May 2026 12:46:08 -0700 Subject: [PATCH] core: preserve last model ids in feedback tags (#21026) ## Why Feedback reports do not currently surface a direct pointer to the last model call, so investigations may require searching through many requests in a session to find the bad response. Preserve the last model-side IDs at response-stream time so immediate feedback reports carry that breadcrumb. ## What changed - Record `last_model_request_id` when a Responses stream exposes an upstream request ID. - Record `last_model_response_id` when the model response completes. - Add unit coverage for the emitted feedback tags. ## Verification - `cargo test -p codex-core client::tests::response_stream_records_last_model_feedback_ids` --- codex-rs/core/src/client.rs | 8 +++ codex-rs/core/src/client_tests.rs | 81 +++++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index c60c4507d..d11fda47d 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -107,6 +107,7 @@ use tracing::warn; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; use crate::client_common::ResponseStream; +use crate::feedback_tags; use crate::flags::CODEX_RS_SSE_FIXTURE; use crate::util::emit_feedback_auth_recovery_tags; use codex_api::map_api_error; @@ -1685,6 +1686,9 @@ where let mut items_added: Vec = Vec::new(); let mut api_stream = api_stream; let upstream_request_id = upstream_request_id.as_deref(); + if let Some(upstream_request_id) = upstream_request_id { + feedback_tags!(last_model_request_id = upstream_request_id); + } loop { let event = tokio::select! { _ = consumer_dropped.cancelled() => { @@ -1721,6 +1725,7 @@ where token_usage, end_turn, }) => { + feedback_tags!(last_model_response_id = &response_id); if let Some(usage) = &token_usage { session_telemetry.sse_event_completed( usage.input_tokens, @@ -1769,6 +1774,9 @@ where extract_response_debug_context_from_api_error(&err); let upstream_request_id = upstream_request_id.or(response_debug_context.request_id.as_deref()); + if let Some(upstream_request_id) = upstream_request_id { + feedback_tags!(last_model_request_id = upstream_request_id); + } let mapped = map_api_error(err); inference_trace_attempt.record_failed( &mapped, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index e56500ba5..13c5603ba 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -31,14 +31,24 @@ use codex_rollout_trace::replay_bundle; use futures::StreamExt; use pretty_assertions::assert_eq; use serde_json::json; +use std::collections::BTreeMap; use std::collections::VecDeque; use std::pin::Pin; use std::sync::Arc; +use std::sync::Mutex; use std::task::Context; use std::task::Poll; use std::time::Duration; use tempfile::TempDir; use tokio::sync::Notify; +use tracing::Event; +use tracing::Subscriber; +use tracing::field::Visit; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::Context as LayerContext; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::registry::LookupSpan; +use tracing_subscriber::util::SubscriberInitExt; fn test_model_client(session_source: SessionSource) -> ModelClient { let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses); @@ -100,6 +110,42 @@ fn test_session_telemetry() -> SessionTelemetry { ) } +#[derive(Default)] +struct TagCollectorVisitor { + tags: BTreeMap, +} + +impl Visit for TagCollectorVisitor { + fn record_str(&mut self, field: &tracing::field::Field, value: &str) { + self.tags + .insert(field.name().to_string(), value.to_string()); + } + + fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { + self.tags + .insert(field.name().to_string(), format!("{value:?}")); + } +} + +#[derive(Clone)] +struct TagCollectorLayer { + tags: Arc>>, +} + +impl Layer for TagCollectorLayer +where + S: Subscriber + for<'a> LookupSpan<'a>, +{ + fn on_event(&self, event: &Event<'_>, _ctx: LayerContext<'_, S>) { + if event.metadata().target() != "feedback_tags" { + return; + } + let mut visitor = TagCollectorVisitor::default(); + event.record(&mut visitor); + self.tags.lock().unwrap().extend(visitor.tags); + } +} + fn started_inference_attempt(temp: &TempDir) -> anyhow::Result { let writer = Arc::new(TraceWriter::create( temp.path(), @@ -316,6 +362,41 @@ async fn dropped_response_stream_traces_cancelled_partial_output() -> anyhow::Re Ok(()) } +#[tokio::test] +async fn response_stream_records_last_model_feedback_ids() { + let tags = Arc::new(Mutex::new(BTreeMap::new())); + let _guard = tracing_subscriber::registry() + .with(TagCollectorLayer { tags: tags.clone() }) + .set_default(); + + let api_stream = futures::stream::iter([ + Ok(ResponseEvent::Created), + Ok(ResponseEvent::Completed { + response_id: "resp-123".to_string(), + token_usage: None, + end_turn: Some(true), + }), + ]); + let (mut stream, _) = super::map_response_events( + Some("req-123".to_string()), + api_stream, + test_session_telemetry(), + InferenceTraceAttempt::disabled(), + ); + + while stream.next().await.is_some() {} + + let tags = tags.lock().unwrap().clone(); + assert_eq!( + tags.get("last_model_request_id").map(String::as_str), + Some("\"req-123\"") + ); + assert_eq!( + tags.get("last_model_response_id").map(String::as_str), + Some("\"resp-123\"") + ); +} + #[tokio::test] async fn dropped_backpressured_response_stream_traces_cancelled_partial_output() -> anyhow::Result<()> {