Merge remote-tracking branch 'origin/master' into codex/tiaoyitiao

server-rs/crates/platform-image/tests/vector_engine.rs

@@ -1,8 +1,20 @@
 use platform_image::vector_engine::{
-    GPT_IMAGE_2_MODEL, VECTOR_ENGINE_PROVIDER, VectorEngineImageSettings,
-    build_vector_engine_image_request_body, vector_engine_images_edit_url,
+    GPT_IMAGE_2_MODEL, ReferenceImage, VECTOR_ENGINE_PROVIDER, VectorEngineImageSettings,
+    build_vector_engine_image_http_client, build_vector_engine_image_request_body,
+    create_vector_engine_image_edit, vector_engine_images_edit_url,
     vector_engine_images_generation_url,
 };
+use std::{
+    sync::{
+        Arc,
+        atomic::{AtomicUsize, Ordering},
+    },
+    time::Duration,
+};
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::TcpListener,
+};
 
 #[test]
 fn vector_engine_module_exposes_provider_protocol_helpers() {
@@ -30,3 +42,70 @@ fn vector_engine_module_exposes_provider_protocol_helpers() {
         "https://vector.example/v1/images/edits"
     );
 }
+
+#[tokio::test]
+async fn vector_engine_image_edit_retries_send_timeout_once_and_succeeds() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("mock server should bind");
+    let server_addr = listener
+        .local_addr()
+        .expect("mock server address should be readable");
+    let request_count = Arc::new(AtomicUsize::new(0));
+    let request_count_for_server = Arc::clone(&request_count);
+
+    let server = tokio::spawn(async move {
+        loop {
+            let Ok((mut stream, _)) = listener.accept().await else {
+                break;
+            };
+            let request_index = request_count_for_server.fetch_add(1, Ordering::SeqCst);
+            tokio::spawn(async move {
+                let mut buffer = [0_u8; 4096];
+                let _ = stream.read(&mut buffer).await;
+                if request_index == 0 {
+                    tokio::time::sleep(Duration::from_millis(120)).await;
+                    return;
+                }
+
+                let body = r#"{"data":[{"b64_json":"iVBORw0KGgpyZXN0"}]}"#;
+                let response = format!(
+                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
+                    body.len(),
+                    body
+                );
+                let _ = stream.write_all(response.as_bytes()).await;
+            });
+        }
+    });
+
+    let settings = VectorEngineImageSettings {
+        base_url: format!("http://{server_addr}/v1"),
+        api_key: "test-key".to_string(),
+        request_timeout_ms: 40,
+    };
+    let http_client =
+        build_vector_engine_image_http_client(&settings).expect("client should build");
+    let reference_image = ReferenceImage {
+        bytes: b"reference".to_vec(),
+        mime_type: "image/png".to_string(),
+        file_name: "reference.png".to_string(),
+    };
+
+    let generated = create_vector_engine_image_edit(
+        &http_client,
+        &settings,
+        "测试提示词",
+        None,
+        "1024x1024",
+        &reference_image,
+        "测试 VectorEngine 图片编辑失败",
+    )
+    .await
+    .expect("second attempt should return generated image");
+
+    assert_eq!(generated.images.len(), 1);
+    assert_eq!(generated.images[0].mime_type, "image/png");
+    assert_eq!(request_count.load(Ordering::SeqCst), 2);
+    server.abort();
+}