Preserve image generation failure metadata for premium usage limits (#57275)

## Why

Image generation usage-limit errors with the `imagegen_premium` limit ID were not recognized as structured `UsageLimitExceeded` failures.

## What changed

Recognize `imagegen_premium` alongside `image_gen`, preserving the limit ID and available reset time in image generation failure metadata.

## Testing

Extend app-server coverage to verify premium limit metadata and reset-time fallback in completion events, thread reads, and resumed threads. Add unit coverage for missing exhausted-window reset times and unrelated or absent limit IDs.

GitOrigin-RevId: 8cf991776403ddce1f45d9780b8040982fd0fcfe

codex-rs/app-server/tests/suite/v2/imagegen_extension.rs

@@ -679,24 +679,37 @@ async fn standalone_image_generation_failure_emits_terminal_item(
     Ok(())
 }
 
+#[test_case("image_gen", "x-image-gen-primary", Some(1_786_150_800); "image_count")]
+#[test_case("imagegen_premium", "x-imagegen-premium-primary", None; "imagegen_premium")]
 #[tokio::test]
-async fn image_generation_usage_limit_preserves_correlated_failure_metadata() -> Result<()> {
+async fn image_generation_usage_limit_preserves_correlated_failure_metadata(
+    limit_id: &str,
+    window_header: &str,
+    body_reset_at: Option<i64>,
+) -> Result<()> {
     let call_id = "image-run-limited";
-    let reset_at = 1_786_150_800;
+    let window_reset_at = 1_786_154_400;
+    let expected_failure = ImageGenerationFailure::UsageLimitExceeded {
+        limit_id: limit_id.to_string(),
+        resets_at: Some(body_reset_at.unwrap_or(window_reset_at)),
+    };
     let server = responses::start_mock_server().await;
     Mock::given(method("POST"))
         .and(path("/api/codex/images/generations"))
         .respond_with(
             ResponseTemplate::new(429)
-                .insert_header("x-codex-active-limit", "image_gen")
-                .insert_header("x-image-gen-primary-used-percent", "100")
-                .insert_header("x-image-gen-primary-window-minutes", "1440")
-                .insert_header("x-image-gen-primary-reset-at", reset_at.to_string())
+                .insert_header("x-codex-active-limit", limit_id)
+                .insert_header(format!("{window_header}-used-percent"), "100")
+                .insert_header(format!("{window_header}-window-minutes"), "1440")
+                .insert_header(
+                    format!("{window_header}-reset-at"),
+                    window_reset_at.to_string(),
+                )
                 .set_body_json(json!({
                     "error": {
                         "type": "usage_limit_reached",
                         "message": "image limit reached",
-                        "resets_at": reset_at,
+                        "resets_at": body_reset_at,
                         "plan_type": "plus"
                     }
                 })),
@@ -748,14 +761,9 @@ async fn image_generation_usage_limit_preserves_correlated_failure_metadata() ->
     let ThreadItemPayload::ImageGeneration(image) = completed.item.payload else {
         panic!("expected failed image-generation item");
     };
+    assert_eq!(image.id, call_id);
     assert_eq!(image.status, "failed");
-    assert_eq!(
-        image.failure,
-        Some(ImageGenerationFailure::UsageLimitExceeded {
-            limit_id: "image_gen".to_string(),
-            resets_at: Some(reset_at),
-        })
-    );
+    assert_eq!(image.failure.as_ref(), Some(&expected_failure));
 
     timeout(
         DEFAULT_READ_TIMEOUT,
@@ -775,16 +783,10 @@ async fn image_generation_usage_limit_preserves_correlated_failure_metadata() ->
         .iter()
         .flat_map(Turn::item_payloads)
         .find_map(|item| match item {
-            ThreadItemPayload::ImageGeneration(item) => item.failure.as_ref(),
+            ThreadItemPayload::ImageGeneration(item) if item.id == call_id => item.failure.as_ref(),
             _ => None,
         });
-    assert_eq!(
-        persisted_failure,
-        Some(&ImageGenerationFailure::UsageLimitExceeded {
-            limit_id: "image_gen".to_string(),
-            resets_at: Some(reset_at),
-        })
-    );
+    assert_eq!(persisted_failure, Some(&expected_failure));
 
     let resume_id = mcp
         .send_thread_resume_request(ThreadResumeParams {
@@ -799,16 +801,10 @@ async fn image_generation_usage_limit_preserves_correlated_failure_metadata() ->
         .iter()
         .flat_map(Turn::item_payloads)
         .find_map(|item| match item {
-            ThreadItemPayload::ImageGeneration(item) => item.failure.as_ref(),
+            ThreadItemPayload::ImageGeneration(item) if item.id == call_id => item.failure.as_ref(),
             _ => None,
         });
-    assert_eq!(
-        resumed_failure,
-        Some(&ImageGenerationFailure::UsageLimitExceeded {
-            limit_id: "image_gen".to_string(),
-            resets_at: Some(reset_at),
-        })
-    );
+    assert_eq!(resumed_failure, Some(&expected_failure));
 
     Ok(())
 }

codex-rs/ext/image-generation/src/tests.rs

@@ -8,8 +8,11 @@ use codex_exec_server::LOCAL_FS;
 use codex_extension_api::ToolOutput;
 use codex_extension_api::ToolPayload;
 use codex_extension_api::ToolSpec;
+use codex_extension_items::image_generation::ImageGenerationFailure;
 use codex_extension_items::image_generation::ImageGenerationInputImage;
 use codex_protocol::ResponseItemId;
+use codex_protocol::error::CodexErr;
+use codex_protocol::error::UsageLimitReachedError;
 use codex_protocol::models::ContentItem;
 use codex_protocol::models::DEFAULT_IMAGE_DETAIL;
 use codex_protocol::models::FunctionCallOutputBody;
@@ -32,6 +35,7 @@ use super::ImagegenArgs;
 use super::PreparedImageRequest;
 use super::imagegen_tool_spec;
 use super::request_for_call_args;
+use super::usage_limit_failure;
 use crate::IMAGE_GEN_NAMESPACE;
 use crate::IMAGEGEN_TOOL_NAME;
 use crate::artifact::image_generation_artifact_path;
@@ -39,6 +43,39 @@ use crate::artifact::image_generation_output_hint;
 
 const RESULT: &str = "cG5n";
 
+#[test]
+fn test_usage_limit_failure_without_an_exhausted_window_reset() {
+    for (limit_id, expected) in [
+        (
+            Some("imagegen_premium"),
+            Some(ImageGenerationFailure::UsageLimitExceeded {
+                limit_id: "imagegen_premium".to_string(),
+                resets_at: None,
+            }),
+        ),
+        (Some("premium"), None),
+        (None, None),
+    ] {
+        let error = CodexErr::UsageLimitReached(UsageLimitReachedError {
+            plan_type: None,
+            resets_at: None,
+            limit_window_minutes: None,
+            rate_limits: Some(Box::new(
+                serde_json::from_value(serde_json::json!({
+                    "limit_id": limit_id,
+                    "primary": {"used_percent": 100.0},
+                    "secondary": {"used_percent": 50.0, "resets_at": 3000}
+                }))
+                .expect("rate limit snapshot"),
+            )),
+            promo_message: None,
+            rate_limit_reached_type: None,
+        });
+
+        assert_eq!(usage_limit_failure(&error), expected);
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn referenced_files_record_the_exact_bytes_read_for_each_request() {
     let directory = tempfile::tempdir().expect("temporary directory");

codex-rs/ext/image-generation/src/tool.rs

@@ -323,7 +323,7 @@ fn usage_limit_failure(error: &CodexErr) -> Option<ImageGenerationFailure> {
     };
     let rate_limits = usage_limit.rate_limits.as_deref()?;
     let limit_id = rate_limits.limit_id.as_deref()?;
-    if limit_id != "image_gen" {
+    if !matches!(limit_id, "image_gen" | "imagegen_premium") {
         return None;
     }