From 929389f59696171f494e3105e56853ad2ec26c98 Mon Sep 17 00:00:00 2001 From: Chris Dong Date: Wed, 9 Sep 2026 00:24:54 +0000 Subject: [PATCH] Preserve per-image generation IDs in image generation analytics (#43953) ## What changed Parse optional `generation_id` values from image API responses and carry the selected image's ID through the image generation tool into analytics events. Keep the ID out of serialized extension items, JSON schemas, and TypeScript types. Responses without an ID remain supported. ## Testing Add coverage for distinct IDs in multi-image responses and responses without IDs. Extend analytics and app-server tests to verify that the selected image's ID reaches analytics, and item tests to verify that it is omitted from serialization and TypeScript types. GitOrigin-RevId: 70a600856990140b76fdbda51a0d73b3414338b1 --- .../analytics/src/analytics_client_tests.rs | 6 +++ codex-rs/analytics/src/events.rs | 1 + codex-rs/analytics/src/reducer.rs | 1 + .../src/protocol/thread_history.rs | 5 +++ .../src/protocol/v2/item.rs | 1 + .../thread_resume_redaction.rs | 1 + .../tests/suite/v2/imagegen_extension.rs | 29 +++++++++++--- codex-rs/codex-api/src/endpoint/images.rs | 40 +++++++++++++++++++ codex-rs/codex-api/src/images.rs | 2 + .../src/tools/handlers/extension_tools.rs | 2 + codex-rs/ext/image-generation/src/tool.rs | 14 ++++++- codex-rs/ext/items/src/image_generation.rs | 5 +++ codex-rs/ext/items/src/tests.rs | 5 +++ .../local/rollout_migration/legacy_event.rs | 1 + .../src/local/rollout_migration_tests.rs | 1 + codex-rs/tui/src/chatwidget/tests/helpers.rs | 1 + 16 files changed, 107 insertions(+), 8 deletions(-) diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index 4f3f4504d1..fb570694c7 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -5336,6 +5336,7 @@ async fn image_generation_events_preserve_transparent_background_metadata() { failure: None, saved_path: None, imagegen_request_id: None, + generation_id: None, }); reducer @@ -5481,6 +5482,7 @@ async fn turn_event_counts_completed_tool_items() { failure: None, saved_path: None, imagegen_request_id: Some("req-imagegen-123".to_string()), + generation_id: Some("gen-image-123".to_string()), }), ]; @@ -5580,6 +5582,10 @@ async fn turn_event_counts_completed_tool_items() { payload["event_params"]["imagegen_request_id"], json!("req-imagegen-123") ); + assert_eq!( + payload["event_params"]["generation_id"], + json!("gen-image-123") + ); let mcp_tool_call_event = out .iter() diff --git a/codex-rs/analytics/src/events.rs b/codex-rs/analytics/src/events.rs index e2b13d60b9..d109ecb918 100644 --- a/codex-rs/analytics/src/events.rs +++ b/codex-rs/analytics/src/events.rs @@ -903,6 +903,7 @@ pub(crate) struct CodexImageGenerationEventParams { pub(crate) saved_path_present: bool, pub(crate) transparent_background: Option, pub(crate) imagegen_request_id: Option, + pub(crate) generation_id: Option, } #[derive(Serialize)] diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index 5dfa55f690..c736552130 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -2932,6 +2932,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option { saved_path_present: item.saved_path.is_some(), transparent_background: item.transparent_background, imagegen_request_id: item.imagegen_request_id.clone(), + generation_id: item.generation_id.clone(), }, }, )) diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 56c90427e4..9485e9e8f5 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -860,6 +860,7 @@ impl ThreadHistoryBuilder { failure: None, saved_path: None, imagegen_request_id: None, + generation_id: None, }); self.upsert_item_in_current_turn(item); } @@ -874,6 +875,7 @@ impl ThreadHistoryBuilder { failure: payload.failure.clone(), saved_path: payload.saved_path.clone(), imagegen_request_id: None, + generation_id: None, }); self.upsert_item_in_current_turn(item); } @@ -2182,6 +2184,7 @@ mod tests { failure: None, saved_path: Some(saved_path.clone()), imagegen_request_id: None, + generation_id: None, }, )), started_at_ms: Some(0), @@ -2215,6 +2218,7 @@ mod tests { failure: None, saved_path: Some(saved_path), imagegen_request_id: None, + generation_id: None, })] ); } @@ -2545,6 +2549,7 @@ mod tests { ), saved_path: Some(test_path_buf("/tmp/ig_123.png").abs()), imagegen_request_id: None, + generation_id: None, }), ], } diff --git a/codex-rs/app-server-protocol/src/protocol/v2/item.rs b/codex-rs/app-server-protocol/src/protocol/v2/item.rs index b983710b03..f56386cae6 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/item.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/item.rs @@ -994,6 +994,7 @@ impl From for ThreadItem { failure: None, saved_path: image.saved_path, imagegen_request_id: None, + generation_id: None, }) } CoreTurnItem::EnteredReviewMode(review) => ThreadItem::EnteredReviewMode { diff --git a/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs b/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs index 86fe94eb36..e9c2880f3d 100644 --- a/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs +++ b/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs @@ -112,6 +112,7 @@ mod tests { failure: None, saved_path: Some(test_path_buf("/tmp/ig-1.png").abs()), imagegen_request_id: None, + generation_id: None, }), ]); diff --git a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs index 1a43813874..b5bae46cd2 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -62,7 +62,23 @@ const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Result<()> { let call_id = "image-run-1"; let server = responses::start_mock_server().await; - mount_image_response(&server).await; + Mock::given(method("POST")) + .and(path("/api/codex/images/generations")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("x-codex-imagegen-request-id", "req-imagegen-123") + .set_body_json(json!({ + "created": 1, + "background": "opaque", + "data": [ + {"b64_json": RESULT, "generation_id": "gen-image-123"}, + {"b64_json": "ignored", "generation_id": "gen-other"}, + ], + })), + ) + .expect(1) + .mount(&server) + .await; let response_mock = responses::mount_sse_sequence( &server, @@ -192,6 +208,10 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul event["event_params"]["imagegen_request_id"], json!("req-imagegen-123") ); + assert_eq!( + event["event_params"]["generation_id"], + json!("gen-image-123") + ); Ok(()) } @@ -430,6 +450,7 @@ async fn standalone_image_generation_failure_emits_terminal_item() -> Result<()> failure: None, saved_path: None, imagegen_request_id: None, + generation_id: None, }) ); @@ -700,7 +721,7 @@ async fn standalone_image_generation_is_exposed_in_code_mode_only() -> Result<() async fn standalone_image_generation_is_callable_from_code_mode_only() -> Result<()> { let call_id = "code-mode-image-run-1"; let server = responses::start_mock_server().await; - mount_image_response(&server).await; + mount_image_response_with_background(&server, "opaque").await; let response_mock = responses::mount_sse_sequence( &server, @@ -926,10 +947,6 @@ async fn wait_for_image_generation_completed( } } -async fn mount_image_response(server: &MockServer) { - mount_image_response_with_background(server, "opaque").await; -} - async fn mount_image_response_with_background(server: &MockServer, background: &str) { Mock::given(method("POST")) .and(path("/api/codex/images/generations")) diff --git a/codex-rs/codex-api/src/endpoint/images.rs b/codex-rs/codex-api/src/endpoint/images.rs index 716e4dec97..9e3b8973eb 100644 --- a/codex-rs/codex-api/src/endpoint/images.rs +++ b/codex-rs/codex-api/src/endpoint/images.rs @@ -191,6 +191,7 @@ mod tests { background: Some(ImageBackground::Opaque), data: vec![ImageData { b64_json: "REDACT".to_string(), + generation_id: None, }], quality: Some(ImageQuality::Medium), size: Some("1024x1536".to_string()), @@ -256,6 +257,45 @@ mod tests { ); } + #[tokio::test] + async fn preserves_distinct_generation_ids_on_image_data() { + let body = serde_json::to_vec(&json!({ + "created": 1, + "data": [ + {"b64_json": "first", "generation_id": "gen-first"}, + {"b64_json": "second", "generation_id": "gen-second"}, + {"b64_json": "legacy"} + ] + })) + .expect("serialize response"); + let transport = CapturingTransport::new(body); + let client = ImagesClient::new(transport, provider(), Arc::new(DummyAuth)); + + let (response, _) = client + .generate( + &ImageGenerationRequest { + prompt: "test".to_string(), + background: None, + model: "gpt-image-2".to_string(), + n: None, + quality: None, + size: None, + }, + HeaderMap::new(), + ) + .await + .expect("image response should parse"); + + assert_eq!( + response + .data + .iter() + .map(|image| image.generation_id.as_deref()) + .collect::>(), + vec![Some("gen-first"), Some("gen-second"), None] + ); + } + #[tokio::test] async fn edit_posts_typed_request_and_parses_image_response() { let transport = CapturingTransport::new(response_body()); diff --git a/codex-rs/codex-api/src/images.rs b/codex-rs/codex-api/src/images.rs index f915a5f78d..96c7c37543 100644 --- a/codex-rs/codex-api/src/images.rs +++ b/codex-rs/codex-api/src/images.rs @@ -67,4 +67,6 @@ pub struct ImageResponse { #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct ImageData { pub b64_json: String, + #[serde(default)] + pub generation_id: Option, } diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index 9a6d56f22e..7998de0c1d 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -551,6 +551,7 @@ mod tests { failure: None, saved_path: None, imagegen_request_id: None, + generation_id: None, }); let expected_completed_item = ExtensionItem::ImageGeneration(ImageGenerationItem { id: "call-image".to_string(), @@ -561,6 +562,7 @@ mod tests { failure: None, saved_path: Some(expected_path.clone()), imagegen_request_id: None, + generation_id: None, }); codex_tools::TurnItemEmitter::emit_started( &emitter, diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index 927cf6f8d9..ac36240cc4 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -158,6 +158,7 @@ impl ImageGenerationTool { failure: None, saved_path: None, imagegen_request_id: None, + generation_id: None, }, EventMsg::ImageGenerationBegin(ImageGenerationBeginEvent { call_id: call.call_id.clone(), @@ -184,10 +185,17 @@ impl ImageGenerationTool { .data .into_iter() .next() - .map(|data| (data.b64_json, transparent_background, imagegen_request_id)) + .map(|data| { + ( + data.b64_json, + transparent_background, + imagegen_request_id, + data.generation_id, + ) + }) .ok_or_else(|| ("image generation returned no image data".to_string(), None)) }); - let (result, transparent_background, imagegen_request_id) = match result { + let (result, transparent_background, imagegen_request_id, generation_id) = match result { Ok(result) => result, Err((message, failure)) => { let item = ImageGenerationItem { @@ -199,6 +207,7 @@ impl ImageGenerationTool { failure, saved_path: None, imagegen_request_id: None, + generation_id: None, }; let legacy_event = legacy_end_event(&item); call.turn_item_emitter @@ -224,6 +233,7 @@ impl ImageGenerationTool { failure: None, saved_path: saved_path.clone(), imagegen_request_id, + generation_id, }; let legacy_event = legacy_end_event(&item); call.turn_item_emitter diff --git a/codex-rs/ext/items/src/image_generation.rs b/codex-rs/ext/items/src/image_generation.rs index 35aa5cbd07..8cc20123c0 100644 --- a/codex-rs/ext/items/src/image_generation.rs +++ b/codex-rs/ext/items/src/image_generation.rs @@ -46,4 +46,9 @@ pub struct ImageGenerationItem { #[schemars(skip)] #[ts(skip)] pub imagegen_request_id: Option, + /// Backend generation ID for this exact image, retained only for in-process analytics. + #[serde(skip)] + #[schemars(skip)] + #[ts(skip)] + pub generation_id: Option, } diff --git a/codex-rs/ext/items/src/tests.rs b/codex-rs/ext/items/src/tests.rs index 1ecc621e49..33da46f236 100644 --- a/codex-rs/ext/items/src/tests.rs +++ b/codex-rs/ext/items/src/tests.rs @@ -18,6 +18,7 @@ fn completed_image_generation_item() -> ExtensionItem { failure: None, saved_path: None, imagegen_request_id: None, + generation_id: None, }) } @@ -96,17 +97,21 @@ fn image_generation_request_id_stays_internal() { panic!("expected image-generation item"); }; image.imagegen_request_id = Some("req-imagegen-123".to_string()); + image.generation_id = Some("gen-image-123".to_string()); let item = ExtensionItem::ImageGeneration(image); let value = serde_json::to_value(&item).expect("serialize extension item"); assert!(value.get("imagegenRequestId").is_none()); + assert!(value.get("generationId").is_none()); let ExtensionItem::ImageGeneration(round_tripped) = serde_json::from_value::(value).expect("deserialize extension item") else { panic!("expected image-generation item"); }; assert_eq!(round_tripped.imagegen_request_id, None); + assert_eq!(round_tripped.generation_id, None); assert!(!ImageGenerationItem::inline().contains("imagegenRequestId")); + assert!(!ImageGenerationItem::inline().contains("generationId")); } #[test] diff --git a/codex-rs/thread-store/src/local/rollout_migration/legacy_event.rs b/codex-rs/thread-store/src/local/rollout_migration/legacy_event.rs index 9f3f8e13c0..a6eb63d75a 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/legacy_event.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/legacy_event.rs @@ -168,6 +168,7 @@ pub(super) fn completed_item( failure: event.failure.clone(), saved_path: event.saved_path.clone(), imagegen_request_id: None, + generation_id: None, })), None, )), diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index 82339f04dc..12ceaf44f7 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration_tests.rs @@ -450,6 +450,7 @@ async fn migration_preserves_image_generation_failure_metadata() { }), saved_path: None, imagegen_request_id: None, + generation_id: None, }; let image_completion = RolloutItem::EventMsg(EventMsg::ImageGenerationEnd(ImageGenerationEndEvent { diff --git a/codex-rs/tui/src/chatwidget/tests/helpers.rs b/codex-rs/tui/src/chatwidget/tests/helpers.rs index 3092a265b6..85884f7336 100644 --- a/codex-rs/tui/src/chatwidget/tests/helpers.rs +++ b/codex-rs/tui/src/chatwidget/tests/helpers.rs @@ -764,6 +764,7 @@ pub(super) fn handle_image_generation_end( failure: None, saved_path, imagegen_request_id: None, + generation_id: None, }), }), /*replay_kind*/ None,