feat(piagent): fall back to concurrent single-image requests when count unsupported

Some image endpoints reject the `n` count parameter (Unknown parameter
tools[0].n / param: n). Detect that error and re-issue the batch as
concurrent count=1 requests, each with a suffixed idempotency key, then
assemble the URLs into a single GeneratedImage.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-07-09 12:26:17 +08:00
parent 568690e2f1
commit 299bac6c1a
2 changed files with 175 additions and 0 deletions
@@ -3,10 +3,12 @@ package piagent
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
)
@@ -170,6 +172,89 @@ func TestImageClientSendsImageCount(t *testing.T) {
}
}
func TestImageClientFallsBackToConcurrentSingleImageRequestsWhenCountUnsupported(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
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
}
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)
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 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)
}
if len(fallbackBodies) != 3 {
t.Fatalf("expected three single-image fallback requests, got %d", len(fallbackBodies))
}
if maxActiveFallbacks < 2 {
t.Fatalf("expected concurrent fallback 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)
}
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