feat(piagent): support image quality and output_format in image client

Add quality and output_format to ImageRequestOptions, plumb them through
the JSON/multipart request bodies, generalize asset URL collection to
non-image formats, and add GenerateLayeredDocument for PSD output.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

server/internal/infrastructure/piagent/design_agent.go

@@ -303,6 +303,9 @@ func (a *DesignAgent) buildImagePrompts(ctx context.Context, req design.AgentReq
 
 func plannedGenerationCount(req design.AgentRequest) int {
 	count := req.TaskPlan.ImageCount
+	if count <= 0 && req.ImageCount > 0 {
+		count = req.ImageCount
+	}
 	if count <= 0 && len(req.TaskPlan.ImageTasks) > 0 {
 		count = len(req.TaskPlan.ImageTasks)
 	}
@@ -516,6 +519,7 @@ func (a *DesignAgent) generateImageWithFallbackCount(ctx context.Context, req de
 	image, err := a.imageAgent.Generate(ctx, prompt, ImageRequestOptions{
 		Model:          req.ImageModel,
 		Size:           req.ImageSize,
+		Quality:        req.ImageQuality,
 		Count:          count,
 		Images:         referenceImages,
 		IdempotencyKey: imageGenerationIdempotencyKey(req, index, prompt, false),
@@ -532,6 +536,7 @@ func (a *DesignAgent) generateImageWithFallbackCount(ctx context.Context, req de
 	fallbackImage, fallbackErr := a.imageAgent.Generate(ctx, fallback, ImageRequestOptions{
 		Model:          req.ImageModel,
 		Size:           req.ImageSize,
+		Quality:        req.ImageQuality,
 		Count:          count,
 		Images:         referenceImages,
 		IdempotencyKey: imageGenerationIdempotencyKey(req, index, fallback, true),
@@ -553,6 +558,7 @@ func imageGenerationIdempotencyKey(req design.AgentRequest, index int, prompt st
 		strings.TrimSpace(req.Prompt),
 		strings.TrimSpace(req.ImageModel),
 		strings.TrimSpace(req.ImageSize),
+		strings.TrimSpace(req.ImageQuality),
 		strconv.Itoa(plannedGenerationCount(req)),
 		strconv.Itoa(index),
 		kind,

server/internal/infrastructure/piagent/image_generation.go

@@ -46,8 +46,10 @@ type GeneratedImage struct {
 type ImageRequestOptions struct {
 	Model          string
 	Size           string
+	Quality        string
 	Count          int
 	Images         []string
+	OutputFormat   string
 	IdempotencyKey string
 }
 
@@ -113,6 +115,68 @@ func normalizeInputImageTransport(value string) string {
 	}
 }
 
+func normalizeImageOutputFormat(value string) string {
+	switch strings.ToLower(strings.TrimSpace(value)) {
+	case "psd", "png", "jpeg", "jpg", "webp":
+		if strings.EqualFold(strings.TrimSpace(value), "jpg") {
+			return "jpeg"
+		}
+		return strings.ToLower(strings.TrimSpace(value))
+	default:
+		return ""
+	}
+}
+
+func normalizeImageQuality(value string) string {
+	switch strings.ToLower(strings.TrimSpace(value)) {
+	case "auto", "low", "medium", "high":
+		return strings.ToLower(strings.TrimSpace(value))
+	default:
+		return ""
+	}
+}
+
+func (c *ImageClient) GenerateLayeredDocument(ctx context.Context, req design.LayeredDocumentRequest) (design.LayeredDocument, error) {
+	imageURL := strings.TrimSpace(req.ImageURL)
+	if imageURL == "" {
+		return design.LayeredDocument{}, fmt.Errorf("layered document image_url is required")
+	}
+	outputFormat := normalizeImageOutputFormat(req.OutputFormat)
+	if outputFormat == "" {
+		outputFormat = "psd"
+	}
+	prompt := strings.TrimSpace(req.Prompt)
+	if prompt == "" {
+		prompt = "Create a layered PSD from the provided image. Preserve the exact canvas and visual design. Put a single clean background layer at the bottom with all text and small foreground elements removed. Put each small standalone foreground element on its own image layer. Keep text as editable text layers when possible. Do not merge text into the background. Return only the layered PSD file."
+	}
+	size := strings.TrimSpace(req.Size)
+	if size == "" {
+		size = "auto"
+	}
+	generated, err := c.Generate(ctx, prompt, ImageRequestOptions{
+		Model:          req.Model,
+		Size:           size,
+		Count:          1,
+		Images:         []string{imageURL},
+		OutputFormat:   outputFormat,
+		IdempotencyKey: strings.TrimSpace(req.IdempotencyKey),
+	})
+	if err != nil {
+		return design.LayeredDocument{}, err
+	}
+	content := strings.TrimSpace(generated.URL)
+	if content == "" && len(generated.URLs) > 0 {
+		content = strings.TrimSpace(generated.URLs[0])
+	}
+	if content == "" {
+		return design.LayeredDocument{}, fmt.Errorf("layered document generation returned empty file")
+	}
+	return design.LayeredDocument{
+		Content:     content,
+		ContentType: outputFormatContentType(outputFormat),
+	}, nil
+}
+
 func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageRequestOptions) (GeneratedImage, error) {
 	prompt = strings.TrimSpace(prompt)
 	if prompt == "" {
@@ -134,10 +198,26 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq
 		count = 10
 	}
 	images := normalizeImageInputs(opts.Images)
+	quality := normalizeImageQuality(opts.Quality)
+	outputFormat := normalizeImageOutputFormat(opts.OutputFormat)
 
 	idempotencyKey := strings.TrimSpace(opts.IdempotencyKey)
+	if count > 1 {
+		logx.Infof(
+			"gpt image generation splitting multi-image request model=%s size=%s count=%d prompt_runes=%d images=%d output_format=%s idempotency_key=%s",
+			model,
+			size,
+			count,
+			utf8.RuneCountInString(prompt),
+			len(images),
+			outputFormat,
+			idempotencyKey,
+		)
+		return c.generateImageCountFallback(ctx, prompt, model, size, quality, count, images, outputFormat, idempotencyKey)
+	}
+
 	streamPlans := []bool{false}
-	if len(images) == 0 && count == 1 {
+	if len(images) == 0 && count == 1 && outputFormat == "" {
 		streamPlans = []bool{true, false}
 	}
 
@@ -145,12 +225,12 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq
 	startedAt := time.Now()
 	promptRunes := utf8.RuneCountInString(prompt)
 	for _, stream := range streamPlans {
-		body, contentType, err := c.imageRequestBody(ctx, model, prompt, size, count, images, stream)
+		body, contentType, err := c.imageRequestBody(ctx, model, prompt, size, quality, count, images, outputFormat, stream)
 		if err != nil {
 			return GeneratedImage{}, err
 		}
 		for attempt := 0; attempt < 3; attempt++ {
-			image, err := c.generateOnce(ctx, body, contentType, idempotencyKey, len(images) > 0, stream, count)
+			image, err := c.generateOnce(ctx, body, contentType, idempotencyKey, len(images) > 0, stream, count, outputFormat)
 			if err == nil {
 				logx.Infof(
 					"gpt image generation succeeded model=%s size=%s count=%d prompt_runes=%d attempt=%d stream=%t elapsed=%s timeout=%s images=%d idempotency_key=%s",
@@ -227,12 +307,12 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq
 			idempotencyKey,
 			lastErr,
 		)
-		return c.generateImageCountFallback(ctx, prompt, model, size, count, images, idempotencyKey)
+		return c.generateImageCountFallback(ctx, prompt, model, size, quality, count, images, outputFormat, idempotencyKey)
 	}
 	return GeneratedImage{}, lastErr
 }
 
-func (c *ImageClient) generateImageCountFallback(ctx context.Context, prompt string, model string, size string, count int, images []string, idempotencyKey string) (GeneratedImage, error) {
+func (c *ImageClient) generateImageCountFallback(ctx context.Context, prompt string, model string, size string, quality string, count int, images []string, outputFormat string, idempotencyKey string) (GeneratedImage, error) {
 	fallbackCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
@@ -258,10 +338,12 @@ func (c *ImageClient) generateImageCountFallback(ctx context.Context, prompt str
 		go func() {
 			defer wg.Done()
 			fallbackOpts := ImageRequestOptions{
-				Model:  model,
+				Model:        model,
-				Size:   size,
+				Size:         size,
-				Count:  1,
+				Quality:      quality,
-				Images: append([]string(nil), images...),
+				Count:        1,
+				Images:       append([]string(nil), images...),
+				OutputFormat: outputFormat,
 			}
 			if idempotencyKey != "" {
 				fallbackOpts.IdempotencyKey = fmt.Sprintf("%s-%d", idempotencyKey, index+1)
@@ -293,12 +375,12 @@ func (c *ImageClient) generateImageCountFallback(ctx context.Context, prompt str
 	return GeneratedImage{URL: urls[0], URLs: urls}, nil
 }
 
-func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) {
+func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt string, size string, quality string, count int, images []string, outputFormat string, stream bool) ([]byte, string, error) {
 	if len(images) > 0 {
 		if c.inputImageTransport == "url" && imageInputsCanUseURLTransport(images) {
-			return jsonImageEditBody(model, prompt, size, count, images, stream)
+			return jsonImageEditBody(model, prompt, size, quality, count, images, outputFormat, stream)
 		}
-		return c.multipartImageEditBody(ctx, model, prompt, size, count, images, stream)
+		return c.multipartImageEditBody(ctx, model, prompt, size, quality, count, images, outputFormat, stream)
 	}
 
 	requestPayload := map[string]any{
@@ -311,6 +393,12 @@ func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt
 	if size != "" {
 		requestPayload["size"] = size
 	}
+	if quality != "" {
+		requestPayload["quality"] = quality
+	}
+	if outputFormat != "" {
+		requestPayload["output_format"] = outputFormat
+	}
 	if stream {
 		requestPayload["stream"] = true
 	}
@@ -350,7 +438,7 @@ func imageInputRequiresFileUpload(value string) bool {
 	return netutil.IsLocalOrPrivateHost(host)
 }
 
-func jsonImageEditBody(model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) {
+func jsonImageEditBody(model string, prompt string, size string, quality string, count int, images []string, outputFormat string, stream bool) ([]byte, string, error) {
 	requestPayload := map[string]any{
 		"model":  model,
 		"prompt": prompt,
@@ -362,6 +450,12 @@ func jsonImageEditBody(model string, prompt string, size string, count int, imag
 	if size != "" {
 		requestPayload["size"] = size
 	}
+	if quality != "" {
+		requestPayload["quality"] = quality
+	}
+	if outputFormat != "" {
+		requestPayload["output_format"] = outputFormat
+	}
 	if stream {
 		requestPayload["stream"] = true
 	}
@@ -369,7 +463,7 @@ func jsonImageEditBody(model string, prompt string, size string, count int, imag
 	return body, "application/json", err
 }
 
-func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) {
+func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string, prompt string, size string, quality string, count int, images []string, outputFormat string, stream bool) ([]byte, string, error) {
 	var body bytes.Buffer
 	writer := multipart.NewWriter(&body)
 	fields := map[string]string{
@@ -379,9 +473,15 @@ func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string,
 	if size != "" {
 		fields["size"] = size
 	}
+	if quality != "" {
+		fields["quality"] = quality
+	}
 	if count > 1 {
 		fields["n"] = strconv.Itoa(count)
 	}
+	if outputFormat != "" {
+		fields["output_format"] = outputFormat
+	}
 	if stream {
 		fields["stream"] = "true"
 	}
@@ -557,7 +657,7 @@ func imageExtension(contentType string, fallback string) string {
 	}
 }
 
-func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType string, idempotencyKey string, edit bool, stream bool, expectedCount int) (GeneratedImage, error) {
+func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType string, idempotencyKey string, edit bool, stream bool, expectedCount int, outputFormat string) (GeneratedImage, error) {
 	endpoint := "/v1/images/generations"
 	if edit {
 		endpoint = "/v1/images/edits"
@@ -603,7 +703,7 @@ func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType
 	if err := json.NewDecoder(resp.Body).Decode(&responsePayload); err != nil {
 		return GeneratedImage{}, err
 	}
-	urls := collectGeneratedImageURLs(responsePayload)
+	urls := collectGeneratedAssetURLs(responsePayload, outputFormat)
 	if len(urls) == 0 {
 		return GeneratedImage{}, fmt.Errorf("image generation returned empty image")
 	}
@@ -668,7 +768,7 @@ func decodeImageGenerationStreamOrJSON(body io.Reader, expectedCount int) (Gener
 		if message := imageStreamFailureMessage(payload); message != "" {
 			return true, fmt.Errorf("image generation stream failed: %s", message)
 		}
-		appendURLs(collectGeneratedImageURLs(payload))
+		appendURLs(collectGeneratedAssetURLs(payload, ""))
 		return len(urls) >= expectedCount, nil
 	}
 
@@ -716,7 +816,7 @@ func decodeImageGenerationStreamOrJSON(body io.Reader, expectedCount int) (Gener
 	if err := json.Unmarshal([]byte(raw.String()), &responsePayload); err != nil {
 		return GeneratedImage{}, err
 	}
-	urls = collectGeneratedImageURLs(responsePayload)
+	urls = collectGeneratedAssetURLs(responsePayload, "")
 	if len(urls) == 0 {
 		return GeneratedImage{}, fmt.Errorf("image generation returned empty image")
 	}
@@ -786,6 +886,10 @@ func imageErrorMessage(value any) string {
 }
 
 func collectGeneratedImageURLs(payload any) []string {
+	return collectGeneratedAssetURLs(payload, "")
+}
+
+func collectGeneratedAssetURLs(payload any, outputFormat string) []string {
 	urls := make([]string, 0)
 	seen := make(map[string]struct{})
 	var walk func(any, string)
@@ -800,7 +904,7 @@ func collectGeneratedImageURLs(payload any) []string {
 				walk(childValue, key)
 			}
 		case string:
-			if imageURL := generatedImageString(key, typed); imageURL != "" {
+			if imageURL := generatedAssetString(key, typed, outputFormat); imageURL != "" {
 				if _, ok := seen[imageURL]; ok {
 					return
 				}
@@ -814,11 +918,15 @@ func collectGeneratedImageURLs(payload any) []string {
 }
 
 func generatedImageString(key string, value string) string {
+	return generatedAssetString(key, value, "")
+}
+
+func generatedAssetString(key string, value string, outputFormat string) string {
 	value = strings.TrimSpace(value)
 	if value == "" {
 		return ""
 	}
-	if strings.HasPrefix(value, "data:image/") {
+	if strings.HasPrefix(value, "data:") {
 		return value
 	}
 	normalizedKey := strings.ToLower(strings.TrimSpace(key))
@@ -826,14 +934,34 @@ func generatedImageString(key string, value string) string {
 		return value
 	}
 	if isGeneratedImageBase64Key(normalizedKey) {
-		return "data:image/png;base64," + strings.TrimPrefix(value, "data:image/png;base64,")
+		return "data:" + outputFormatContentType(outputFormat) + ";base64," + stripDataURLPrefix(value)
 	}
 	return ""
 }
 
+func stripDataURLPrefix(value string) string {
+	if _, payload, ok := strings.Cut(value, ","); ok && strings.HasPrefix(strings.TrimSpace(value), "data:") {
+		return payload
+	}
+	return value
+}
+
+func outputFormatContentType(format string) string {
+	switch normalizeImageOutputFormat(format) {
+	case "psd":
+		return "image/vnd.adobe.photoshop"
+	case "jpeg":
+		return "image/jpeg"
+	case "webp":
+		return "image/webp"
+	default:
+		return "image/png"
+	}
+}
+
 func isGeneratedImageURLKey(key string) bool {
 	switch key {
-	case "url", "image_url", "imageurl", "image", "output_url", "public_url", "publicurl":
+	case "url", "image_url", "imageurl", "image", "output_url", "public_url", "publicurl", "file_url", "fileurl", "download_url", "downloadurl":
 		return true
 	default:
 		return false
@@ -842,7 +970,7 @@ func isGeneratedImageURLKey(key string) bool {
 
 func isGeneratedImageBase64Key(key string) bool {
 	switch key {
-	case "b64_json", "b64json", "base64", "image_base64", "imagebase64", "result":
+	case "b64_json", "b64json", "base64", "image_base64", "imagebase64", "file_base64", "filebase64", "file_b64", "fileb64", "result":
 		return true
 	default:
 		return false
@@ -867,14 +995,13 @@ func isImageGenerationCountFallbackError(err error) bool {
 		return false
 	}
 	message := strings.ToLower(err.Error())
-	if !strings.Contains(message, "unknown parameter") {
+	return strings.Contains(message, "n > 1") && strings.Contains(message, "stream=false") ||
-		return false
+		strings.Contains(message, "n > 1") && strings.Contains(message, "non-streaming") ||
-	}
+		strings.Contains(message, "unknown parameter") && (strings.Contains(message, "tools[0].n") ||
-	return strings.Contains(message, "tools[0].n") ||
+			strings.Contains(message, "param:n") ||
-		strings.Contains(message, "param:n") ||
+			strings.Contains(message, "param: n") ||
-		strings.Contains(message, "param: n") ||
+			strings.Contains(message, "parameter: 'n'") ||
-		strings.Contains(message, "parameter: 'n'") ||
+			strings.Contains(message, `parameter: "n"`))
-		strings.Contains(message, `parameter: "n"`)
 }
 
 func shouldFallbackImageGenerationStream(err error, idempotencyKey string) bool {

server/internal/infrastructure/piagent/image_generation_test.go

@@ -11,6 +11,8 @@ import (
 	"sync"
 	"testing"
 	"time"
+
+	"img_infinite_canvas/internal/domain/design"
 )
 
 func TestNewImageClientDefaultsToTenMinuteTimeout(t *testing.T) {
@@ -145,36 +147,140 @@ func TestImageClientSendsExplicitSize(t *testing.T) {
 	}
 }
 
-func TestImageClientSendsImageCount(t *testing.T) {
+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-1.png"},{"url":"https://example.com/image-2.png"}]}`))
+		_, _ = 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 two images", ImageRequestOptions{Model: "gpt-image-2", Count: 2})
+	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", Quality: "high"}); err != nil {
-	if err != nil {
 		t.Fatal(err)
 	}
-	if requestBody["n"] != float64(2) {
+	if requestBody["quality"] != "high" {
-		t.Fatalf("expected image request n=2, got %#v", requestBody["n"])
+		t.Fatalf("expected quality=high, got %#v", requestBody["quality"])
-	}
-	if requestBody["stream"] != nil {
-		t.Fatalf("expected multi-image request to avoid stream, got %#v", requestBody["stream"])
-	}
-	if len(image.URLs) != 2 {
-		t.Fatalf("expected two generated images, got %#v", image.URLs)
 	}
 }
 
-func TestImageClientFallsBackToConcurrentSingleImageRequestsWhenCountUnsupported(t *testing.T) {
+func TestImageClientSendsOutputFormatAndParsesPSDBase64(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":"cHNkLWJ5dGVz"}]}`))
+	}))
+	defer server.Close()
+
+	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
+	image, err := client.Generate(context.Background(), "create layered file", ImageRequestOptions{Model: "gpt-image-2", OutputFormat: "psd"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if requestBody["output_format"] != "psd" {
+		t.Fatalf("expected output_format=psd, got %#v", requestBody["output_format"])
+	}
+	if requestBody["stream"] != nil {
+		t.Fatalf("expected psd output to avoid streaming, got %#v", requestBody["stream"])
+	}
+	if image.URL != "data:image/vnd.adobe.photoshop;base64,cHNkLWJ5dGVz" {
+		t.Fatalf("expected psd data url, got %#v", image)
+	}
+}
+
+func TestImageClientGenerateLayeredDocumentUsesPSDOutput(t *testing.T) {
+	var requestBody map[string]any
+	var requestPath string
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requestPath = r.URL.Path
+		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/layers.psd"}]}`))
+	}))
+	defer server.Close()
+
+	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key", InputImageTransport: "url"})
+	document, err := client.GenerateLayeredDocument(context.Background(), design.LayeredDocumentRequest{
+		ImageURL: "https://example.com/source.png",
+		Prompt:   "separate layers",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if requestPath != "/v1/images/edits" {
+		t.Fatalf("expected image edit endpoint, got %s", requestPath)
+	}
+	if requestBody["output_format"] != "psd" || requestBody["size"] != "auto" {
+		t.Fatalf("expected psd auto request, got %#v", requestBody)
+	}
+	if document.Content != "https://example.com/layers.psd" || document.ContentType != "image/vnd.adobe.photoshop" {
+		t.Fatalf("unexpected layered document: %#v", document)
+	}
+}
+
+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
-	var initialBody map[string]any
 	fallbackBodies := make([]map[string]any, 0, 3)
 	fallbackKeys := make([]string, 0, 3)
 	activeFallbacks := 0
@@ -188,13 +294,6 @@ func TestImageClientFallsBackToConcurrentSingleImageRequestsWhenCountUnsupported
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
-		if requestBody["n"] != nil {
-			mu.Lock()
-			initialBody = requestBody
-			mu.Unlock()
-			http.Error(w, `{"error":{"message":"Unknown parameter: 'tools[0].n'.","param":"tools[0].n","type":"invalid_request_error"}}`, http.StatusBadRequest)
-			return
-		}
 
 		mu.Lock()
 		fallbackBodies = append(fallbackBodies, requestBody)
@@ -229,22 +328,19 @@ func TestImageClientFallsBackToConcurrentSingleImageRequestsWhenCountUnsupported
 	}
 	mu.Lock()
 	defer mu.Unlock()
-	if initialBody["n"] != float64(3) {
-		t.Fatalf("expected first request n=3, got %#v", initialBody["n"])
-	}
 	if len(image.URLs) != 3 {
-		t.Fatalf("expected three fallback images, got %#v", image.URLs)
+		t.Fatalf("expected three generated images, got %#v", image.URLs)
 	}
 	if len(fallbackBodies) != 3 {
-		t.Fatalf("expected three single-image fallback requests, got %d", len(fallbackBodies))
+		t.Fatalf("expected three single-image requests, got %d", len(fallbackBodies))
 	}
 	if maxActiveFallbacks < 2 {
-		t.Fatalf("expected concurrent fallback requests, max active was %d", maxActiveFallbacks)
+		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 fallback request %d to omit n, got %#v", index, body)
+			t.Fatalf("expected split request %d to omit n, got %#v", index, body)
 		}
 		keySet[fallbackKeys[index]] = true
 	}