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 e81f9e7be3..a0792ae98a 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/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, +) -> 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(()) } diff --git a/codex-rs/ext/image-generation/src/tests.rs b/codex-rs/ext/image-generation/src/tests.rs index 8adc858b5d..5e286ab96e 100644 --- a/codex-rs/ext/image-generation/src/tests.rs +++ b/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"); diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index fb154371f7..26f3badf3d 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -323,7 +323,7 @@ fn usage_limit_failure(error: &CodexErr) -> Option { }; 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; }