moteva/server/internal/infrastructure/piagent/image_generation_test.go

package piagent

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestNewImageClientDefaultsToTenMinuteTimeout(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL: "http://example.com",
		APIKey:  "test-key",
	})
	if client == nil {
		t.Fatal("expected image client")
	}
	if client.client.Timeout != 10*time.Minute {
		t.Fatalf("expected 10 minute timeout, got %s", client.client.Timeout)
	}
}

func TestNewImageClientUsesConfiguredTimeout(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL:        "http://example.com",
		APIKey:         "test-key",
		TimeoutSeconds: 180,
	})
	if client == nil {
		t.Fatal("expected image client")
	}
	if client.client.Timeout != 3*time.Minute {
		t.Fatalf("expected configured 3 minute timeout, got %s", client.client.Timeout)
	}
}

func TestNewImageClientCapsConfiguredTimeoutAtTenMinutes(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL:        "http://example.com",
		APIKey:         "test-key",
		TimeoutSeconds: 900,
	})
	if client == nil {
		t.Fatal("expected image client")
	}
	if client.client.Timeout != 10*time.Minute {
		t.Fatalf("expected configured timeout to be capped at 10 minutes, got %s", client.client.Timeout)
	}
}

func TestNewImageClientDisablesKeepAliveReuse(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL: "http://example.com",
		APIKey:  "test-key",
	})
	transport, ok := client.client.Transport.(*http.Transport)
	if !ok {
		t.Fatalf("expected custom http transport, got %T", client.client.Transport)
	}
	if !transport.DisableKeepAlives {
		t.Fatal("expected image client to disable keep-alive reuse")
	}
}

func TestNewImageClientDefaultsInputImagesToFileTransport(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL: "http://example.com",
		APIKey:  "test-key",
	})
	if client == nil {
		t.Fatal("expected image client")
	}
	if client.inputImageTransport != "file" {
		t.Fatalf("expected file input image transport, got %q", client.inputImageTransport)
	}
}

func TestImageClientSendsTransportDefaultSizeWhenUnset(t *testing.T) {
	var requestBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"}); err != nil {
		t.Fatal(err)
	}
	if requestBody["size"] != "1024x1024" {
		t.Fatalf("expected transport default size, got %#v", requestBody["size"])
	}
	if requestBody["stream"] != true {
		t.Fatalf("expected image generations to use stream transport, got %#v", requestBody["stream"])
	}
	if _, ok := requestBody["n"]; ok {
		t.Fatalf("expected single image request to omit n for compatible transports, got %#v", requestBody["n"])
	}
}

func TestImageClientSendsIdempotencyKey(t *testing.T) {
	var idempotencyKey string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		idempotencyKey = r.Header.Get("Idempotency-Key")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", IdempotencyKey: "project-thread-image-0"}); err != nil {
		t.Fatal(err)
	}
	if idempotencyKey != "project-thread-image-0" {
		t.Fatalf("expected idempotency key header, got %q", idempotencyKey)
	}
}

func TestImageClientSendsExplicitSize(t *testing.T) {
	var requestBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", Size: "2048x1152"}); err != nil {
		t.Fatal(err)
	}
	if requestBody["size"] != "2048x1152" {
		t.Fatalf("expected explicit size, got %#v", requestBody["size"])
	}
}

func TestImageClientSendsQuality(t *testing.T) {
	var requestBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", Quality: "high"}); err != nil {
		t.Fatal(err)
	}
	if requestBody["quality"] != "high" {
		t.Fatalf("expected quality=high, got %#v", requestBody["quality"])
	}
}

func TestNormalizeImageOutputFormatRejectsPSD(t *testing.T) {
	if got := normalizeImageOutputFormat("psd"); got != "" {
		t.Fatalf("expected unsupported PSD output format to be omitted, got %q", got)
	}
}

