Handle Codex image stream rate limits (#346)

apps/aether-gateway/src/ai_pipeline/finalize/internal/stream_rewrite.rs

@@ -148,6 +148,7 @@ struct OpenAiImageStreamState {
     latest_image: Option<OpenAiImageFrame>,
     emitted_partial_count: u64,
     saw_upstream_partial: bool,
+    emitted_failure: bool,
 }
 
 #[derive(Clone)]
@@ -208,6 +209,7 @@ impl OpenAiImageStreamState {
             .or(event_name.as_deref())
             .unwrap_or_default();
         match event_type {
+            "error" | "response.failed" => self.handle_failed(report_context, &event),
             "response.image_generation_call.partial_image" => {
                 self.handle_image_generation_partial(report_context, &event)
             }
@@ -222,6 +224,9 @@ impl OpenAiImageStreamState {
         report_context: &Value,
         event: &Value,
     ) -> Result<Vec<u8>, GatewayError> {
+        if self.emitted_failure {
+            return Ok(Vec::new());
+        }
         if requested_partial_images(report_context) == 0 {
             return Ok(Vec::new());
         }
@@ -262,6 +267,9 @@ impl OpenAiImageStreamState {
         report_context: &Value,
         event: &Value,
     ) -> Result<Vec<u8>, GatewayError> {
+        if self.emitted_failure {
+            return Ok(Vec::new());
+        }
         let Some(item) = event.get("item").and_then(Value::as_object) else {
             return Ok(Vec::new());
         };
@@ -303,6 +311,9 @@ impl OpenAiImageStreamState {
         report_context: &Value,
         event: &Value,
     ) -> Result<Vec<u8>, GatewayError> {
+        if self.emitted_failure {
+            return Ok(Vec::new());
+        }
         if self.latest_image.is_none() {
             if let Some(result) = completed_response_image_result(event) {
                 self.latest_image = Some(OpenAiImageFrame {
@@ -334,6 +345,82 @@ impl OpenAiImageStreamState {
             }),
         )
     }
+
+    fn handle_failed(
+        &mut self,
+        report_context: &Value,
+        event: &Value,
+    ) -> Result<Vec<u8>, GatewayError> {
+        if self.emitted_failure {
+            return Ok(Vec::new());
+        }
+        self.emitted_failure = true;
+        let error = image_failure_error(event);
+        encode_json_sse(
+            Some(image_failed_event_name(report_context)),
+            &serde_json::json!({
+                "type": image_failed_event_name(report_context),
+                "error": error,
+            }),
+        )
+    }
+}
+
+fn image_failure_error(event: &Value) -> Value {
+    let mut error = event
+        .get("error")
+        .or_else(|| event.get("response").and_then(|value| value.get("error")))
+        .and_then(Value::as_object)
+        .cloned()
+        .unwrap_or_default();
+
+    if !error.contains_key("message") {
+        if let Some(message) = event
+            .get("message")
+            .and_then(Value::as_str)
+            .or_else(|| {
+                event
+                    .get("response")
+                    .and_then(|value| value.get("error"))
+                    .and_then(|value| value.get("message"))
+                    .and_then(Value::as_str)
+            })
+            .map(str::trim)
+            .filter(|value| !value.is_empty())
+        {
+            error.insert("message".to_string(), Value::String(message.to_string()));
+        }
+    }
+    if !error.contains_key("code") {
+        if let Some(code) = event
+            .get("code")
+            .or_else(|| {
+                event
+                    .get("response")
+                    .and_then(|value| value.get("error"))
+                    .and_then(|value| value.get("code"))
+            })
+            .cloned()
+        {
+            error.insert("code".to_string(), code);
+        }
+    }
+    if !error.contains_key("type") {
+        let inferred_type = error
+            .get("code")
+            .and_then(Value::as_str)
+            .filter(|value| !value.is_empty())
+            .unwrap_or("upstream_error");
+        error.insert("type".to_string(), Value::String(inferred_type.to_string()));
+    }
+    if !error.contains_key("message") {
+        error.insert(
+            "message".to_string(),
+            Value::String("Image generation failed".to_string()),
+        );
+    }
+
+    Value::Object(error)
 }
 
 fn completed_response_image_result(event: &Value) -> Option<&str> {
@@ -373,6 +460,14 @@ fn image_completed_event_name(report_context: &Value) -> &'static str {
     }
 }
 
+fn image_failed_event_name(report_context: &Value) -> &'static str {
+    if image_request_operation(report_context) == Some("edit") {
+        "image_edit.failed"
+    } else {
+        "image_generation.failed"
+    }
+}
+
 fn image_request_operation(report_context: &Value) -> Option<&str> {
     report_context
         .get("image_request")

apps/aether-gateway/src/ai_pipeline/finalize/tests_stream.rs

@@ -257,6 +257,77 @@ fn openai_image_stream_rewriter_reads_final_image_from_completed_response_output
     assert!(completed_text.contains("\"total_tokens\":3"));
 }
 
+#[test]
+fn openai_image_stream_rewriter_maps_upstream_error_to_generation_failed() {
+    let report_context = json!({
+        "provider_api_format": "openai:image",
+        "client_api_format": "openai:image",
+        "needs_conversion": false,
+        "image_request": {
+            "operation": "generate"
+        }
+    });
+    let mut rewriter =
+        maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
+
+    let output = rewriter
+        .push_chunk(
+            concat!(
+                "event: error\n",
+                "data: {\"type\":\"error\",\"error\":{\"type\":\"input-images\",\"code\":\"rate_limit_exceeded\",\"message\":\"Rate limit reached for gpt-image-2\",\"param\":null}}\n\n",
+                "event: response.failed\n",
+                "data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"Rate limit reached for gpt-image-2\"}}}\n\n"
+            )
+            .as_bytes(),
+        )
+        .expect("rewrite should succeed");
+    let output_text = utf8(output);
+    assert!(output_text.contains("event: image_generation.failed"));
+    assert_eq!(
+        output_text
+            .matches("event: image_generation.failed")
+            .count(),
+        1
+    );
+    assert!(output_text.contains("\"type\":\"image_generation.failed\""));
+    assert!(output_text.contains("\"type\":\"input-images\""));
+    assert!(output_text.contains("\"code\":\"rate_limit_exceeded\""));
+    assert!(output_text.contains("\"message\":\"Rate limit reached for gpt-image-2\""));
+    assert!(!output_text.contains("response.failed"));
+    assert!(rewriter.finish().expect("finish should succeed").is_empty());
+}
+
+#[test]
+fn openai_image_stream_rewriter_maps_response_failed_to_edit_failed() {
+    let report_context = json!({
+        "provider_api_format": "openai:image",
+        "client_api_format": "openai:image",
+        "needs_conversion": false,
+        "image_request": {
+            "operation": "edit"
+        }
+    });
+    let mut rewriter =
+        maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
+
+    let output = rewriter
+        .push_chunk(
+            concat!(
+                "event: response.failed\n",
+                "data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"slow down\"}}}\n\n"
+            )
+            .as_bytes(),
+        )
+        .expect("rewrite should succeed");
+    let output_text = utf8(output);
+    assert!(output_text.contains("event: image_edit.failed"));
+    assert!(output_text.contains("\"type\":\"image_edit.failed\""));
+    assert!(output_text.contains("\"type\":\"rate_limit_exceeded\""));
+    assert!(output_text.contains("\"code\":\"rate_limit_exceeded\""));
+    assert!(output_text.contains("\"message\":\"slow down\""));
+    assert!(rewriter.finish().expect("finish should succeed").is_empty());
+}
+
 #[test]
 fn openai_image_stream_rewriter_emits_partial_and_completed_events_for_edit() {
     let report_context = json!({