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
This commit is contained in:
Chris Dong
2026-09-09 00:34:45 +00:00
committed by copyberry
parent 102e1763b9
commit 929389f596
16 changed files with 107 additions and 8 deletions
@@ -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()
+1
View File
@@ -903,6 +903,7 @@ pub(crate) struct CodexImageGenerationEventParams {
pub(crate) saved_path_present: bool,
pub(crate) transparent_background: Option<bool>,
pub(crate) imagegen_request_id: Option<String>,
pub(crate) generation_id: Option<String>,
}
#[derive(Serialize)]
+1
View File
@@ -2932,6 +2932,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option<TrackEventRequest> {
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(),
},
},
))
@@ -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,
}),
],
}
@@ -994,6 +994,7 @@ impl From<CoreTurnItem> for ThreadItem {
failure: None,
saved_path: image.saved_path,
imagegen_request_id: None,
generation_id: None,
})
}
CoreTurnItem::EnteredReviewMode(review) => ThreadItem::EnteredReviewMode {
@@ -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,
}),
]);
@@ -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"))
+40
View File
@@ -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<_>>(),
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());
+2
View File
@@ -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<String>,
}
@@ -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,
+12 -2
View File
@@ -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
@@ -46,4 +46,9 @@ pub struct ImageGenerationItem {
#[schemars(skip)]
#[ts(skip)]
pub imagegen_request_id: Option<String>,
/// Backend generation ID for this exact image, retained only for in-process analytics.
#[serde(skip)]
#[schemars(skip)]
#[ts(skip)]
pub generation_id: Option<String>,
}
+5
View File
@@ -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::<ExtensionItem>(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]
@@ -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,
)),
@@ -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 {
@@ -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,