| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231 |
- package service
- import (
- "bytes"
- "context"
- "errors"
- "fmt"
- "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/imroc/req/v3"
- "github.com/stretchr/testify/require"
- "github.com/tidwall/gjson"
- )
- type failingOpenAIImageWriter struct {
- gin.ResponseWriter
- failAfter int
- writes int
- }
- type openAIImagesReadErrorBody struct {
- err error
- }
- func (b *openAIImagesReadErrorBody) Read([]byte) (int, error) { return 0, b.err }
- func (b *openAIImagesReadErrorBody) Close() error { return nil }
- 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_AllowsGrokImageModels(t *testing.T) {
- gin.SetMode(gin.TestMode)
- for _, model := range []string{"grok-imagine", "grok-imagine-image", "grok-imagine-image-quality", "grok-imagine-edit"} {
- t.Run(model, func(t *testing.T) {
- body := []byte(fmt.Sprintf(`{"model":%q,"prompt":"draw a cat","response_format":"b64_json"}`, model))
- 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, model, parsed.Model)
- require.Equal(t, OpenAIImagesCapabilityNative, parsed.RequiredCapability)
- })
- }
- }
- 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",
- }, openAIUpstreamErrorBodyReadLimit)
- require.NoError(t, err)
- require.Equal(t, []byte("ABC"), data)
- }
- func TestNewOpenAIImageStatusError_UsesProvidedReadLimit(t *testing.T) {
- padding := strings.Repeat("x", int(openAIUpstreamErrorBodyReadLimit)+1024)
- body := fmt.Sprintf(`{"error":{"padding":"%s","message":"diagnostic-marker"}}`, padding)
- resp := &req.Response{Response: &http.Response{
- StatusCode: http.StatusBadGateway,
- Header: http.Header{},
- Body: io.NopCloser(strings.NewReader(body)),
- }}
- err := newOpenAIImageStatusError(resp, "download image bytes failed", int64(len(body)))
- require.Error(t, err)
- require.Equal(t, "diagnostic-marker", err.Error())
- var statusErr *openAIImageStatusError
- require.ErrorAs(t, err, &statusErr)
- require.Len(t, statusErr.ResponseBody, len(body))
- }
- func TestOpenAIUpstreamErrorBodyReadLimitForConfig_RespectsDiagnosticLimit(t *testing.T) {
- cfg := &config.Config{Gateway: config.GatewayConfig{
- LogUpstreamErrorBody: true,
- LogUpstreamErrorBodyMaxBytes: int(openAIUpstreamErrorBodyReadLimit) + 1024,
- }}
- require.Equal(t, int64(cfg.Gateway.LogUpstreamErrorBodyMaxBytes), openAIUpstreamErrorBodyReadLimitForConfig(cfg))
- }
- 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 TestAccountSupportsOpenAIImageCapability_SetupTokenSupportsNative(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeSetupToken,
- }
- require.True(t, account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityBasic))
- require.True(t, account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityNative))
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings))
- }
- func TestAccountSupportsOpenAIImageCapability_EmptyRequirementDoesNotRejectGrok(t *testing.T) {
- account := &Account{
- Platform: PlatformGrok,
- Type: AccountTypeOAuth,
- }
- require.True(t, account.SupportsOpenAIImageCapability(""))
- require.False(t, account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityBasic))
- }
- func TestAccountSupportsOpenAIEndpointCapability(t *testing.T) {
- t.Run("OpenAI APIKey 默认兼容 chat、embeddings 和 alpha search", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings))
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
- })
- t.Run("OpenAI OAuth 默认仅兼容 chat", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings))
- })
- t.Run("alpha search 允许 OpenAI OAuth/PAT 与 APIKey 账号,拒绝 Grok", func(t *testing.T) {
- // OAuth/PAT 走 chatgpt.com Codex 端点,APIKey 走 {base_url}/v1/alpha/search,
- // 两类都能承接独立搜索(APIKey 被排除曾导致纯 APIKey 分组搜索失效的回归)。
- apiKey := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- }
- oauth := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- }
- grok := &Account{
- Platform: PlatformGrok,
- Type: AccountTypeAPIKey,
- }
- require.True(t, apiKey.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
- require.True(t, oauth.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
- require.False(t, grok.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
- })
- 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))
- // chat 能力隐含放行 alpha search(OAuth/APIKey 语义一致)。
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings))
- })
- t.Run("OAuth 显式列表沿用 chat 能力放行 alpha search", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Credentials: map[string]any{
- "openai_capabilities": []any{"chat_completions"},
- },
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
- })
- 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("空 openai_capabilities({})与未配置一致,不排除 OAuth 文本调度", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Credentials: map[string]any{
- "openai_capabilities": map[string]any{},
- },
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityResponses))
- })
- t.Run("空 openai_capabilities([]any)与未配置一致,不排除 OAuth 文本调度", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Credentials: map[string]any{
- "openai_capabilities": []any{},
- },
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
- })
- t.Run("空 openai_capabilities([]string)与未配置一致,不排除 OAuth 文本调度", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Credentials: map[string]any{
- "openai_capabilities": []string{},
- },
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
- })
- t.Run("非空但全 false 的 map 仍按显式禁用处理,不默认放行", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Credentials: map[string]any{
- "openai_capabilities": map[string]any{"chat_completions": false},
- },
- }
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
- })
- t.Run("类型异常(字符串)仍视为已配置但不含能力,不默认放行", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Credentials: map[string]any{
- "openai_capabilities": "chat_completions",
- },
- }
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
- })
- t.Run("未知能力不应默认放行", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- }
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapability("unknown")))
- })
- t.Run("responses 能力:未探测的 APIKey 默认放行", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityResponses))
- })
- t.Run("responses 能力:探测确认不支持的 APIKey 被排除", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Extra: map[string]any{"openai_responses_supported": false},
- }
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityResponses))
- // 非生图路径仍可选中(只要求 chat_completions)。
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
- })
- t.Run("responses 能力:探测确认支持的 APIKey 放行", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Extra: map[string]any{"openai_responses_supported": true},
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityResponses))
- })
- t.Run("responses 能力:force_chat_completions 覆盖排除 APIKey", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Extra: map[string]any{"openai_responses_mode": "force_chat_completions"},
- }
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityResponses))
- })
- t.Run("responses 能力:OAuth 账号不受探测标记影响", func(t *testing.T) {
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Extra: map[string]any{"openai_responses_supported": false},
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityResponses))
- })
- t.Run("responses 能力:仍需通过 chat_completions 配置集校验", func(t *testing.T) {
- // 未探测(默认支持 responses),但显式能力集未声明 chat_completions。
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Credentials: map[string]any{
- "openai_capabilities": []any{"embeddings"},
- },
- }
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityResponses))
- })
- }
- 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\":{\"input_tokens\":46,\"output_tokens\":2459,\"output_tokens_details\":{\"image_tokens\":2459},\"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, 46, result.Usage.InputTokens)
- require.Equal(t, 2459, result.Usage.OutputTokens)
- require.Equal(t, 2459, 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 TestParseOpenAIImagesSSEUsageBytes_ToolUsagePrecedenceAndFallback(t *testing.T) {
- svc := &OpenAIGatewayService{}
- fallback := OpenAIUsage{InputTokens: 3, OutputTokens: 4, ImageOutputTokens: 2}
- tests := []struct {
- name string
- toolUsage string
- want OpenAIUsage
- }{
- {
- name: "valid tool usage takes atomic precedence",
- toolUsage: `{"input_tokens":4.6e1,"output_tokens":2459e0,"output_tokens_details":{"image_tokens":24590e-1}}`,
- want: OpenAIUsage{InputTokens: 46, OutputTokens: 2459, ImageOutputTokens: 2459},
- },
- {name: "absent", want: fallback},
- {name: "malformed field", toolUsage: `{"input_tokens":"46","output_tokens":2459,"output_tokens_details":{"image_tokens":2459}}`, want: fallback},
- {name: "fractional field", toolUsage: `{"input_tokens":46,"output_tokens":2459.5,"output_tokens_details":{"image_tokens":2459}}`, want: fallback},
- {name: "negative field", toolUsage: `{"input_tokens":46,"output_tokens":2459,"output_tokens_details":{"image_tokens":-1}}`, want: fallback},
- {name: "overflow field", toolUsage: `{"input_tokens":46,"output_tokens":9223372036854775808,"output_tokens_details":{"image_tokens":2459}}`, want: fallback},
- {name: "incomplete object", toolUsage: `{"input_tokens":46,"output_tokens":2459}`, want: fallback},
- {name: "hostile huge exponent", toolUsage: `{"input_tokens":1e1000000000,"output_tokens":2459,"output_tokens_details":{"image_tokens":2459}}`, want: fallback},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- toolUsageField := ""
- if tt.toolUsage != "" {
- toolUsageField = `,"tool_usage":{"image_gen":` + tt.toolUsage + `}`
- }
- payload := []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":3,"output_tokens":4,"output_tokens_details":{"image_tokens":2}}` + toolUsageField + `}}`)
- var got OpenAIUsage
- svc.parseOpenAIImagesSSEUsageBytes(payload, &got)
- require.Equal(t, tt.want, got)
- })
- }
- }
- func TestParseOpenAIImagesSSEUsageBytes_MalformedCompletedDoesNotOverrideUsage(t *testing.T) {
- svc := &OpenAIGatewayService{}
- var usage OpenAIUsage
- svc.parseOpenAIImagesSSEUsageBytes([]byte(`{"type":"response.output_item.done","item":{"type":"image_generation_call","result":"aW1hZ2U="}}`), &usage)
- svc.parseOpenAIImagesSSEUsageBytes([]byte(`{"type":"response.completed","response":{"usage":{"input_tokens":3,"output_tokens":4,"output_tokens_details":{"image_tokens":2}}}}`), &usage)
- svc.parseOpenAIImagesSSEUsageBytes([]byte(`{"type":"response.completed","response":{"tool_usage":{"image_gen":{"input_tokens":46,"output_tokens":2459,"output_tokens_details":{"image_tokens":2459}}}}} trailing`), &usage)
- require.Equal(t, OpenAIUsage{InputTokens: 3, OutputTokens: 4, ImageOutputTokens: 2}, usage)
- }
- func TestBoundedJSONNonNegativeInt(t *testing.T) {
- tests := []struct {
- name string
- raw string
- want int
- ok bool
- }{
- {name: "scale reduction before accumulation", raw: `10000000000000000000e-19`, want: 1, ok: true},
- {name: "decimal scale reduction", raw: `10000000000000000000.0e-19`, want: 1, ok: true},
- {name: "fractional after scale reduction", raw: `10000000000000000001e-19`, ok: false},
- {name: "overflow after scale reduction", raw: `92233720368547758080e-1`, ok: false},
- {name: "zero with negative exponent", raw: `0e-100`, want: 0, ok: true},
- {name: "zero beyond exponent bound", raw: `0e101`, want: 0, ok: true},
- {name: "zero padded decimal beyond exponent bound", raw: `0.000000e+000000000000000000000000000000000000000000000000101`, want: 0, ok: true},
- {name: "zero padded exponent", raw: `1e0000`, want: 1, ok: true},
- {name: "negative zero syntax", raw: `-0e101`, ok: false},
- {name: "hostile exponent", raw: `1e-1000`, ok: false},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- got, ok := boundedJSONNonNegativeInt(gjson.Parse(tt.raw))
- require.Equal(t, tt.ok, ok)
- require.Equal(t, tt.want, got)
- })
- }
- }
- func TestOpenAIGatewayServiceForwardImages_OAuthUpstreamHTTPErrorSurfacesRealError(t *testing.T) {
- gin.SetMode(gin.TestMode)
- body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","response_format":"b64_json"}`)
- // The non-failover upstream error path is shared by /generations and /edits;
- // use /generations here so the request parses without an uploaded image.
- 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.StatusBadRequest,
- Header: http.Header{
- "Content-Type": []string{"application/json"},
- "X-Request-Id": []string{"req_img_badreq"},
- },
- Body: io.NopCloser(strings.NewReader(
- `{"error":{"message":"Invalid value for 'size': expected one of 1024x1024, 1536x1024.","type":"invalid_request_error","param":"size","code":"unknown_parameter"}}`,
- )),
- },
- }
- 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, "invalid_request_error", upstreamErr.ErrorType)
- require.Equal(t, "unknown_parameter", upstreamErr.Code)
- // The client must receive the actual upstream status code and message instead
- // of a generic 502 "Upstream request failed".
- require.Equal(t, http.StatusBadRequest, rec.Code)
- require.Equal(t, "invalid_request_error", gjson.Get(rec.Body.String(), "error.type").String())
- require.Equal(t, "unknown_parameter", gjson.Get(rec.Body.String(), "error.code").String())
- require.Equal(t, "size", gjson.Get(rec.Body.String(), "error.param").String())
- require.Contains(t, gjson.Get(rec.Body.String(), "error.message").String(), "Invalid value for 'size'")
- }
- 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_OAuthNonStreamServerErrorReturnsFailoverBeforeFlush(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{
- httpUpstream: &httpUpstreamRecorder{
- resp: &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{
- "Content-Type": []string{"text/event-stream"},
- "X-Request-Id": []string{"req_img_server_error"},
- },
- Body: io.NopCloser(strings.NewReader(
- "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000021}}\n\n" +
- "data: {\"type\":\"error\",\"error\":{\"type\":\"server_error\",\"code\":\"server_error\",\"message\":\"The image service is temporarily unavailable.\"}}\n\n",
- )),
- },
- },
- }
- parsed, err := svc.ParseOpenAIImagesRequest(c, body)
- require.NoError(t, err)
- account := &Account{
- ID: 21,
- Name: "openai-oauth-server-error",
- 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 failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
- require.Contains(t, string(failoverErr.ResponseBody), "temporarily unavailable")
- require.False(t, c.Writer.Written())
- require.Empty(t, rec.Body.String())
- rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
- require.True(t, ok)
- events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
- require.True(t, ok)
- require.Len(t, events, 1)
- require.Equal(t, "failover", events[0].Kind)
- require.Equal(t, account.ID, events[0].AccountID)
- require.Equal(t, http.StatusBadGateway, events[0].UpstreamStatusCode)
- }
- func TestOpenAIGatewayServiceForwardImages_OAuth429PreservesRetryAfterCooldown(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{httpUpstream: &httpUpstreamRecorder{resp: &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{"Retry-After": []string{"1"}, "X-Request-Id": []string{"req_img_oauth_429"}},
- Body: io.NopCloser(strings.NewReader(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"rate limited"}}`)),
- }}}
- parsed, err := svc.ParseOpenAIImagesRequest(c, body)
- require.NoError(t, err)
- account := &Account{ID: 22, 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 failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.False(t, failoverErr.RetryableOnSameAccount)
- require.Zero(t, failoverErr.SameAccountRetryDelay)
- require.True(t, failoverErr.SameAccountRetryDeadline.IsZero())
- require.Equal(t, "1", failoverErr.ResponseHeaders.Get("Retry-After"))
- }
- func TestOpenAIImagesOAuthBodyReadTransportErrorFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{
- "X-Request-Id": []string{"req_h2_read_failure"},
- "X-Upstream": []string{"preserved"},
- },
- Body: &openAIImagesReadErrorBody{err: errors.New("stream error: stream ID 11; INTERNAL_ERROR; received from peer")},
- }
- account := &Account{ID: 5400, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- svc := &OpenAIGatewayService{}
- _, _, _, readErr := svc.handleOpenAIImagesOAuthNonStreamingResponse(resp, c, "b64_json", "gpt-image-2")
- require.Error(t, readErr)
- err := svc.handleOpenAIImagesOAuthResponseError(context.Background(), c, account, "gpt-image-2", "https://api.openai.com/v1/responses", resp, OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c), readErr)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
- require.JSONEq(t, `{"error":{"type":"upstream_error","code":"upstream_http2_stream_error","message":"Upstream HTTP/2 stream failed"}}`, string(failoverErr.ResponseBody))
- require.Equal(t, "req_h2_read_failure", failoverErr.ResponseHeaders.Get("x-request-id"))
- require.Equal(t, "preserved", failoverErr.ResponseHeaders.Get("x-upstream"))
- resp.Header.Set("X-Upstream", "mutated")
- require.Equal(t, "preserved", failoverErr.ResponseHeaders.Get("x-upstream"))
- rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
- require.True(t, ok)
- events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
- require.True(t, ok)
- require.Len(t, events, 1)
- require.Equal(t, "failover", events[0].Kind)
- require.Equal(t, "req_h2_read_failure", events[0].UpstreamRequestID)
- require.Equal(t, "Upstream HTTP/2 stream failed", events[0].Message)
- }
- func TestOpenAIImagesOAuthBodyReadErrorsNotMisclassified(t *testing.T) {
- tests := []struct {
- name string
- err error
- }{
- {name: "context canceled", err: context.Canceled},
- {name: "response too large", err: fmt.Errorf("%w: limit=1", ErrUpstreamResponseBodyTooLarge)},
- {name: "semantic error", err: &OpenAIImagesUpstreamError{StatusCode: http.StatusBadRequest, ErrorType: "invalid_request_error", Code: "invalid_value", Message: "bad image request"}},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
- resp := &http.Response{StatusCode: http.StatusOK, Header: make(http.Header)}
- err := tt.err
- if tt.name != "semantic error" && shouldClassifyOpenAIUpstreamStreamReadError(err, c.Request.Context()) {
- err = newOpenAIUpstreamStreamReadError(err)
- }
- got := (&OpenAIGatewayService{}).handleOpenAIImagesOAuthResponseError(context.Background(), c, &Account{Platform: PlatformOpenAI}, "gpt-image-2", "", resp, OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c), err)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(got, &failoverErr))
- require.ErrorIs(t, got, tt.err)
- })
- }
- }
- func TestOpenAIImagesOAuthTransportErrorAfterDownstreamWriteDoesNotFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
- before := OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)
- _, writeErr := c.Writer.Write([]byte("downstream image bytes"))
- require.NoError(t, writeErr)
- classifiedErr := newOpenAIUpstreamStreamReadError(errors.New("unexpected EOF"))
- account := &Account{ID: 5401, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- resp := &http.Response{Header: http.Header{"X-Request-Id": []string{"req_after_write"}}}
- err := (&OpenAIGatewayService{}).handleOpenAIImagesOAuthResponseError(context.Background(), c, account, "gpt-image-2", "", resp, before, classifiedErr)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr))
- require.ErrorIs(t, err, classifiedErr)
- rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
- require.True(t, ok)
- events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
- require.True(t, ok)
- require.Len(t, events, 1)
- require.Equal(t, "retry_exhausted_failover", events[0].Kind)
- }
- func TestShouldClassifyOpenAIUpstreamStreamReadErrorTransportStrings(t *testing.T) {
- for _, message := range []string{"unexpected EOF", "connection reset by peer", "broken pipe", "use of closed network connection"} {
- t.Run(message, func(t *testing.T) {
- require.True(t, shouldClassifyOpenAIUpstreamStreamReadError(errors.New(message)))
- })
- }
- canceledCtx, cancel := context.WithCancel(context.Background())
- cancel()
- require.False(t, shouldClassifyOpenAIUpstreamStreamReadError(errors.New("unexpected EOF"), canceledCtx))
- }
- func TestOpenAIGatewayServiceForwardImages_OAuthStreamServerErrorAfterFlushDoesNotFailover(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{
- httpUpstream: &httpUpstreamRecorder{
- resp: &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{
- "Content-Type": []string{"text/event-stream"},
- "X-Request-Id": []string{"req_img_server_error_after_partial"},
- },
- 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\":\"error\",\"error\":{\"type\":\"server_error\",\"code\":\"server_error\",\"message\":\"The image service failed after partial output.\"}}\n\n",
- )),
- },
- },
- }
- parsed, err := svc.ParseOpenAIImagesRequest(c, body)
- require.NoError(t, err)
- account := &Account{
- ID: 22,
- Name: "openai-oauth-partial-server-error",
- 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 failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr))
- var upstreamErr *OpenAIImagesUpstreamError
- require.ErrorAs(t, err, &upstreamErr)
- require.True(t, IsOpenAIImagesRetryableUpstreamError(upstreamErr))
- require.True(t, c.Writer.Written())
- require.Contains(t, rec.Body.String(), "event: image_generation.partial_image")
- require.Contains(t, rec.Body.String(), "event: error")
- require.Contains(t, rec.Body.String(), "failed after partial output")
- rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
- require.True(t, ok)
- events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
- require.True(t, ok)
- require.Len(t, events, 1)
- require.Equal(t, "retry_exhausted_failover", events[0].Kind)
- require.Equal(t, account.ID, events[0].AccountID)
- }
- func TestOpenAIImagesSSEClientErrorsAreNotRetryable(t *testing.T) {
- tests := []struct {
- name string
- payload string
- wantStatus int
- }{
- {
- name: "invalid request",
- payload: `{"type":"error","error":{"type":"invalid_request_error","code":"invalid_value","message":"bad size"}}`,
- wantStatus: http.StatusBadRequest,
- },
- {
- name: "content policy",
- payload: `{"type":"error","error":{"type":"image_generation_user_error","code":"content_policy_violation","message":"blocked"}}`,
- wantStatus: http.StatusBadRequest,
- },
- {
- name: "rate limit remains distinct from server error",
- payload: `{"type":"error","error":{"type":"rate_limit_exceeded","code":"rate_limit_exceeded","message":"try again"}}`,
- wantStatus: http.StatusTooManyRequests,
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- upstreamErr := openAIImagesUpstreamErrorFromSSEPayload([]byte(tt.payload))
- require.NotNil(t, upstreamErr)
- require.Equal(t, tt.wantStatus, upstreamErr.StatusCode)
- require.False(t, IsOpenAIImagesRetryableUpstreamError(upstreamErr))
- })
- }
- }
- 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_APIKeyAccessStateUsesTypedFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- body := []byte(`{"model":"gpt-image-1","prompt":"draw a cat"}`)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
- c.Request.Header.Set("Content-Type", "application/json")
- upstreamBody := []byte(`{"error":{"code":"organization_suspended","message":"Organization has been suspended"}}`)
- upstream := &httpUpstreamRecorder{resp: &http.Response{
- StatusCode: http.StatusForbidden,
- Header: http.Header{
- "Content-Type": []string{"application/json"},
- "X-Request-Id": []string{"req_images_access_state"},
- },
- Body: io.NopCloser(bytes.NewReader(upstreamBody)),
- }}
- svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
- parsed, err := svc.ParseOpenAIImagesRequest(c, body)
- require.NoError(t, err)
- account := &Account{
- ID: 51,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Credentials: map[string]any{
- "api_key": "sk-test",
- },
- }
- result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "")
- require.Nil(t, result)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusForbidden, failoverErr.StatusCode)
- require.Equal(t, GatewayFailureStageAccountAuth, failoverErr.Stage)
- require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
- require.Equal(t, OpenAIUpstreamAccessStateReason, failoverErr.Reason)
- require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
- require.Equal(t, http.StatusBadGateway, failoverErr.ClientStatusCode)
- require.Equal(t, openAIUpstreamAccessUnavailableClientMessage, failoverErr.ClientMessage)
- require.False(t, failoverErr.RetryableOnSameAccount)
- require.Equal(t, "req_images_access_state", failoverErr.ResponseHeaders.Get("x-request-id"))
- require.False(t, c.Writer.Written())
- }
- 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\":{\"input_tokens\":46,\"output_tokens\":2459,\"output_tokens_details\":{\"image_tokens\":2459},\"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)
- require.Equal(t, OpenAIUsage{InputTokens: 46, OutputTokens: 2459, ImageOutputTokens: 2459}, result.Usage)
- 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, `{"input_tokens":46,"output_tokens":2459,"output_tokens_details":{"image_tokens":2459},"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_ForcesImageToolChoice(t *testing.T) {
- parsed := &OpenAIImagesRequest{
- Endpoint: openAIImagesGenerationsEndpoint,
- Model: "gpt-image-2",
- Prompt: "draw a cat",
- }
- body, err := buildOpenAIImagesResponsesRequest(parsed, "gpt-image-2")
- require.NoError(t, err)
- require.NotNil(t, body)
- require.Equal(t, "image_generation", gjson.GetBytes(body, "tool_choice.type").String())
- require.Equal(t, "image_generation", gjson.GetBytes(body, "tools.0.type").String())
- require.Equal(t, "gpt-image-2", gjson.GetBytes(body, "tools.0.model").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 TestBuildOpenAIImagesResponsesRequest_RequiresVerbatimUserPrompt(t *testing.T) {
- prompt := "画一个蓝色马克杯,杯身只写“SkelOT”,保持大小写;白色背景,不要增加其他文字。"
- parsed := &OpenAIImagesRequest{
- Endpoint: openAIImagesGenerationsEndpoint,
- Model: "gpt-image-2",
- Prompt: prompt,
- N: 1,
- }
- body, err := buildOpenAIImagesResponsesRequest(parsed, "gpt-image-2")
- require.NoError(t, err)
- require.Equal(t, openAIImagesVerbatimPromptInstructions, gjson.GetBytes(body, "instructions").String())
- require.Equal(t, prompt, gjson.GetBytes(body, "input.0.content.0.text").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)
- }
|