func TestImageClientSendsOutputFormatAndParsesWebPBase64(t *testing.T) {
	var requestBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"b64_json":"d2VicC1ieXRlcw=="}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", OutputFormat: "webp"})
	if err != nil {
		t.Fatal(err)
	}
	if requestBody["output_format"] != "webp" {
		t.Fatalf("expected output_format=webp, got %#v", requestBody["output_format"])
	}
	if requestBody["stream"] != nil {
		t.Fatalf("expected explicit output format to avoid streaming, got %#v", requestBody["stream"])
	}
	if image.URL != "data:image/webp;base64,d2VicC1ieXRlcw==" {
		t.Fatalf("expected WebP data URL, got %#v", image)
	}
}

func TestImageClientSplitsImageCountIntoStreamedSingleRequests(t *testing.T) {
	var mu sync.Mutex
	requestBodies := make([]map[string]any, 0, 2)
	requestKeys := make([]string, 0, 2)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		key := r.Header.Get("Idempotency-Key")
		imageIndex := strings.TrimPrefix(key, "batch-key-")
		if imageIndex == key {
			imageIndex = "unknown"
		}
		var requestBody map[string]any
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		mu.Lock()
		requestBodies = append(requestBodies, requestBody)
		requestKeys = append(requestKeys, key)
		mu.Unlock()
		w.Header().Set("Content-Type", "application/json")
		_, _ = fmt.Fprintf(w, `{"data":[{"url":"https://example.com/image-%s.png"}]}`, imageIndex)
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "create two images", ImageRequestOptions{Model: "gpt-image-2", Count: 2, IdempotencyKey: "batch-key"})
	if err != nil {
		t.Fatal(err)
	}
	mu.Lock()
	defer mu.Unlock()
	if len(requestBodies) != 2 {
		t.Fatalf("expected two single-image requests, got %d", len(requestBodies))
	}
	for index, body := range requestBodies {
		if body["n"] != nil {
			t.Fatalf("expected split request %d to omit n, got %#v", index, body)
		}
		if body["stream"] != true {
			t.Fatalf("expected split request %d to keep stream=true, got %#v", index, body)
		}
	}
	keySet := map[string]bool{}
	for _, key := range requestKeys {
		keySet[key] = true
	}
	if !keySet["batch-key-1"] || !keySet["batch-key-2"] {
		t.Fatalf("expected suffixed idempotency keys, got %#v", requestKeys)
	}
	if len(image.URLs) != 2 || image.URLs[0] != "https://example.com/image-1.png" || image.URLs[1] != "https://example.com/image-2.png" {
		t.Fatalf("expected ordered generated images, got %#v", image.URLs)
	}
}

func TestImageClientSplitsImageCountRequestsConcurrently(t *testing.T) {
	var mu sync.Mutex
	fallbackBodies := make([]map[string]any, 0, 3)
	fallbackKeys := make([]string, 0, 3)
	activeFallbacks := 0
	maxActiveFallbacks := 0
	releaseFallbacks := make(chan struct{})

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var requestBody map[string]any
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Errorf("decode request body: %v", err)
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		mu.Lock()
		fallbackBodies = append(fallbackBodies, requestBody)
		fallbackKeys = append(fallbackKeys, r.Header.Get("Idempotency-Key"))
		index := len(fallbackBodies)
		activeFallbacks++
		if activeFallbacks > maxActiveFallbacks {
			maxActiveFallbacks = activeFallbacks
		}
		if index == 3 {
			close(releaseFallbacks)
		}
		mu.Unlock()

		select {
		case <-releaseFallbacks:
		case <-time.After(2 * time.Second):
		}

		mu.Lock()
		activeFallbacks--
		mu.Unlock()
		w.Header().Set("Content-Type", "application/json")
		_, _ = fmt.Fprintf(w, `{"data":[{"url":"https://example.com/fallback-%d.png"}]}`, index)
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "create three images", ImageRequestOptions{Model: "gpt-image-2", Count: 3, IdempotencyKey: "batch-key"})
	if err != nil {
		t.Fatal(err)
	}
	mu.Lock()
	defer mu.Unlock()
	if len(image.URLs) != 3 {
		t.Fatalf("expected three generated images, got %#v", image.URLs)
	}
	if len(fallbackBodies) != 3 {
		t.Fatalf("expected three single-image requests, got %d", len(fallbackBodies))
	}
	if maxActiveFallbacks < 2 {
		t.Fatalf("expected concurrent split requests, max active was %d", maxActiveFallbacks)
	}
	keySet := map[string]bool{}
	for index, body := range fallbackBodies {
		if body["n"] != nil {
			t.Fatalf("expected split request %d to omit n, got %#v", index, body)
		}
		keySet[fallbackKeys[index]] = true
	}
	for _, key := range []string{"batch-key-1", "batch-key-2", "batch-key-3"} {
		if !keySet[key] {
			t.Fatalf("expected fallback idempotency key %q, got %#v", key, fallbackKeys)
		}
	}
}

