package service import ( "bytes" "context" "errors" "io" "mime/multipart" "net/http" "net/http/httptest" "net/textproto" "strings" "testing" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" ) type failingOpenAIImageWriter struct { gin.ResponseWriter failAfter int writes int } func (w *failingOpenAIImageWriter) Write(p []byte) (int, error) { if w.writes >= w.failAfter { return 0, errors.New("write failed: client disconnected") } w.writes++ return w.ResponseWriter.Write(p) } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_JSON(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","size":"1024x1024","quality":"high","stream":true}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) require.NotNil(t, parsed) require.Equal(t, "/v1/images/generations", parsed.Endpoint) require.Equal(t, "gpt-image-2", parsed.Model) require.Equal(t, "draw a cat", parsed.Prompt) require.True(t, parsed.Stream) require.Equal(t, "1024x1024", parsed.Size) require.Equal(t, "1K", parsed.SizeTier) require.Equal(t, OpenAIImagesCapabilityNative, parsed.RequiredCapability) require.False(t, parsed.Multipart) } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_MultipartEdit(t *testing.T) { gin.SetMode(gin.TestMode) var body bytes.Buffer writer := multipart.NewWriter(&body) require.NoError(t, writer.WriteField("model", "gpt-image-2")) require.NoError(t, writer.WriteField("prompt", "replace background")) require.NoError(t, writer.WriteField("size", "1536x1024")) part, err := writer.CreateFormFile("image", "source.png") require.NoError(t, err) _, err = part.Write([]byte("fake-image-bytes")) require.NoError(t, err) require.NoError(t, writer.Close()) req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes())) req.Header.Set("Content-Type", writer.FormDataContentType()) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body.Bytes()) require.NoError(t, err) require.NotNil(t, parsed) require.Equal(t, "/v1/images/edits", parsed.Endpoint) require.True(t, parsed.Multipart) require.Equal(t, "gpt-image-2", parsed.Model) require.Equal(t, "replace background", parsed.Prompt) require.Equal(t, "1536x1024", parsed.Size) require.Equal(t, "2K", parsed.SizeTier) require.Len(t, parsed.Uploads, 1) require.Equal(t, OpenAIImagesCapabilityNative, parsed.RequiredCapability) } func TestOpenAIImagesRequestModerationBody_JSONEditIncludesInputImageURLs(t *testing.T) { parsed := &OpenAIImagesRequest{ Endpoint: openAIImagesEditsEndpoint, Prompt: "replace background", InputImageURLs: []string{"https://example.com/source.png"}, MaskImageURL: "https://example.com/mask.png", } input := ExtractContentModerationInput(ContentModerationProtocolOpenAIImages, parsed.ModerationBody()) require.Equal(t, "replace background", input.Text) require.Equal(t, []string{"https://example.com/source.png", "https://example.com/mask.png"}, input.Images) } func TestOpenAIImagesRequestModerationBody_MultipartEditIncludesUploadsInMemory(t *testing.T) { parsed := &OpenAIImagesRequest{ Endpoint: openAIImagesEditsEndpoint, Prompt: "replace background", Uploads: []OpenAIImagesUpload{{ FieldName: "image", FileName: "source.png", ContentType: "image/png", Data: []byte("fake-image-bytes"), }}, MaskUpload: &OpenAIImagesUpload{ FieldName: "mask", FileName: "mask.png", ContentType: "image/png", Data: []byte("fake-mask-bytes"), }, } input := ExtractContentModerationInput(ContentModerationProtocolOpenAIImages, parsed.ModerationBody()) require.Equal(t, "replace background", input.Text) require.Equal(t, []string{ "data:image/png;base64,ZmFrZS1pbWFnZS1ieXRlcw==", "data:image/png;base64,ZmFrZS1tYXNrLWJ5dGVz", }, input.Images) log := (&ContentModerationService{}).buildLog(ContentModerationCheckInput{}, defaultContentModerationConfig(), ContentModerationActionAllow, false, "", 0, nil, input.ExcerptText(), nil, nil, "") require.Equal(t, "replace background", log.InputExcerpt) require.NotContains(t, log.InputExcerpt, "ZmFrZS") } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_NormalizesOfficialAndCustomSizes(t *testing.T) { gin.SetMode(gin.TestMode) tests := []struct { size string wantTier string }{ {size: "1024x1024", wantTier: "1K"}, {size: "1536x1024", wantTier: "2K"}, {size: "1024x1536", wantTier: "2K"}, {size: "2048x2048", wantTier: "2K"}, {size: "2048x1152", wantTier: "2K"}, {size: "3840x2160", wantTier: "4K"}, {size: "2160x3840", wantTier: "4K"}, {size: "1024X768", wantTier: "1K"}, {size: "1280x768", wantTier: "2K"}, {size: "2560x1440", wantTier: "4K"}, {size: "2560x1600", wantTier: "4K"}, {size: "auto", wantTier: "2K"}, } svc := &OpenAIGatewayService{} for _, tt := range tests { t.Run(tt.size, func(t *testing.T) { body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","size":"` + tt.size + `"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) require.NotNil(t, parsed) require.Equal(t, tt.size, parsed.Size) require.Equal(t, tt.wantTier, parsed.SizeTier) }) } } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_UnknownSizesDoNotBlockPassthrough(t *testing.T) { gin.SetMode(gin.TestMode) tests := []struct { size string wantTier string }{ {size: "2048x1153", wantTier: "2K"}, {size: "4096x1024", wantTier: "4K"}, {size: "3840x1024", wantTier: "4K"}, {size: "512x512", wantTier: "1K"}, {size: "invalid", wantTier: "2K"}, {size: "999999999999999999999999999x2", wantTier: "2K"}, } svc := &OpenAIGatewayService{} for _, tt := range tests { t.Run(tt.size, func(t *testing.T) { body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","size":"` + tt.size + `"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) require.NotNil(t, parsed) require.Equal(t, tt.size, parsed.Size) require.Equal(t, tt.wantTier, parsed.SizeTier) }) } } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_LegacyImageModelUnknownSizePassthrough(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-1.5","prompt":"draw a cat","size":"2048x1152"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) require.NotNil(t, parsed) require.Equal(t, "2048x1152", parsed.Size) require.Equal(t, "2K", parsed.SizeTier) } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_MultipartEditWithMaskAndNativeOptions(t *testing.T) { gin.SetMode(gin.TestMode) var body bytes.Buffer writer := multipart.NewWriter(&body) require.NoError(t, writer.WriteField("model", "gpt-image-2")) require.NoError(t, writer.WriteField("prompt", "replace foreground")) require.NoError(t, writer.WriteField("output_format", "png")) require.NoError(t, writer.WriteField("input_fidelity", "high")) require.NoError(t, writer.WriteField("output_compression", "80")) require.NoError(t, writer.WriteField("partial_images", "2")) imageHeader := make(textproto.MIMEHeader) imageHeader.Set("Content-Disposition", `form-data; name="image"; filename="source.png"`) imageHeader.Set("Content-Type", "image/png") imagePart, err := writer.CreatePart(imageHeader) require.NoError(t, err) _, err = imagePart.Write([]byte("source-image-bytes")) require.NoError(t, err) maskHeader := make(textproto.MIMEHeader) maskHeader.Set("Content-Disposition", `form-data; name="mask"; filename="mask.png"`) maskHeader.Set("Content-Type", "image/png") maskPart, err := writer.CreatePart(maskHeader) require.NoError(t, err) _, err = maskPart.Write([]byte("mask-image-bytes")) require.NoError(t, err) require.NoError(t, writer.Close()) req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes())) req.Header.Set("Content-Type", writer.FormDataContentType()) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body.Bytes()) require.NoError(t, err) require.NotNil(t, parsed) require.Len(t, parsed.Uploads, 1) require.NotNil(t, parsed.MaskUpload) require.True(t, parsed.HasMask) require.Equal(t, "png", parsed.OutputFormat) require.Equal(t, "high", parsed.InputFidelity) require.NotNil(t, parsed.OutputCompression) require.Equal(t, 80, *parsed.OutputCompression) require.NotNil(t, parsed.PartialImages) require.Equal(t, 2, *parsed.PartialImages) require.Equal(t, OpenAIImagesCapabilityNative, parsed.RequiredCapability) } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_PromptOnlyDefaultsRemainBasic(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"prompt":"draw a cat"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) require.NotNil(t, parsed) require.Equal(t, "gpt-image-2", parsed.Model) require.Equal(t, OpenAIImagesCapabilityBasic, parsed.RequiredCapability) } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_ExplicitSizeRequiresNativeCapability(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"prompt":"draw a cat","size":"1024x1024"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) require.NotNil(t, parsed) require.Equal(t, OpenAIImagesCapabilityNative, parsed.RequiredCapability) } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_RejectsNonImageModel(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-5.4","prompt":"draw a cat"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.Nil(t, parsed) require.ErrorContains(t, err, `images endpoint requires an image model, got "gpt-5.4"`) } func TestOpenAIGatewayServiceParseOpenAIImagesRequest_JSONEditURLs(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{ "model":"gpt-image-2", "prompt":"replace the background", "images":[{"image_url":"https://example.com/source.png"}], "mask":{"image_url":"https://example.com/mask.png"}, "input_fidelity":"high", "output_compression":90, "partial_images":2, "response_format":"url" }`) req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) require.NotNil(t, parsed) require.Equal(t, []string{"https://example.com/source.png"}, parsed.InputImageURLs) require.Equal(t, "https://example.com/mask.png", parsed.MaskImageURL) require.Equal(t, "high", parsed.InputFidelity) require.NotNil(t, parsed.OutputCompression) require.Equal(t, 90, *parsed.OutputCompression) require.NotNil(t, parsed.PartialImages) require.Equal(t, 2, *parsed.PartialImages) require.True(t, parsed.HasMask) require.Equal(t, OpenAIImagesCapabilityNative, parsed.RequiredCapability) } func TestCollectOpenAIImagePointers_RecognizesDirectAssets(t *testing.T) { items := collectOpenAIImagePointers([]byte(`{ "revised_prompt": "cat astronaut", "parts": [ {"b64_json":"QUJD"}, {"download_url":"https://files.example.com/image.png?sig=1"}, {"asset_pointer":"file-service://file_123"} ] }`)) require.Len(t, items, 3) var sawBase64, sawURL, sawPointer bool for _, item := range items { if item.B64JSON == "QUJD" { sawBase64 = true require.Equal(t, "cat astronaut", item.Prompt) } if item.DownloadURL == "https://files.example.com/image.png?sig=1" { sawURL = true } if item.Pointer == "file-service://file_123" { sawPointer = true } } require.True(t, sawBase64) require.True(t, sawURL) require.True(t, sawPointer) } func TestResolveOpenAIImageBytes_PrefersInlineBase64(t *testing.T) { data, err := resolveOpenAIImageBytes(context.Background(), nil, nil, "", openAIImagePointerInfo{ B64JSON: "data:image/png;base64,QUJD", }) require.NoError(t, err) require.Equal(t, []byte("ABC"), data) } func TestAccountSupportsOpenAIImageCapability_OAuthSupportsNative(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, Type: AccountTypeOAuth, } require.True(t, account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityBasic)) require.True(t, account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityNative)) } func TestAccountSupportsOpenAIEndpointCapability(t *testing.T) { t.Run("OpenAI APIKey 默认兼容 chat 和 embeddings", func(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, Type: AccountTypeAPIKey, } require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions)) require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings)) }) t.Run("OpenAI OAuth 默认仅兼容 chat", func(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, Type: AccountTypeOAuth, } require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions)) require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings)) }) t.Run("显式列表支持同时声明 chat 和 embeddings", func(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "openai_capabilities": []any{"chat_completions", "embeddings"}, }, } require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions)) require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings)) }) t.Run("显式列表只声明 chat 时不支持 embeddings", func(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "openai_capabilities": []any{"chat_completions"}, }, } require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions)) require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings)) }) t.Run("显式 map 支持单独关闭 chat 并开启 embeddings", func(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "openai_capabilities": map[string]any{ "chat_completions": false, "embeddings": true, }, }, } require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions)) require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings)) }) t.Run("未知能力不应默认放行", func(t *testing.T) { account := &Account{ Platform: PlatformOpenAI, Type: AccountTypeAPIKey, } require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapability("unknown"))) }) } func TestBuildOpenAIImagesURL_HandlesVersionedBaseURL(t *testing.T) { require.Equal(t, "https://image-upstream.example/v1/images/generations", buildOpenAIImagesURL("https://image-upstream.example/v1", openAIImagesGenerationsEndpoint), ) require.Equal(t, "https://open.bigmodel.cn/api/paas/v4/images/generations", buildOpenAIImagesURL("https://open.bigmodel.cn/api/paas/v4", openAIImagesGenerationsEndpoint), ) require.Equal(t, "https://image-upstream.example/v1/images/edits", buildOpenAIImagesURL("https://image-upstream.example/v1/", openAIImagesEditsEndpoint), ) require.Equal(t, "https://image-upstream.example/v1/images/generations", buildOpenAIImagesURL("https://image-upstream.example", openAIImagesGenerationsEndpoint), ) require.Equal(t, "https://image-upstream.example/v1/images/generations", buildOpenAIImagesURL("https://image-upstream.example/v1/images/generations", openAIImagesGenerationsEndpoint), ) } type openAIImageTestSSEEvent struct { Name string Data string } func parseOpenAIImageTestSSEEvents(body string) []openAIImageTestSSEEvent { chunks := strings.Split(body, "\n\n") events := make([]openAIImageTestSSEEvent, 0, len(chunks)) for _, chunk := range chunks { chunk = strings.TrimSpace(chunk) if chunk == "" { continue } var event openAIImageTestSSEEvent for _, line := range strings.Split(chunk, "\n") { switch { case strings.HasPrefix(line, "event: "): event.Name = strings.TrimSpace(strings.TrimPrefix(line, "event: ")) case strings.HasPrefix(line, "data: "): event.Data = strings.TrimSpace(strings.TrimPrefix(line, "data: ")) } } if event.Name != "" || event.Data != "" { events = append(events, event) } } return events } func findOpenAIImageTestSSEEvent(events []openAIImageTestSSEEvent, name string) (openAIImageTestSSEEvent, bool) { for _, event := range events { if event.Name == name { return event, true } } return openAIImageTestSSEEvent{}, false } func TestOpenAIGatewayServiceForwardImages_OAuthPassesNAndReturnsAllImages(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","size":"1024x1024","quality":"high","n":3}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req c.Set("api_key", &APIKey{ID: 42}) svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) upstream := &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_123"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000000,\"usage\":{\"input_tokens\":11,\"output_tokens\":22,\"input_tokens_details\":{\"cached_tokens\":3},\"output_tokens_details\":{\"image_tokens\":7}},\"tool_usage\":{\"image_gen\":{\"images\":3}},\"output\":[{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2UtMQ==\",\"revised_prompt\":\"draw a cat 1\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"},{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2UtMg==\",\"revised_prompt\":\"draw a cat 2\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"},{\"type\":\"image_generation_call\",\"result\":\"aW1hZ2UtMw==\",\"revised_prompt\":\"draw a cat 3\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"}]}}\n\n" + "data: [DONE]\n\n", )), }, } svc.httpUpstream = upstream account := &Account{ ID: 1, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", "chatgpt_account_id": "acct-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.Equal(t, "gpt-image-2", result.Model) require.Equal(t, "gpt-image-2", result.UpstreamModel) require.Equal(t, 3, result.ImageCount) require.Equal(t, 11, result.Usage.InputTokens) require.Equal(t, 22, result.Usage.OutputTokens) require.Equal(t, 7, result.Usage.ImageOutputTokens) require.NotNil(t, upstream.lastReq) require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String()) require.Equal(t, "chatgpt.com", upstream.lastReq.Host) require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.lastReq.Context())) require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type")) require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept")) require.Equal(t, "acct-123", upstream.lastReq.Header.Get("chatgpt-account-id")) require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta")) require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String()) require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) require.Equal(t, "image_generation", gjson.GetBytes(upstream.lastBody, "tools.0.type").String()) require.Equal(t, "generate", gjson.GetBytes(upstream.lastBody, "tools.0.action").String()) require.Equal(t, "gpt-image-2", gjson.GetBytes(upstream.lastBody, "tools.0.model").String()) require.Equal(t, "1024x1024", gjson.GetBytes(upstream.lastBody, "tools.0.size").String()) require.Equal(t, "high", gjson.GetBytes(upstream.lastBody, "tools.0.quality").String()) require.Equal(t, int64(3), gjson.GetBytes(upstream.lastBody, "tools.0.n").Int()) require.Equal(t, "draw a cat", gjson.GetBytes(upstream.lastBody, "input.0.content.0.text").String()) require.Equal(t, http.StatusOK, rec.Code) require.Equal(t, "gpt-image-2", gjson.Get(rec.Body.String(), "model").String()) require.Len(t, gjson.Get(rec.Body.String(), "data").Array(), 3) require.Equal(t, "aW1hZ2UtMQ==", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) require.Equal(t, "aW1hZ2UtMg==", gjson.Get(rec.Body.String(), "data.1.b64_json").String()) require.Equal(t, "aW1hZ2UtMw==", gjson.Get(rec.Body.String(), "data.2.b64_json").String()) require.Equal(t, "draw a cat 1", gjson.Get(rec.Body.String(), "data.0.revised_prompt").String()) require.Equal(t, "draw a cat 3", gjson.Get(rec.Body.String(), "data.2.revised_prompt").String()) } func TestOpenAIGatewayServiceForwardImages_OAuthNonStreamModerationBlockedReturnsClientError(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw blocked image","response_format":"b64_json"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req c.Set("api_key", &APIKey{ID: 42}) svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) svc.httpUpstream = &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_blocked"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000020}}\n\n" + "data: {\"type\":\"error\",\"error\":{\"type\":\"image_generation_user_error\",\"code\":\"moderation_blocked\",\"message\":\"Your request was rejected by the safety system. safety_violations=[sexual].\"}}\n\n" + "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_blocked\",\"status\":\"failed\",\"error\":{\"type\":\"image_generation_user_error\",\"code\":\"moderation_blocked\",\"message\":\"Your request was rejected by the safety system. safety_violations=[sexual].\"}}}\n\n", )), }, } account := &Account{ ID: 1, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.Nil(t, result) var upstreamErr *OpenAIImagesUpstreamError require.ErrorAs(t, err, &upstreamErr) require.Equal(t, http.StatusBadRequest, upstreamErr.StatusCode) require.Equal(t, "moderation_blocked", upstreamErr.Code) require.Equal(t, http.StatusBadRequest, rec.Code) require.Equal(t, "image_generation_user_error", gjson.Get(rec.Body.String(), "error.type").String()) require.Equal(t, "moderation_blocked", gjson.Get(rec.Body.String(), "error.code").String()) require.Contains(t, gjson.Get(rec.Body.String(), "error.message").String(), "safety system") } func TestOpenAIGatewayServiceForwardImages_APIKeyGenerationUsesConfiguredV1BaseURL(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","response_format":"b64_json"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{ cfg: &config.Config{}, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"application/json"}, "X-Request-Id": []string{"req_img_apikey"}, }, Body: io.NopCloser(strings.NewReader(`{"created":1710000007,"data":[{"b64_json":"aGVsbG8=","revised_prompt":"draw a cat"}]}`)), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) account := &Account{ ID: 6, Name: "openai-apikey", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "api_key": "test-api-key", "base_url": "https://image-upstream.example/v1", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.Equal(t, 1, result.ImageCount) require.Equal(t, "gpt-image-2", result.Model) require.Equal(t, "gpt-image-2", result.UpstreamModel) upstream, ok := svc.httpUpstream.(*httpUpstreamRecorder) require.True(t, ok) require.NotNil(t, upstream.lastReq) require.Equal(t, "https://image-upstream.example/v1/images/generations", upstream.lastReq.URL.String()) require.Equal(t, "Bearer test-api-key", upstream.lastReq.Header.Get("Authorization")) require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type")) require.Equal(t, "gpt-image-2", gjson.GetBytes(upstream.lastBody, "model").String()) require.Equal(t, http.StatusOK, rec.Code) require.Equal(t, "aGVsbG8=", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) } func TestOpenAIGatewayServiceForwardImages_APIKeyStreamJSONResponseBillsImage(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"response_format":"b64_json"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{ cfg: &config.Config{}, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"application/json"}, "X-Request-Id": []string{"req_img_stream_json"}, }, Body: io.NopCloser(strings.NewReader(`{"created":1710000008,"usage":{"input_tokens":12,"output_tokens":21,"output_tokens_details":{"image_tokens":9}},"data":[{"b64_json":"aGVsbG8=","revised_prompt":"draw a cat"}]}`)), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) account := &Account{ ID: 7, Name: "openai-apikey", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "api_key": "test-api-key", "base_url": "https://image-upstream.example/v1", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.True(t, result.Stream) require.Equal(t, 1, result.ImageCount) require.Equal(t, 12, result.Usage.InputTokens) require.Equal(t, 21, result.Usage.OutputTokens) require.Equal(t, 9, result.Usage.ImageOutputTokens) require.Equal(t, http.StatusOK, rec.Code) require.Equal(t, "aGVsbG8=", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) } func TestOpenAIGatewayServiceForwardImages_APIKeyStreamRawJSONEventStreamFallbackBillsImage(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"response_format":"b64_json"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{ cfg: &config.Config{}, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_stream_json_mislabeled"}, }, Body: io.NopCloser(strings.NewReader(`{"created":1710000009,"usage":{"input_tokens":10,"output_tokens":18,"output_tokens_details":{"image_tokens":8}},"data":[{"b64_json":"ZmluYWw="}]}`)), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) account := &Account{ ID: 8, Name: "openai-apikey", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "api_key": "test-api-key", "base_url": "https://image-upstream.example/v1", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.True(t, result.Stream) require.Equal(t, 1, result.ImageCount) require.Equal(t, 10, result.Usage.InputTokens) require.Equal(t, 18, result.Usage.OutputTokens) require.Equal(t, 8, result.Usage.ImageOutputTokens) require.Equal(t, "ZmluYWw=", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) } func TestOpenAIGatewayServiceForwardImages_APIKeyStreamMultilineSSEDataBillsImage(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"response_format":"b64_json"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{ cfg: &config.Config{}, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_stream_multiline"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"image_generation.completed\",\n" + "data: \"usage\":{\"input_tokens\":10,\"output_tokens\":18,\"output_tokens_details\":{\"image_tokens\":8}},\n" + "data: \"b64_json\":\"ZmluYWw=\",\"output_format\":\"png\"}\n\n" + "data: [DONE]\n\n", )), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) account := &Account{ ID: 8, Name: "openai-apikey", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "api_key": "test-api-key", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.True(t, result.Stream) require.Equal(t, 1, result.ImageCount) require.Equal(t, 10, result.Usage.InputTokens) require.Equal(t, 18, result.Usage.OutputTokens) require.Equal(t, 8, result.Usage.ImageOutputTokens) } func TestExtractOpenAIImagesBillableCountFromJSONBytes_CompletedEvent(t *testing.T) { body := []byte(`{"type":"image_generation.completed","b64_json":"ZmluYWw=","usage":{"input_tokens":10,"output_tokens":18}}`) require.Equal(t, 1, extractOpenAIImagesBillableCountFromJSONBytes(body)) } func TestOpenAIGatewayServiceForwardImages_APIKeyEditUsesConfiguredV1BaseURL(t *testing.T) { gin.SetMode(gin.TestMode) var body bytes.Buffer writer := multipart.NewWriter(&body) require.NoError(t, writer.WriteField("model", "gpt-image-2")) require.NoError(t, writer.WriteField("prompt", "replace background")) imagePart, err := writer.CreateFormFile("image", "source.png") require.NoError(t, err) _, err = imagePart.Write([]byte("png-image-content")) require.NoError(t, err) require.NoError(t, writer.Close()) req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes())) req.Header.Set("Content-Type", writer.FormDataContentType()) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{ cfg: &config.Config{}, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"application/json"}, "X-Request-Id": []string{"req_img_edit_apikey"}, }, Body: io.NopCloser(strings.NewReader(`{"created":1710000008,"data":[{"b64_json":"ZWRpdGVk","revised_prompt":"replace background"}]}`)), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body.Bytes()) require.NoError(t, err) account := &Account{ ID: 7, Name: "openai-apikey", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "api_key": "test-api-key", "base_url": "https://image-upstream.example/v1/", }, } result, err := svc.ForwardImages(context.Background(), c, account, body.Bytes(), parsed, "") require.NoError(t, err) require.NotNil(t, result) require.Equal(t, 1, result.ImageCount) upstream, ok := svc.httpUpstream.(*httpUpstreamRecorder) require.True(t, ok) require.NotNil(t, upstream.lastReq) require.Equal(t, "https://image-upstream.example/v1/images/edits", upstream.lastReq.URL.String()) require.Equal(t, "Bearer test-api-key", upstream.lastReq.Header.Get("Authorization")) require.Contains(t, upstream.lastReq.Header.Get("Content-Type"), "multipart/form-data") require.Contains(t, string(upstream.lastBody), `name="model"`) require.Contains(t, string(upstream.lastBody), "gpt-image-2") require.Equal(t, http.StatusOK, rec.Code) require.Equal(t, "ZWRpdGVk", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) } func TestOpenAIGatewayServiceForwardImages_OAuthStreamingTransformsEvents(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"response_format":"url"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) upstream := &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_stream"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000001,\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"background\":\"auto\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"}]}}\n\n" + "data: {\"type\":\"response.image_generation_call.partial_image\",\"partial_image_b64\":\"cGFydGlhbA==\",\"partial_image_index\":0,\"output_format\":\"png\",\"background\":\"auto\"}\n\n" + "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000001,\"usage\":{\"input_tokens\":5,\"output_tokens\":9,\"output_tokens_details\":{\"image_tokens\":4}},\"tool_usage\":{\"image_gen\":{\"images\":1}},\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"background\":\"auto\",\"output_format\":\"png\",\"quality\":\"high\",\"size\":\"1024x1024\"}],\"output\":[{\"type\":\"image_generation_call\",\"result\":\"ZmluYWw=\",\"output_format\":\"png\"}]}}\n\n" + "data: [DONE]\n\n", )), }, } svc.httpUpstream = upstream account := &Account{ ID: 2, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.True(t, result.Stream) require.Equal(t, 1, result.ImageCount) events := parseOpenAIImageTestSSEEvents(rec.Body.String()) partial, ok := findOpenAIImageTestSSEEvent(events, "image_generation.partial_image") require.True(t, ok) require.Equal(t, "image_generation.partial_image", gjson.Get(partial.Data, "type").String()) require.Equal(t, int64(1710000001), gjson.Get(partial.Data, "created_at").Int()) require.Equal(t, "cGFydGlhbA==", gjson.Get(partial.Data, "b64_json").String()) require.Equal(t, "data:image/png;base64,cGFydGlhbA==", gjson.Get(partial.Data, "url").String()) require.Equal(t, "gpt-image-2", gjson.Get(partial.Data, "model").String()) require.Equal(t, "png", gjson.Get(partial.Data, "output_format").String()) require.Equal(t, "high", gjson.Get(partial.Data, "quality").String()) require.Equal(t, "1024x1024", gjson.Get(partial.Data, "size").String()) require.Equal(t, "auto", gjson.Get(partial.Data, "background").String()) completed, ok := findOpenAIImageTestSSEEvent(events, "image_generation.completed") require.True(t, ok) require.Equal(t, "image_generation.completed", gjson.Get(completed.Data, "type").String()) require.Equal(t, int64(1710000001), gjson.Get(completed.Data, "created_at").Int()) require.Equal(t, "ZmluYWw=", gjson.Get(completed.Data, "b64_json").String()) require.Equal(t, "data:image/png;base64,ZmluYWw=", gjson.Get(completed.Data, "url").String()) require.Equal(t, "gpt-image-2", gjson.Get(completed.Data, "model").String()) require.Equal(t, "png", gjson.Get(completed.Data, "output_format").String()) require.Equal(t, "high", gjson.Get(completed.Data, "quality").String()) require.Equal(t, "1024x1024", gjson.Get(completed.Data, "size").String()) require.Equal(t, "auto", gjson.Get(completed.Data, "background").String()) require.JSONEq(t, `{"images":1}`, gjson.Get(completed.Data, "usage").Raw) require.False(t, gjson.Get(completed.Data, "revised_prompt").Exists()) } func TestOpenAIGatewayServiceForwardImages_APIKeyStreamingDrainsAfterClientDisconnect(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req c.Writer = &failingOpenAIImageWriter{ResponseWriter: c.Writer, failAfter: 1} svc := &OpenAIGatewayService{ cfg: &config.Config{ Gateway: config.GatewayConfig{ ImageStreamDataIntervalTimeout: 1, ImageStreamKeepaliveInterval: 0, }, }, httpUpstream: &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_stream_disconnect_apikey"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"image_generation.partial_image\",\"b64_json\":\"cGFydGlhbA==\"}\n\n" + "data: {\"type\":\"image_generation.completed\",\"usage\":{\"input_tokens\":3,\"output_tokens\":4,\"output_tokens_details\":{\"image_tokens\":2}},\"b64_json\":\"ZmluYWw=\",\"output_format\":\"png\"}\n\n" + "data: [DONE]\n\n", )), }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) account := &Account{ ID: 8, Name: "openai-apikey", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "api_key": "test-api-key", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.Equal(t, 1, result.ImageCount) require.Equal(t, 3, result.Usage.InputTokens) require.Equal(t, 4, result.Usage.OutputTokens) require.Equal(t, 2, result.Usage.ImageOutputTokens) } func TestOpenAIGatewayServiceForwardImages_OAuthEditsMultipartUsesResponsesAPI(t *testing.T) { gin.SetMode(gin.TestMode) var body bytes.Buffer writer := multipart.NewWriter(&body) require.NoError(t, writer.WriteField("model", "gpt-image-2")) require.NoError(t, writer.WriteField("prompt", "replace background with aurora")) require.NoError(t, writer.WriteField("input_fidelity", "high")) require.NoError(t, writer.WriteField("output_format", "webp")) require.NoError(t, writer.WriteField("quality", "high")) imageHeader := make(textproto.MIMEHeader) imageHeader.Set("Content-Disposition", `form-data; name="image"; filename="source.png"`) imageHeader.Set("Content-Type", "image/png") imagePart, err := writer.CreatePart(imageHeader) require.NoError(t, err) _, err = imagePart.Write([]byte("png-image-content")) require.NoError(t, err) maskHeader := make(textproto.MIMEHeader) maskHeader.Set("Content-Disposition", `form-data; name="mask"; filename="mask.png"`) maskHeader.Set("Content-Type", "image/png") maskPart, err := writer.CreatePart(maskHeader) require.NoError(t, err) _, err = maskPart.Write([]byte("png-mask-content")) require.NoError(t, err) require.NoError(t, writer.Close()) req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes())) req.Header.Set("Content-Type", writer.FormDataContentType()) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req c.Set("api_key", &APIKey{ID: 100}) svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body.Bytes()) require.NoError(t, err) upstream := &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_edit_123"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000002,\"usage\":{\"input_tokens\":13,\"output_tokens\":21,\"output_tokens_details\":{\"image_tokens\":8}},\"tool_usage\":{\"image_gen\":{\"images\":1}},\"output\":[{\"type\":\"image_generation_call\",\"result\":\"ZWRpdGVk\",\"revised_prompt\":\"replace background with aurora\",\"output_format\":\"webp\",\"quality\":\"high\"}]}}\n\n" + "data: [DONE]\n\n", )), }, } svc.httpUpstream = upstream account := &Account{ ID: 3, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body.Bytes(), parsed, "") require.NoError(t, err) require.NotNil(t, result) require.Equal(t, 1, result.ImageCount) require.Equal(t, "gpt-image-2", gjson.GetBytes(upstream.lastBody, "tools.0.model").String()) require.Equal(t, "edit", gjson.GetBytes(upstream.lastBody, "tools.0.action").String()) require.False(t, gjson.GetBytes(upstream.lastBody, "tools.0.input_fidelity").Exists()) require.Equal(t, "webp", gjson.GetBytes(upstream.lastBody, "tools.0.output_format").String()) require.True(t, strings.HasPrefix(gjson.GetBytes(upstream.lastBody, "input.0.content.1.image_url").String(), "data:image/png;base64,")) require.True(t, strings.HasPrefix(gjson.GetBytes(upstream.lastBody, "tools.0.input_image_mask.image_url").String(), "data:image/png;base64,")) require.Equal(t, "replace background with aurora", gjson.GetBytes(upstream.lastBody, "input.0.content.0.text").String()) require.Equal(t, "ZWRpdGVk", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) require.Equal(t, "replace background with aurora", gjson.Get(rec.Body.String(), "data.0.revised_prompt").String()) } func TestOpenAIGatewayServiceForwardImages_OAuthEditsStreamingTransformsEvents(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{ "model":"gpt-image-2", "prompt":"replace background with aurora", "images":[{"image_url":"https://example.com/source.png"}], "mask":{"image_url":"https://example.com/mask.png"}, "stream":true, "response_format":"url" }`) req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) upstream := &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000003,\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"background\":\"transparent\",\"output_format\":\"webp\",\"quality\":\"high\",\"size\":\"1024x1024\"}]}}\n\n" + "data: {\"type\":\"response.image_generation_call.partial_image\",\"partial_image_b64\":\"cGFydGlhbA==\",\"partial_image_index\":0,\"output_format\":\"webp\",\"background\":\"transparent\"}\n\n" + "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000003,\"usage\":{\"input_tokens\":7,\"output_tokens\":10,\"output_tokens_details\":{\"image_tokens\":5}},\"tool_usage\":{\"image_gen\":{\"images\":1}},\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"background\":\"transparent\",\"output_format\":\"webp\",\"quality\":\"high\",\"size\":\"1024x1024\"}],\"output\":[{\"type\":\"image_generation_call\",\"result\":\"ZWRpdGVk\",\"revised_prompt\":\"replace background with aurora\",\"output_format\":\"webp\"}]}}\n\n" + "data: [DONE]\n\n", )), }, } svc.httpUpstream = upstream account := &Account{ ID: 4, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.Equal(t, 1, result.ImageCount) require.Equal(t, "edit", gjson.GetBytes(upstream.lastBody, "tools.0.action").String()) require.Equal(t, "https://example.com/source.png", gjson.GetBytes(upstream.lastBody, "input.0.content.1.image_url").String()) require.Equal(t, "https://example.com/mask.png", gjson.GetBytes(upstream.lastBody, "tools.0.input_image_mask.image_url").String()) events := parseOpenAIImageTestSSEEvents(rec.Body.String()) partial, ok := findOpenAIImageTestSSEEvent(events, "image_edit.partial_image") require.True(t, ok) require.Equal(t, "image_edit.partial_image", gjson.Get(partial.Data, "type").String()) require.Equal(t, int64(1710000003), gjson.Get(partial.Data, "created_at").Int()) require.Equal(t, "cGFydGlhbA==", gjson.Get(partial.Data, "b64_json").String()) require.Equal(t, "data:image/webp;base64,cGFydGlhbA==", gjson.Get(partial.Data, "url").String()) require.Equal(t, "gpt-image-2", gjson.Get(partial.Data, "model").String()) require.Equal(t, "webp", gjson.Get(partial.Data, "output_format").String()) require.Equal(t, "high", gjson.Get(partial.Data, "quality").String()) require.Equal(t, "1024x1024", gjson.Get(partial.Data, "size").String()) require.Equal(t, "transparent", gjson.Get(partial.Data, "background").String()) completed, ok := findOpenAIImageTestSSEEvent(events, "image_edit.completed") require.True(t, ok) require.Equal(t, "image_edit.completed", gjson.Get(completed.Data, "type").String()) require.Equal(t, int64(1710000003), gjson.Get(completed.Data, "created_at").Int()) require.Equal(t, "ZWRpdGVk", gjson.Get(completed.Data, "b64_json").String()) require.Equal(t, "data:image/webp;base64,ZWRpdGVk", gjson.Get(completed.Data, "url").String()) require.Equal(t, "gpt-image-2", gjson.Get(completed.Data, "model").String()) require.Equal(t, "webp", gjson.Get(completed.Data, "output_format").String()) require.Equal(t, "high", gjson.Get(completed.Data, "quality").String()) require.Equal(t, "1024x1024", gjson.Get(completed.Data, "size").String()) require.Equal(t, "transparent", gjson.Get(completed.Data, "background").String()) require.JSONEq(t, `{"images":1}`, gjson.Get(completed.Data, "usage").Raw) require.False(t, gjson.Get(completed.Data, "revised_prompt").Exists()) } func TestBuildOpenAIImagesResponsesRequest_PassesThroughNForMultiImageModels(t *testing.T) { parsed := &OpenAIImagesRequest{ Endpoint: openAIImagesGenerationsEndpoint, Model: "gpt-image-2", Prompt: "draw a cat", N: 2, } body, err := buildOpenAIImagesResponsesRequest(parsed, "gpt-image-2") require.NoError(t, err) require.NotNil(t, body) require.Equal(t, int64(2), gjson.GetBytes(body, "tools.0.n").Int()) require.Equal(t, "gpt-image-2", gjson.GetBytes(body, "tools.0.model").String()) require.Equal(t, "draw a cat", gjson.GetBytes(body, "input.0.content.0.text").String()) } func TestBuildOpenAIImagesResponsesRequest_DoesNotPassNForDallE3(t *testing.T) { parsed := &OpenAIImagesRequest{ Endpoint: openAIImagesGenerationsEndpoint, Model: "dall-e-3", Prompt: "draw a cat", N: 2, } body, err := buildOpenAIImagesResponsesRequest(parsed, "dall-e-3") require.NoError(t, err) require.NotNil(t, body) require.False(t, gjson.GetBytes(body, "tools.0.n").Exists()) require.Equal(t, "dall-e-3", gjson.GetBytes(body, "tools.0.model").String()) } func TestBuildOpenAIImagesResponsesRequest_StripsInputFidelity(t *testing.T) { parsed := &OpenAIImagesRequest{ Endpoint: openAIImagesEditsEndpoint, Model: "gpt-image-2", Prompt: "replace background", InputFidelity: "high", InputImageURLs: []string{ "https://example.com/source.png", }, } body, err := buildOpenAIImagesResponsesRequest(parsed, "gpt-image-2") require.NoError(t, err) require.NotNil(t, body) require.False(t, gjson.GetBytes(body, "tools.0.input_fidelity").Exists()) require.Equal(t, "edit", gjson.GetBytes(body, "tools.0.action").String()) } func TestCollectOpenAIImagesFromResponsesBody_FallsBackToOutputItemDone(t *testing.T) { body := []byte( "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000004}}\n\n" + "data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"result\":\"aGVsbG8=\",\"revised_prompt\":\"draw a cat\",\"output_format\":\"png\",\"quality\":\"high\"}}\n\n" + "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000004,\"tool_usage\":{\"image_gen\":{\"images\":1}},\"output\":[]}}\n\n" + "data: [DONE]\n\n", ) results, createdAt, usageRaw, firstMeta, foundFinal, err := collectOpenAIImagesFromResponsesBody(body) require.NoError(t, err) require.True(t, foundFinal) require.Equal(t, int64(1710000004), createdAt) require.Len(t, results, 1) require.Equal(t, "aGVsbG8=", results[0].Result) require.Equal(t, "draw a cat", results[0].RevisedPrompt) require.Equal(t, "png", firstMeta.OutputFormat) require.JSONEq(t, `{"images":1}`, string(usageRaw)) } func TestCollectOpenAIImagesFromResponsesBody_MultilineSSE(t *testing.T) { body := []byte( "data: {\"type\":\"response.completed\",\n" + "data: \"response\":{\"created_at\":1710000010,\"usage\":{\"input_tokens\":5,\"output_tokens\":9,\"output_tokens_details\":{\"image_tokens\":4}},\"tool_usage\":{\"image_gen\":{\"images\":1}},\"output\":[{\"type\":\"image_generation_call\",\"result\":\"ZmluYWw=\",\"output_format\":\"png\"}]}}\n\n" + "data: [DONE]\n\n", ) results, createdAt, usageRaw, firstMeta, foundFinal, err := collectOpenAIImagesFromResponsesBody(body) require.NoError(t, err) require.True(t, foundFinal) require.Equal(t, int64(1710000010), createdAt) require.Len(t, results, 1) require.Equal(t, "ZmluYWw=", results[0].Result) require.Equal(t, "png", firstMeta.OutputFormat) require.JSONEq(t, `{"images":1}`, string(usageRaw)) } func TestOpenAIGatewayServiceForwardImages_OAuthStreamingHandlesOutputItemDoneFallback(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"response_format":"url"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) upstream := &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_stream_output_item_done"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"result\":\"ZmluYWw=\",\"revised_prompt\":\"draw a cat\",\"output_format\":\"png\"}}\n\n" + "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000005,\"usage\":{\"input_tokens\":5,\"output_tokens\":9,\"output_tokens_details\":{\"image_tokens\":4}},\"tool_usage\":{\"image_gen\":{\"images\":1}},\"output\":[]}}\n\n" + "data: [DONE]\n\n", )), }, } svc.httpUpstream = upstream account := &Account{ ID: 5, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.True(t, result.Stream) require.Equal(t, 1, result.ImageCount) events := parseOpenAIImageTestSSEEvents(rec.Body.String()) completed, ok := findOpenAIImageTestSSEEvent(events, "image_generation.completed") require.True(t, ok) require.Equal(t, "image_generation.completed", gjson.Get(completed.Data, "type").String()) require.Equal(t, int64(1710000005), gjson.Get(completed.Data, "created_at").Int()) require.Equal(t, "ZmluYWw=", gjson.Get(completed.Data, "b64_json").String()) require.Equal(t, "data:image/png;base64,ZmluYWw=", gjson.Get(completed.Data, "url").String()) require.Equal(t, "gpt-image-2", gjson.Get(completed.Data, "model").String()) require.JSONEq(t, `{"images":1}`, gjson.Get(completed.Data, "usage").Raw) require.NotContains(t, rec.Body.String(), "event: error") } func TestOpenAIGatewayServiceForwardImages_OAuthStreamingHandlesMultilineSSE(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"response_format":"b64_json"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req svc := &OpenAIGatewayService{} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) svc.httpUpstream = &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_stream_multiline_oauth"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.completed\",\n" + "data: \"response\":{\"created_at\":1710000011,\"usage\":{\"input_tokens\":6,\"output_tokens\":10,\"output_tokens_details\":{\"image_tokens\":5}},\"tool_usage\":{\"image_gen\":{\"images\":1}},\"output\":[{\"type\":\"image_generation_call\",\"result\":\"TXVsdGlsaW5l\",\"output_format\":\"png\"}]}}\n\n" + "data: [DONE]\n\n", )), }, } account := &Account{ ID: 11, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.True(t, result.Stream) require.Equal(t, 1, result.ImageCount) require.Equal(t, 6, result.Usage.InputTokens) require.Equal(t, 10, result.Usage.OutputTokens) require.Equal(t, 5, result.Usage.ImageOutputTokens) events := parseOpenAIImageTestSSEEvents(rec.Body.String()) completed, ok := findOpenAIImageTestSSEEvent(events, "image_generation.completed") require.True(t, ok) require.Equal(t, "TXVsdGlsaW5l", gjson.Get(completed.Data, "b64_json").String()) require.JSONEq(t, `{"images":1}`, gjson.Get(completed.Data, "usage").Raw) require.NotContains(t, rec.Body.String(), "event: error") } func TestOpenAIGatewayServiceForwardImages_OAuthStreamingDrainsAfterClientDisconnect(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"response_format":"url"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = req c.Writer = &failingOpenAIImageWriter{ResponseWriter: c.Writer, failAfter: 1} svc := &OpenAIGatewayService{ cfg: &config.Config{ Gateway: config.GatewayConfig{ ImageStreamDataIntervalTimeout: 1, ImageStreamKeepaliveInterval: 0, }, }, } parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) upstream := &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, Header: http.Header{ "Content-Type": []string{"text/event-stream"}, "X-Request-Id": []string{"req_img_stream_disconnect_oauth"}, }, Body: io.NopCloser(strings.NewReader( "data: {\"type\":\"response.image_generation_call.partial_image\",\"partial_image_b64\":\"cGFydGlhbA==\",\"partial_image_index\":0,\"output_format\":\"png\"}\n\n" + "data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000009,\"usage\":{\"input_tokens\":5,\"output_tokens\":9,\"output_tokens_details\":{\"image_tokens\":4}},\"tool_usage\":{\"image_gen\":{\"images\":1}},\"output\":[{\"type\":\"image_generation_call\",\"result\":\"ZmluYWw=\",\"output_format\":\"png\"}]}}\n\n" + "data: [DONE]\n\n", )), }, } svc.httpUpstream = upstream account := &Account{ ID: 9, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ "access_token": "token-123", }, } result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") require.NoError(t, err) require.NotNil(t, result) require.True(t, result.Stream) require.Equal(t, 1, result.ImageCount) require.Equal(t, 5, result.Usage.InputTokens) require.Equal(t, 9, result.Usage.OutputTokens) require.Equal(t, 4, result.Usage.ImageOutputTokens) }