func TestImageClientUploadsDataReferenceImagesAsMultipart(t *testing.T) {
	var requestPath string
	var requestContentType string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestPath = r.URL.Path
		requestContentType = r.Header.Get("Content-Type")
		if err := r.ParseMultipartForm(32 << 20); err != nil {
			t.Errorf("parse multipart form: %v", err)
			return
		}
		file, header, err := r.FormFile("image")
		if err != nil {
			t.Errorf("expected image file field: %v", err)
			return
		}
		defer file.Close()
		data, err := io.ReadAll(file)
		if err != nil {
			t.Errorf("read image file: %v", err)
			return
		}
		if header.Filename != "reference-1.png" || header.Header.Get("Content-Type") != "image/png" || string(data) != "reference bytes" {
			t.Errorf("unexpected uploaded image: filename=%q content_type=%q data=%q", header.Filename, header.Header.Get("Content-Type"), string(data))
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	_, err := client.Generate(context.Background(), "extract product from reference", ImageRequestOptions{
		Model:  "gpt-image-2",
		Images: []string{"data:image/png;base64,cmVmZXJlbmNlIGJ5dGVz"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if requestPath != "/v1/images/edits" {
		t.Fatalf("expected image edit endpoint for references, got %s", requestPath)
	}
	if !strings.HasPrefix(requestContentType, "multipart/form-data") {
		t.Fatalf("expected multipart edit request, got %s", requestContentType)
	}
}

func TestImageClientSendsReferenceImagesByURLWhenConfigured(t *testing.T) {
	var requestBody map[string]any
	var requestPath string
	var requestContentType string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestPath = r.URL.Path
		requestContentType = r.Header.Get("Content-Type")
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key", InputImageTransport: "url"})
	_, err := client.Generate(context.Background(), "extract product from reference", ImageRequestOptions{
		Model:  "gpt-image-2",
		Images: []string{" https://example.com/reference.png ", "", "https://example.com/reference.png", "https://example.com/second.webp"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if requestPath != "/v1/images/edits" {
		t.Fatalf("expected image edit endpoint for references, got %s", requestPath)
	}
	if requestContentType != "application/json" {
		t.Fatalf("expected url transport to use JSON, got %s", requestContentType)
	}
	images, ok := requestBody["images"].([]any)
	if !ok {
		t.Fatalf("expected images array in request, got %#v", requestBody["images"])
	}
	if len(images) != 2 {
		t.Fatalf("expected normalized reference images, got %#v", images)
	}
	first, ok := images[0].(map[string]any)
	if !ok || first["image_url"] != "https://example.com/reference.png" {
		t.Fatalf("expected first image_url reference, got %#v", images[0])
	}
	second, ok := images[1].(map[string]any)
	if !ok || second["image_url"] != "https://example.com/second.webp" {
		t.Fatalf("expected second image_url reference, got %#v", images[1])
	}
}

func TestImageClientUploadsLocalhostReferenceImagesEvenWhenURLTransportConfigured(t *testing.T) {
	var sawMultipartEdit bool
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/reference.png":
			w.Header().Set("Content-Type", "image/png")
			_, _ = w.Write([]byte("local image bytes"))
		case "/v1/images/edits":
			sawMultipartEdit = true
			if !strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/form-data") {
				t.Errorf("expected multipart edit request for local image, got %s", r.Header.Get("Content-Type"))
			}
			if err := r.ParseMultipartForm(32 << 20); err != nil {
				t.Errorf("parse multipart form: %v", err)
				return
			}
			file, header, err := r.FormFile("image")
			if err != nil {
				t.Errorf("expected image file field: %v", err)
				return
			}
			defer file.Close()
			data, err := io.ReadAll(file)
			if err != nil {
				t.Errorf("read image file: %v", err)
				return
			}
			if header.Filename != "reference.png" || header.Header.Get("Content-Type") != "image/png" || string(data) != "local image bytes" {
				t.Errorf("unexpected uploaded image: filename=%q content_type=%q data=%q", header.Filename, header.Header.Get("Content-Type"), string(data))
			}
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/edited.png"}]}`))
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key", InputImageTransport: "url"})
	image, err := client.Generate(context.Background(), "edit the local image", ImageRequestOptions{
		Model:  "gpt-image-2",
		Images: []string{server.URL + "/reference.png"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/edited.png" {
		t.Fatalf("expected edited image, got %#v", image)
	}
	if !sawMultipartEdit {
		t.Fatal("expected multipart image edit request")
	}
}

func TestImageClientUploadsInternalReferenceImagesAsMultipart(t *testing.T) {
	var sawEdit bool
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/reference.png":
			w.Header().Set("Content-Type", "image/png")
			_, _ = w.Write([]byte("internal image bytes"))
		case "/v1/images/edits":
			sawEdit = true
			if !strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/form-data") {
				t.Errorf("expected multipart edit request, got %s", r.Header.Get("Content-Type"))
			}
			if err := r.ParseMultipartForm(32 << 20); err != nil {
				t.Errorf("parse multipart form: %v", err)
				return
			}
			if r.FormValue("model") != "gpt-image-2" || r.FormValue("prompt") != "edit the internal image" || r.FormValue("size") != "1024x1024" {
				t.Errorf("unexpected multipart fields: model=%q prompt=%q size=%q", r.FormValue("model"), r.FormValue("prompt"), r.FormValue("size"))
			}
			file, header, err := r.FormFile("image")
			if err != nil {
				t.Errorf("expected image file field: %v", err)
				return
			}
			defer file.Close()
			data, err := io.ReadAll(file)
			if err != nil {
				t.Errorf("read image file: %v", err)
				return
			}
			if header.Filename != "reference.png" || header.Header.Get("Content-Type") != "image/png" || string(data) != "internal image bytes" {
				t.Errorf("unexpected uploaded image: filename=%q content_type=%q data=%q", header.Filename, header.Header.Get("Content-Type"), string(data))
			}
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/edited.png"}]}`))
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "edit the internal image", ImageRequestOptions{
		Model:  "gpt-image-2",
		Size:   "1024x1024",
		Images: []string{server.URL + "/reference.png"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/edited.png" {
		t.Fatalf("expected edited image, got %#v", image)
	}
	if !sawEdit {
		t.Fatal("expected image edit request")
	}
}

func TestImageClientParsesAlternateImageResponseShapes(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"output":[{"image_url":"https://example.com/generated.webp"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/generated.webp" {
		t.Fatalf("expected generated image url, got %#v", image)
	}
}

func TestImageClientParsesStreamingImageResponse(t *testing.T) {
	var requestBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = w.Write([]byte("event: image_generation.completed\n"))
		_, _ = w.Write([]byte(`data: {"type":"image_generation.completed","result":"Z2VuZXJhdGVkIGJ5dGVz"}` + "\n\n"))
		_, _ = w.Write([]byte("data: [DONE]\n\n"))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"})
	if err != nil {
		t.Fatal(err)
	}
	if requestBody["stream"] != true {
		t.Fatalf("expected stream request body, got %#v", requestBody)
	}
	if image.URL != "data:image/png;base64,Z2VuZXJhdGVkIGJ5dGVz" {
		t.Fatalf("expected streamed base64 image, got %#v", image)
	}
}

func TestImageClientFallsBackWhenStreamIsUnsupported(t *testing.T) {
	calls := 0
	streamValues := make([]any, 0, 2)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		var requestBody map[string]any
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		streamValues = append(streamValues, requestBody["stream"])
		if calls == 1 {
			http.Error(w, `{"error":{"message":"unknown parameter: stream"}}`, http.StatusBadRequest)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/image.png" {
		t.Fatalf("expected fallback image, got %#v", image)
	}
	if calls != 2 {
		t.Fatalf("expected stream then json fallback, got %d calls", calls)
	}
	if len(streamValues) != 2 || streamValues[0] != true || streamValues[1] != nil {
		t.Fatalf("expected first request stream=true and fallback without stream, got %#v", streamValues)
	}
}

func TestImageClientFallsBackAfterIdempotentStreamFailures(t *testing.T) {
	calls := 0
	streamValues := make([]any, 0, 4)
	keys := make([]string, 0, 4)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		keys = append(keys, r.Header.Get("Idempotency-Key"))
		var requestBody map[string]any
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		streamValues = append(streamValues, requestBody["stream"])
		if calls <= 3 {
			http.Error(w, `{"error":{"message":"temporary stream failure"}}`, http.StatusServiceUnavailable)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/recovered.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{
		Model:          "gpt-image-2",
		IdempotencyKey: "same-generation",
	})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/recovered.png" {
		t.Fatalf("expected fallback image, got %#v", image)
	}
	if calls != 4 {
		t.Fatalf("expected three stream attempts then json fallback, got %d calls", calls)
	}
	if len(streamValues) != 4 || streamValues[0] != true || streamValues[1] != true || streamValues[2] != true || streamValues[3] != nil {
		t.Fatalf("expected stream retries then json fallback, got %#v", streamValues)
	}
	for _, key := range keys {
		if key != "same-generation" {
			t.Fatalf("expected idempotency key to be preserved, got %#v", keys)
		}
	}
}

func TestImageClientRetriesTransientFailures(t *testing.T) {
	calls := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		if calls == 1 {
			http.Error(w, `{"error":"temporary overload"}`, http.StatusServiceUnavailable)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"}); err != nil {
		t.Fatal(err)
	}
	if calls != 2 {
		t.Fatalf("expected retry after transient failure, got %d calls", calls)
	}
}

func TestImageClientDoesNotRetryRequestTimeout(t *testing.T) {
	calls := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		time.Sleep(150 * time.Millisecond)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	client.client.Timeout = 20 * time.Millisecond
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"}); err == nil {
		t.Fatal("expected timeout error")
	}
	if calls != 1 {
		t.Fatalf("expected timeout to fail fast without retry, got %d calls", calls)
	}
}

func TestImageClientRetriesIdempotentResponseHeaderTimeout(t *testing.T) {
	calls := 0
	var keys []string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		keys = append(keys, r.Header.Get("Idempotency-Key"))
		if calls == 1 {
			time.Sleep(80 * time.Millisecond)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/recovered.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	transport, ok := client.client.Transport.(*http.Transport)
	if !ok {
		t.Fatalf("expected http transport, got %T", client.client.Transport)
	}
	transport.ResponseHeaderTimeout = 20 * time.Millisecond
	client.client.Timeout = time.Second

	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", IdempotencyKey: "same-generation"})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/recovered.png" {
		t.Fatalf("expected recovered image, got %#v", image)
	}
	if calls != 2 {
		t.Fatalf("expected one retry after timeout, got %d calls", calls)
	}
	if len(keys) != 2 || keys[0] != "same-generation" || keys[1] != "same-generation" {
		t.Fatalf("expected retry to keep idempotency key, got %#v", keys)
	}
}