| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546354735483549355035513552355335543555355635573558355935603561356235633564356535663567356835693570357135723573357435753576357735783579358035813582358335843585358635873588358935903591359235933594359535963597359835993600360136023603360436053606360736083609361036113612361336143615361636173618361936203621362236233624362536263627362836293630363136323633363436353636363736383639364036413642364336443645364636473648364936503651365236533654365536563657365836593660366136623663366436653666366736683669367036713672367336743675367636773678367936803681368236833684368536863687368836893690369136923693369436953696369736983699370037013702370337043705370637073708370937103711371237133714371537163717371837193720372137223723372437253726372737283729373037313732373337343735373637373738373937403741374237433744374537463747374837493750375137523753375437553756375737583759376037613762376337643765376637673768376937703771377237733774377537763777377837793780378137823783378437853786378737883789379037913792379337943795379637973798379938003801380238033804380538063807380838093810381138123813381438153816381738183819382038213822382338243825382638273828382938303831383238333834383538363837383838393840384138423843384438453846384738483849385038513852385338543855385638573858385938603861386238633864386538663867386838693870387138723873387438753876387738783879388038813882388338843885388638873888388938903891389238933894389538963897389838993900390139023903390439053906390739083909391039113912391339143915391639173918391939203921392239233924392539263927392839293930393139323933393439353936393739383939394039413942394339443945394639473948394939503951395239533954395539563957395839593960396139623963396439653966 |
- package service
- import (
- "bufio"
- "bytes"
- "context"
- "errors"
- "fmt"
- "io"
- "net/http"
- "net/http/httptest"
- "strings"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/model"
- "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
- "github.com/cespare/xxhash/v2"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/require"
- "github.com/tidwall/gjson"
- )
- // 编译期接口断言
- var _ AccountRepository = (*stubOpenAIAccountRepo)(nil)
- var _ GatewayCache = (*stubGatewayCache)(nil)
- type stubOpenAIAccountRepo struct {
- AccountRepository
- accounts []Account
- }
- type tempUnschedulableOpenAIAccountRepo struct {
- stubOpenAIAccountRepo
- modelRateLimitAccountID int64
- modelRateLimitKey string
- }
- func (r *tempUnschedulableOpenAIAccountRepo) SetModelRateLimit(_ context.Context, accountID int64, modelKey string, _ time.Time, _ ...string) error {
- r.modelRateLimitAccountID = accountID
- r.modelRateLimitKey = modelKey
- return nil
- }
- type snapshotUpdateAccountRepo struct {
- stubOpenAIAccountRepo
- updateExtraCalls chan map[string]any
- }
- func (r *snapshotUpdateAccountRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
- if r.updateExtraCalls != nil {
- copied := make(map[string]any, len(updates))
- for k, v := range updates {
- copied[k] = v
- }
- r.updateExtraCalls <- copied
- }
- return nil
- }
- func (r stubOpenAIAccountRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
- for i := range r.accounts {
- if r.accounts[i].ID == id {
- return &r.accounts[i], nil
- }
- }
- return nil, errors.New("account not found")
- }
- func (r stubOpenAIAccountRepo) GetByIDs(ctx context.Context, ids []int64) ([]*Account, error) {
- if len(ids) == 0 {
- return []*Account{}, nil
- }
- index := make(map[int64]*Account, len(r.accounts))
- for i := range r.accounts {
- account := &r.accounts[i]
- index[account.ID] = account
- }
- out := make([]*Account, 0, len(ids))
- seen := make(map[int64]struct{}, len(ids))
- for _, id := range ids {
- if _, ok := seen[id]; ok {
- continue
- }
- seen[id] = struct{}{}
- if account, ok := index[id]; ok {
- out = append(out, account)
- }
- }
- return out, nil
- }
- func (r stubOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
- var result []Account
- for _, acc := range r.accounts {
- if acc.Platform == platform {
- result = append(result, acc)
- }
- }
- return result, nil
- }
- func (r stubOpenAIAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
- var result []Account
- for _, acc := range r.accounts {
- if acc.Platform == platform {
- result = append(result, acc)
- }
- }
- return result, nil
- }
- func (r stubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
- return r.ListSchedulableByPlatform(ctx, platform)
- }
- func TestOpenAIGatewayService_ForwardAsAnthropic_CapacityShedReturnsRequestScopedFailoverWithoutCommit(t *testing.T) {
- gin.SetMode(gin.TestMode)
- body := []byte(`{"model":"gpt-5.4","max_tokens":32,"messages":[{"role":"user","content":"hello"}],"stream":false}`)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
- c.Request.Header.Set("Content-Type", "application/json")
- upstreamBody := []byte(`{"error":{"message":"Our servers are currently overloaded. Please try again later."}}`)
- upstream := &httpUpstreamRecorder{responses: []*http.Response{
- {
- StatusCode: http.StatusBadRequest,
- Header: http.Header{"Content-Type": []string{"application/json"}},
- Body: io.NopCloser(bytes.NewReader(upstreamBody)),
- },
- {
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"text/event-stream"}},
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- `data: {"type":"response.output_text.delta","delta":"ok"}`,
- "",
- `data: {"type":"response.completed","response":{"id":"resp_second","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1}}}`,
- "",
- }, "\n"))),
- },
- }}
- repo := &tempUnschedulableOpenAIAccountRepo{}
- rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
- svc := &OpenAIGatewayService{
- cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{
- Enabled: false, AllowInsecureHTTP: true,
- }}},
- httpUpstream: upstream,
- rateLimitService: rateLimitService,
- }
- account := &Account{
- ID: 5099, Name: "temporary-unschedulable", Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey, Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "base_url": "http://upstream.example",
- "model_provider": "env-openai",
- "temp_unschedulable_enabled": true,
- "temp_unschedulable_rules": []any{map[string]any{
- "error_code": float64(http.StatusBadRequest),
- "keywords": []any{"our servers are currently overloaded", "please try again later"},
- "duration_minutes": float64(1),
- }},
- },
- }
- _, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusBadRequest, failoverErr.StatusCode)
- require.True(t, failoverErr.ShouldRetryNextAccount())
- require.True(t, failoverErr.RetryableOnSameAccount)
- require.True(t, failoverErr.RequestScopedTransient)
- require.Zero(t, repo.modelRateLimitAccountID, "request-scoped capacity shedding must not change account health")
- require.Empty(t, repo.modelRateLimitKey)
- require.False(t, IsResponseCommitted(c))
- require.Equal(t, http.StatusOK, rec.Code)
- require.Empty(t, rec.Body.String())
- secondRec := httptest.NewRecorder()
- secondContext, _ := gin.CreateTestContext(secondRec)
- secondContext.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
- secondContext.Request.Header.Set("Content-Type", "application/json")
- secondAccount := *account
- secondAccount.ID = 5100
- secondAccount.Name = "healthy-failover-account"
- result, secondErr := svc.ForwardAsAnthropic(context.Background(), secondContext, &secondAccount, body, "", "")
- require.NoError(t, secondErr)
- require.NotNil(t, result)
- require.Equal(t, "resp_second", result.ResponseID)
- require.NotEmpty(t, secondRec.Body.String())
- }
- func TestFailoverOpenAIUpstreamHTTPError_NilContextSkipsTempUnschedulablePolicy(t *testing.T) {
- repo := &tempUnschedulableOpenAIAccountRepo{}
- svc := &OpenAIGatewayService{
- rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil),
- }
- account := &Account{
- ID: 5099, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
- Credentials: map[string]any{
- "temp_unschedulable_enabled": true,
- "temp_unschedulable_rules": []any{map[string]any{
- "error_code": float64(http.StatusBadRequest),
- "keywords": []any{"custom temporary outage"},
- "duration_minutes": float64(1),
- }},
- },
- }
- body := []byte(`{"error":{"message":"Custom temporary outage."}}`)
- resp := &http.Response{StatusCode: http.StatusBadRequest, Header: http.Header{}}
- got := svc.failoverOpenAIUpstreamHTTPError(
- context.Background(), nil, account, resp, body,
- "Custom temporary outage.", "gpt-5.4",
- )
- require.Nil(t, got)
- require.Zero(t, repo.modelRateLimitAccountID)
- require.Empty(t, repo.modelRateLimitKey)
- }
- type groupAwareStubOpenAIAccountRepo struct {
- stubOpenAIAccountRepo
- }
- func (r groupAwareStubOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
- var result []Account
- for _, acc := range r.accounts {
- if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, &groupID) {
- result = append(result, acc)
- }
- }
- return result, nil
- }
- func (r groupAwareStubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
- var result []Account
- for _, acc := range r.accounts {
- if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, nil) {
- result = append(result, acc)
- }
- }
- return result, nil
- }
- type stubConcurrencyCache struct {
- ConcurrencyCache
- loadBatchErr error
- loadMap map[int64]*AccountLoadInfo
- acquireResults map[int64]bool
- waitCounts map[int64]int
- skipDefaultLoad bool
- }
- type cancelReadCloser struct{}
- func (c cancelReadCloser) Read(p []byte) (int, error) { return 0, context.Canceled }
- func (c cancelReadCloser) Close() error { return nil }
- type errReadCloser struct {
- err error
- }
- func (r errReadCloser) Read([]byte) (int, error) { return 0, r.err }
- func (r errReadCloser) Close() error { return nil }
- type openAIStreamReadThenErrorCloser struct {
- reader *strings.Reader
- err error
- }
- func (r *openAIStreamReadThenErrorCloser) Read(p []byte) (int, error) {
- if r.reader != nil && r.reader.Len() > 0 {
- return r.reader.Read(p)
- }
- return 0, r.err
- }
- func (r *openAIStreamReadThenErrorCloser) Close() error { return nil }
- type failingGinWriter struct {
- gin.ResponseWriter
- failAfter int
- writes int
- }
- func (w *failingGinWriter) Write(p []byte) (int, error) {
- if w.writes >= w.failAfter {
- return 0, errors.New("write failed")
- }
- w.writes++
- return w.ResponseWriter.Write(p)
- }
- func (c stubConcurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
- if c.acquireResults != nil {
- if result, ok := c.acquireResults[accountID]; ok {
- return result, nil
- }
- }
- return true, nil
- }
- func (c stubConcurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int64, requestID string) error {
- return nil
- }
- func (c stubConcurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []AccountWithConcurrency) (map[int64]*AccountLoadInfo, error) {
- if c.loadBatchErr != nil {
- return nil, c.loadBatchErr
- }
- out := make(map[int64]*AccountLoadInfo, len(accounts))
- if c.skipDefaultLoad && c.loadMap != nil {
- for _, acc := range accounts {
- if load, ok := c.loadMap[acc.ID]; ok {
- out[acc.ID] = load
- }
- }
- return out, nil
- }
- for _, acc := range accounts {
- if c.loadMap != nil {
- if load, ok := c.loadMap[acc.ID]; ok {
- out[acc.ID] = load
- continue
- }
- }
- out[acc.ID] = &AccountLoadInfo{AccountID: acc.ID, LoadRate: 0}
- }
- return out, nil
- }
- func TestOpenAIGatewayService_GenerateSessionHash_Priority(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- svc := &OpenAIGatewayService{}
- bodyWithKey := []byte(`{"prompt_cache_key":"ses_aaa"}`)
- // 1) session_id header wins
- c.Request.Header.Set("session_id", "sess-123")
- c.Request.Header.Set("conversation_id", "conv-456")
- h1 := svc.GenerateSessionHash(c, bodyWithKey)
- if h1 == "" {
- t.Fatalf("expected non-empty hash")
- }
- // 2) conversation_id used when session_id absent
- c.Request.Header.Del("session_id")
- h2 := svc.GenerateSessionHash(c, bodyWithKey)
- if h2 == "" {
- t.Fatalf("expected non-empty hash")
- }
- if h1 == h2 {
- t.Fatalf("expected different hashes for different keys")
- }
- // 3) prompt_cache_key used when both headers absent
- c.Request.Header.Del("conversation_id")
- h3 := svc.GenerateSessionHash(c, bodyWithKey)
- if h3 == "" {
- t.Fatalf("expected non-empty hash")
- }
- if h2 == h3 {
- t.Fatalf("expected different hashes for different keys")
- }
- // 4) empty when no signals
- h4 := svc.GenerateSessionHash(c, []byte(`{}`))
- if h4 != "" {
- t.Fatalf("expected empty hash when no signals")
- }
- }
- func TestOpenAIGatewayService_ClientSessionHeaderPriority(t *testing.T) {
- gin.SetMode(gin.TestMode)
- recorder := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(recorder)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Set("api_key", &APIKey{ID: 901, Group: &Group{Platform: PlatformGrok}})
- headers := []struct {
- name string
- value string
- }{
- {name: "session-id", value: "codex-session"},
- {name: "session_id", value: "generic-session"},
- {name: "conversation_id", value: "generic-conversation"},
- {name: openCodeSessionAffinityHeader, value: "opencode-affinity"},
- {name: openCodeSessionIDHeader, value: "opencode-session-id"},
- {name: openCodeNativeSessionHeader, value: "opencode-native-session"},
- {name: codeBuddyConversationHeader, value: "codebuddy-conversation"},
- {name: grokConversationIDHeader, value: "grok-conversation"},
- }
- for _, header := range headers {
- c.Request.Header.Set(header.name, header.value)
- }
- svc := &OpenAIGatewayService{}
- body := []byte(`{"prompt_cache_key":"body-session"}`)
- for _, header := range headers {
- require.Equal(t, header.value, svc.ExtractSessionID(c, body), header.name)
- require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String(header.value)), svc.GenerateExplicitSessionHash(c, body), header.name)
- if header.name != grokConversationIDHeader {
- require.Equal(t, header.value, explicitOpenAISessionID(c, body), header.name)
- }
- c.Request.Header.Del(header.name)
- }
- require.Equal(t, "body-session", svc.ExtractSessionID(c, body))
- }
- func TestOpenAIGatewayService_CodexSessionIDKeepsReconnectHashStable(t *testing.T) {
- gin.SetMode(gin.TestMode)
- recorder := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(recorder)
- c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
- c.Request.Header.Set("session-id", "codex-reconnect-session")
- svc := &OpenAIGatewayService{}
- warmup := []byte(`{
- "type":"response.create",
- "model":"gpt-5.6-sol",
- "generate":false,
- "tools":[{"type":"custom","name":"exec"}],
- "input":[{"role":"user","content":"warmup"}]
- }`)
- business := []byte(`{
- "type":"response.create",
- "model":"gpt-5.6-sol",
- "input":[{"role":"user","content":"install codex"}]
- }`)
- require.Equal(t, svc.GenerateSessionHash(c, warmup), svc.GenerateSessionHash(c, business))
- require.Equal(t, "codex-reconnect-session", svc.ExtractSessionID(c, business))
- }
- func TestOpenAIGatewayService_ClientSessionHeadersIgnorePerRequestIDs(t *testing.T) {
- gin.SetMode(gin.TestMode)
- recorder := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(recorder)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- for name, value := range map[string]string{
- "X-Conversation-Request-ID": "request-rotates-every-turn",
- "X-Conversation-Message-ID": "message-rotates-every-turn",
- "X-Request-ID": "generic-request-id",
- } {
- c.Request.Header.Set(name, value)
- }
- svc := &OpenAIGatewayService{}
- require.Empty(t, explicitOpenAIHeaderSessionID(c))
- require.Empty(t, svc.ExtractSessionID(c, nil))
- require.Empty(t, svc.GenerateExplicitSessionHash(c, nil))
- }
- func TestOpenAIGatewayService_GenerateSessionHash_UsesXXHash64(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("session_id", "sess-fixed-value")
- svc := &OpenAIGatewayService{}
- got := svc.GenerateSessionHash(c, nil)
- want := fmt.Sprintf("%016x", xxhash.Sum64String("sess-fixed-value"))
- require.Equal(t, want, got)
- }
- func TestOpenAIGatewayService_GenerateSessionHash_AttachesLegacyHashToContext(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("session_id", "sess-legacy-check")
- svc := &OpenAIGatewayService{}
- sessionHash := svc.GenerateSessionHash(c, nil)
- require.NotEmpty(t, sessionHash)
- require.NotNil(t, c.Request)
- require.NotNil(t, c.Request.Context())
- require.NotEmpty(t, openAILegacySessionHashFromContext(c.Request.Context()))
- }
- func TestExtractOpenAIResponseIDFromJSONBytes(t *testing.T) {
- require.Equal(t, "resp_json", extractOpenAIResponseIDFromJSONBytes([]byte(`{"id":"resp_json"}`)))
- require.Equal(t, "resp_sse", extractOpenAIResponseIDFromJSONBytes([]byte(`{"type":"response.completed","response":{"id":"resp_sse"}}`)))
- require.Empty(t, extractOpenAIResponseIDFromJSONBytes([]byte(`{"response":{}}`)))
- require.Empty(t, extractOpenAIResponseIDFromJSONBytes([]byte(`not-json`)))
- }
- // 复现 #4386:gpt-image-2 /v1/images/edits 的 usage 携带 input_tokens_details.image_tokens,
- // 提取器须将图片输入 token 单独填入 ImageInputTokens(此前被丢弃并入 InputTokens 按文本价计费)。
- func TestExtractOpenAIUsage_CapturesImageInputTokens(t *testing.T) {
- body := []byte(`{"usage":{"input_tokens":371,"input_tokens_details":{"image_tokens":352,"text_tokens":19},"output_tokens":439,"output_tokens_details":{"image_tokens":439,"text_tokens":0},"total_tokens":810}}`)
- usage, ok := extractOpenAIUsageFromJSONBytes(body)
- require.True(t, ok)
- require.Equal(t, 371, usage.InputTokens)
- require.Equal(t, 352, usage.ImageInputTokens)
- require.Equal(t, 439, usage.OutputTokens)
- require.Equal(t, 439, usage.ImageOutputTokens)
- // prompt_tokens_details 回退路径(部分上游用 prompt_tokens 口径)。
- promptStyle := []byte(`{"usage":{"prompt_tokens":100,"prompt_tokens_details":{"image_tokens":80}}}`)
- pu, ok := extractOpenAIUsageFromJSONBytes(promptStyle)
- require.True(t, ok)
- require.Equal(t, 100, pu.InputTokens)
- require.Equal(t, 80, pu.ImageInputTokens)
- // 纯文本请求:无 image_tokens 时 ImageInputTokens 为 0,行为不变。
- textOnly := []byte(`{"usage":{"input_tokens":50,"output_tokens":10}}`)
- tu, ok := extractOpenAIUsageFromJSONBytes(textOnly)
- require.True(t, ok)
- require.Zero(t, tu.ImageInputTokens)
- }
- func TestExtractOpenAIUsage_ReadsClineDataEnvelope(t *testing.T) {
- body := []byte(`{"data":{"choices":[{"message":{"content":"OK"}}],"usage":{"prompt_tokens":8,"completion_tokens":27,"total_tokens":35,"prompt_tokens_details":{"cached_tokens":4}}},"success":true}`)
- usage, ok := extractOpenAIUsageFromJSONBytes(body)
- require.True(t, ok)
- require.Equal(t, 8, usage.InputTokens)
- require.Equal(t, 27, usage.OutputTokens)
- require.Equal(t, 4, usage.CacheReadInputTokens)
- }
- func TestExtractOpenAIUsage_ReadsWrappedResponsesDataEnvelope(t *testing.T) {
- body := []byte(`{"data":{"response":{"usage":{"input_tokens":11,"output_tokens":5,"total_tokens":16,"input_tokens_details":{"cached_tokens":2}}}}}`)
- usage, ok := extractOpenAIUsageFromJSONBytes(body)
- require.True(t, ok)
- require.Equal(t, 11, usage.InputTokens)
- require.Equal(t, 5, usage.OutputTokens)
- require.Equal(t, 2, usage.CacheReadInputTokens)
- }
- func TestExtractOpenAIUsage_PreservesResponseUsagePriority(t *testing.T) {
- body := []byte(`{"data":{"usage":{"prompt_tokens":100,"completion_tokens":50}},"response":{"usage":{"input_tokens":11,"output_tokens":5}}}`)
- usage, ok := extractOpenAIUsageFromJSONBytes(body)
- require.True(t, ok)
- require.Equal(t, 11, usage.InputTokens)
- require.Equal(t, 5, usage.OutputTokens)
- }
- func TestOpenAIGatewayService_BindHTTPResponseAccount(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- groupID := int64(4201)
- c.Set("api_key", &APIKey{ID: 501, GroupID: &groupID})
- SetOpenAIHTTPResponseOwner(c, 601, 501)
- svc := &OpenAIGatewayService{}
- account := &Account{ID: 37001, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- svc.bindHTTPResponseAccount(context.Background(), c, account, "resp_http_001")
- got, err := svc.getOpenAIWSStateStore().GetResponseAccount(context.Background(), groupID, "resp_http_001")
- require.NoError(t, err)
- require.Equal(t, account.ID, got)
- owned, err := svc.ValidateOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_http_001", 601, 501)
- require.NoError(t, err)
- require.True(t, owned)
- owned, err = svc.ValidateOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_http_001", 601, 502)
- require.NoError(t, err)
- require.True(t, owned, "API keys owned by the same downstream user remain interoperable")
- owned, err = svc.ValidateOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_http_001", 602, 501)
- require.NoError(t, err)
- require.False(t, owned)
- owned, err = svc.ValidateOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_unknown", 601, 501)
- require.NoError(t, err)
- require.False(t, owned)
- }
- func TestOpenAIGatewayService_GenerateExplicitSessionHash_SkipsContentFallback(t *testing.T) {
- gin.SetMode(gin.TestMode)
- svc := &OpenAIGatewayService{}
- body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`)
- t.Run("stateless image body stays unstuck", func(t *testing.T) {
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
- require.Empty(t, svc.GenerateExplicitSessionHash(c, body))
- require.Empty(t, openAILegacySessionHashFromContext(c.Request.Context()))
- })
- t.Run("prompt_cache_key is explicit", func(t *testing.T) {
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
- got := svc.GenerateExplicitSessionHash(c, []byte(`{"model":"gpt-image-2","prompt_cache_key":"image-session"}`))
- require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String("image-session")), got)
- require.NotEmpty(t, openAILegacySessionHashFromContext(c.Request.Context()))
- })
- t.Run("header overrides body", func(t *testing.T) {
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
- c.Request.Header.Set("session_id", "header-session")
- got := svc.GenerateExplicitSessionHash(c, []byte(`{"prompt_cache_key":"body-session"}`))
- require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String("header-session")), got)
- })
- }
- func TestOpenAIGatewayService_GenerateSessionHashWithFallback(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- svc := &OpenAIGatewayService{}
- seed := "openai_ws_ingress:9:100:200"
- got := svc.GenerateSessionHashWithFallback(c, []byte(`{}`), seed)
- want := fmt.Sprintf("%016x", xxhash.Sum64String(seed))
- require.Equal(t, want, got)
- require.NotEmpty(t, openAILegacySessionHashFromContext(c.Request.Context()))
- empty := svc.GenerateSessionHashWithFallback(c, []byte(`{}`), " ")
- require.Equal(t, "", empty)
- }
- func TestOpenAIGatewayService_GenerateSessionHash_ContentFallback(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil)
- svc := &OpenAIGatewayService{}
- body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"You are helpful."},{"role":"user","content":"Hello"}]}`)
- hash := svc.GenerateSessionHash(c, body)
- require.NotEmpty(t, hash, "content-based fallback should produce a hash")
- hash2 := svc.GenerateSessionHash(c, body)
- require.Equal(t, hash, hash2, "same content should produce same hash")
- bodyExtended := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"You are helpful."},{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi!"},{"role":"user","content":"How are you?"}]}`)
- hashExtended := svc.GenerateSessionHash(c, bodyExtended)
- require.Equal(t, hash, hashExtended, "hash should be stable across later turns")
- bodyDifferent := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"Different question"}]}`)
- hashDifferent := svc.GenerateSessionHash(c, bodyDifferent)
- require.NotEqual(t, hash, hashDifferent, "different content should produce different hash")
- }
- func TestOpenAIGatewayService_GenerateSessionHash_ExplicitSignalWinsOverContent(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil)
- svc := &OpenAIGatewayService{}
- body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"Hello"}]}`)
- contentHash := svc.GenerateSessionHash(c, body)
- require.NotEmpty(t, contentHash)
- c.Request.Header.Set("session_id", "explicit-session")
- explicitHash := svc.GenerateSessionHash(c, body)
- require.NotEmpty(t, explicitHash)
- require.NotEqual(t, contentHash, explicitHash, "explicit session_id should override content fallback")
- }
- func TestOpenAIGatewayService_GenerateSessionHash_EmptyBodyStillEmpty(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil)
- svc := &OpenAIGatewayService{}
- require.Empty(t, svc.GenerateSessionHash(c, []byte(`{}`)))
- require.Empty(t, svc.GenerateSessionHash(c, nil))
- }
- func (c stubConcurrencyCache) GetAccountWaitingCount(ctx context.Context, accountID int64) (int, error) {
- if c.waitCounts != nil {
- if count, ok := c.waitCounts[accountID]; ok {
- return count, nil
- }
- }
- return 0, nil
- }
- type stubGatewayCache struct {
- sessionBindings map[string]int64
- deletedSessions map[string]int
- }
- func (c *stubGatewayCache) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) {
- if id, ok := c.sessionBindings[sessionHash]; ok {
- return id, nil
- }
- return 0, errors.New("not found")
- }
- func (c *stubGatewayCache) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error {
- if c.sessionBindings == nil {
- c.sessionBindings = make(map[string]int64)
- }
- c.sessionBindings[sessionHash] = accountID
- return nil
- }
- func (c *stubGatewayCache) RefreshSessionTTL(ctx context.Context, groupID int64, sessionHash string, ttl time.Duration) error {
- return nil
- }
- func (c *stubGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error {
- if c.sessionBindings == nil {
- return nil
- }
- if c.deletedSessions == nil {
- c.deletedSessions = make(map[string]int)
- }
- c.deletedSessions[sessionHash]++
- delete(c.sessionBindings, sessionHash)
- return nil
- }
- func (c *stubGatewayCache) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
- return nil
- }
- func (c *stubGatewayCache) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
- return nil, nil
- }
- func (c *stubGatewayCache) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
- return true, nil
- }
- func (c *stubGatewayCache) ReleaseGrokVideoBilled(_ context.Context, _ string) error {
- return nil
- }
- func (c *stubGatewayCache) SetReasoningContent(_ context.Context, _ string, _ string, _ time.Duration) error {
- return nil
- }
- func (c *stubGatewayCache) GetReasoningContent(_ context.Context, _ string) (string, error) {
- return "", ErrReasoningContentNotFound
- }
- func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) {
- now := time.Now()
- resetAt := now.Add(10 * time.Minute)
- groupID := int64(1)
- rateLimited := Account{
- ID: 1,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- RateLimitResetAt: &resetAt,
- }
- available := Account{
- ID: 2,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 1,
- }
- svc := &OpenAIGatewayService{
- accountRepo: stubOpenAIAccountRepo{accounts: []Account{rateLimited, available}},
- concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.2", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.Account == nil {
- t.Fatalf("expected selection with account")
- }
- if selection.Account.ID != available.ID {
- t.Fatalf("expected account %d, got %d", available.ID, selection.Account.ID)
- }
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_ImageRateLimitSkipsOnlyImageRequests(t *testing.T) {
- future := time.Now().Add(10 * time.Minute).Format(time.RFC3339)
- groupID := int64(1)
- imageLimited := Account{
- ID: 1,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- modelRateLimitsKey: map[string]any{
- openAIImageGenerationRateLimitKey: map[string]any{
- "rate_limit_reset_at": future,
- },
- },
- },
- }
- available := Account{
- ID: 2,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 1,
- }
- svc := &OpenAIGatewayService{
- accountRepo: stubOpenAIAccountRepo{accounts: []Account{imageLimited, available}},
- concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
- }
- imageSelection, err := svc.SelectAccountWithLoadAwareness(WithOpenAIImageGenerationIntent(context.Background()), &groupID, "", "gpt-5.4", nil)
- require.NoError(t, err)
- require.NotNil(t, imageSelection)
- require.Equal(t, available.ID, imageSelection.Account.ID)
- if imageSelection.ReleaseFunc != nil {
- imageSelection.ReleaseFunc()
- }
- textSelection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.4", nil)
- require.NoError(t, err)
- require.NotNil(t, textSelection)
- require.Equal(t, imageLimited.ID, textSelection.Account.ID)
- if textSelection.ReleaseFunc != nil {
- textSelection.ReleaseFunc()
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulableWhenNoConcurrencyService(t *testing.T) {
- now := time.Now()
- resetAt := now.Add(10 * time.Minute)
- groupID := int64(1)
- rateLimited := Account{
- ID: 1,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- RateLimitResetAt: &resetAt,
- }
- available := Account{
- ID: 2,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 1,
- }
- svc := &OpenAIGatewayService{
- accountRepo: stubOpenAIAccountRepo{accounts: []Account{rateLimited, available}},
- // concurrencyService is nil, forcing the non-load-batch selection path.
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.2", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.Account == nil {
- t.Fatalf("expected selection with account")
- }
- if selection.Account.ID != available.ID {
- t.Fatalf("expected account %d, got %d", available.ID, selection.Account.ID)
- }
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAISelectAccountForModelWithExclusions_StickyUnschedulableClearsSession(t *testing.T) {
- sessionHash := "session-1"
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusDisabled, Schedulable: true, Concurrency: 1},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
- },
- }
- cache := &stubGatewayCache{
- sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- }
- acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
- }
- if acc == nil || acc.ID != 2 {
- t.Fatalf("expected account 2, got %+v", acc)
- }
- if cache.deletedSessions["openai:"+sessionHash] != 1 {
- t.Fatalf("expected sticky session to be deleted")
- }
- if cache.sessionBindings["openai:"+sessionHash] != 2 {
- t.Fatalf("expected sticky session to bind to account 2")
- }
- }
- func TestOpenAISelectAccountForModelWithExclusions_StickyOutsideGroupClearsSession(t *testing.T) {
- sessionHash := "session-outside-group"
- groupID := int64(1001)
- repo := groupAwareStubOpenAIAccountRepo{
- stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, AccountGroups: []AccountGroup{{GroupID: groupID}}},
- },
- },
- }
- cache := &stubGatewayCache{
- sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- }
- acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), &groupID, sessionHash, "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
- }
- if acc == nil || acc.ID != 2 {
- t.Fatalf("expected account 2, got %+v", acc)
- }
- if cache.deletedSessions["openai:"+sessionHash] != 1 {
- t.Fatalf("expected sticky session to be deleted")
- }
- if cache.sessionBindings["openai:"+sessionHash] != 2 {
- t.Fatalf("expected sticky session to bind to account 2")
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_StickyUnschedulableClearsSession(t *testing.T) {
- sessionHash := "session-2"
- groupID := int64(1)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusDisabled, Schedulable: true, Concurrency: 1},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
- },
- }
- cache := &stubGatewayCache{
- sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
- t.Fatalf("expected account 2, got %+v", selection)
- }
- if cache.deletedSessions["openai:"+sessionHash] != 1 {
- t.Fatalf("expected sticky session to be deleted")
- }
- if cache.sessionBindings["openai:"+sessionHash] != 2 {
- t.Fatalf("expected sticky session to bind to account 2")
- }
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_StickyOutsideGroupClearsSession(t *testing.T) {
- sessionHash := "session-load-outside-group"
- groupID := int64(1002)
- repo := groupAwareStubOpenAIAccountRepo{
- stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, AccountGroups: []AccountGroup{{GroupID: groupID}}},
- },
- },
- }
- cache := &stubGatewayCache{
- sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
- t.Fatalf("expected account 2, got %+v", selection)
- }
- if cache.deletedSessions["openai:"+sessionHash] != 1 {
- t.Fatalf("expected sticky session to be deleted")
- }
- if cache.sessionBindings["openai:"+sessionHash] != 2 {
- t.Fatalf("expected sticky session to bind to account 2")
- }
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAISelectAccountForModelWithExclusions_NoModelSupport(t *testing.T) {
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {
- ID: 1,
- Platform: PlatformOpenAI,
- Status: StatusActive,
- Schedulable: true,
- Credentials: map[string]any{"model_mapping": map[string]any{"gpt-3.5-turbo": "gpt-3.5-turbo"}},
- },
- },
- }
- cache := &stubGatewayCache{}
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- }
- acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, "", "gpt-4", nil)
- if err == nil {
- t.Fatalf("expected error for unsupported model")
- }
- if acc != nil {
- t.Fatalf("expected nil account for unsupported model")
- }
- if !strings.Contains(err.Error(), "supporting model") {
- t.Fatalf("unexpected error: %v", err)
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_LoadBatchErrorFallback(t *testing.T) {
- groupID := int64(1)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 2},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- },
- }
- cache := &stubGatewayCache{}
- concurrencyCache := stubConcurrencyCache{
- loadBatchErr: errors.New("load batch failed"),
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "fallback", "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.Account == nil {
- t.Fatalf("expected selection")
- }
- if selection.Account.ID != 2 {
- t.Fatalf("expected account 2, got %d", selection.Account.ID)
- }
- if cache.sessionBindings["openai:fallback"] != 2 {
- t.Fatalf("expected sticky session updated")
- }
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_NoSlotFallbackWait(t *testing.T) {
- groupID := int64(1)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- },
- }
- cache := &stubGatewayCache{}
- concurrencyCache := stubConcurrencyCache{
- acquireResults: map[int64]bool{1: false},
- loadMap: map[int64]*AccountLoadInfo{
- 1: {AccountID: 1, LoadRate: 10},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.WaitPlan == nil {
- t.Fatalf("expected wait plan fallback")
- }
- if selection.Account == nil || selection.Account.ID != 1 {
- t.Fatalf("expected account 1")
- }
- }
- func TestOpenAISelectAccountForModelWithExclusions_SetsStickyBinding(t *testing.T) {
- sessionHash := "bind"
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- },
- }
- cache := &stubGatewayCache{}
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- }
- acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
- }
- if acc == nil || acc.ID != 1 {
- t.Fatalf("expected account 1")
- }
- if cache.sessionBindings["openai:"+sessionHash] != 1 {
- t.Fatalf("expected sticky session binding")
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_StickyWaitPlan(t *testing.T) {
- sessionHash := "sticky-wait"
- groupID := int64(1)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- },
- }
- cache := &stubGatewayCache{
- sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
- }
- concurrencyCache := stubConcurrencyCache{
- acquireResults: map[int64]bool{1: false},
- waitCounts: map[int64]int{1: 0},
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.WaitPlan == nil {
- t.Fatalf("expected sticky wait plan")
- }
- if selection.Account == nil || selection.Account.ID != 1 {
- t.Fatalf("expected account 1")
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_StickyCapacitySpilloverKeepsBinding(t *testing.T) {
- sessionHash := "sticky-spillover"
- groupID := int64(1)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 6, Priority: 1, GroupIDs: []int64{groupID}},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 6, Priority: 1, GroupIDs: []int64{groupID}},
- },
- }
- cache := &stubGatewayCache{
- sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
- }
- concurrencyCache := stubConcurrencyCache{
- acquireResults: map[int64]bool{1: false, 2: true},
- waitCounts: map[int64]int{1: 1},
- loadMap: map[int64]*AccountLoadInfo{
- 1: {AccountID: 1, LoadRate: 100},
- 2: {AccountID: 2, LoadRate: 10},
- },
- }
- cfg := &config.Config{RunMode: config.RunModeStandard}
- cfg.Gateway.Scheduling.LoadBatchEnabled = true
- cfg.Gateway.Scheduling.StickySessionMaxWaiting = 1
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- cfg: cfg,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(2), selection.Account.ID, "capacity spillover should use the other account for this request")
- require.True(t, selection.Acquired)
- require.Equal(t, int64(1), cache.sessionBindings["openai:"+sessionHash], "capacity spillover must not migrate the durable sticky binding")
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_PrefersLowerLoad(t *testing.T) {
- groupID := int64(1)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- },
- }
- cache := &stubGatewayCache{}
- concurrencyCache := stubConcurrencyCache{
- loadMap: map[int64]*AccountLoadInfo{
- 1: {AccountID: 1, LoadRate: 80},
- 2: {AccountID: 2, LoadRate: 10},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "load", "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
- t.Fatalf("expected account 2")
- }
- if cache.sessionBindings["openai:load"] != 2 {
- t.Fatalf("expected sticky session updated")
- }
- }
- func TestOpenAISelectAccountForModelWithExclusions_StickyExcludedFallback(t *testing.T) {
- sessionHash := "excluded"
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 2},
- },
- }
- cache := &stubGatewayCache{
- sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- }
- excluded := map[int64]struct{}{1: {}}
- acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", excluded)
- if err != nil {
- t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
- }
- if acc == nil || acc.ID != 2 {
- t.Fatalf("expected account 2")
- }
- }
- func TestOpenAISelectAccountForModelWithExclusions_StickyNonOpenAI(t *testing.T) {
- sessionHash := "non-openai"
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformAnthropic, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 2},
- },
- }
- cache := &stubGatewayCache{
- sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- }
- acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
- }
- if acc == nil || acc.ID != 2 {
- t.Fatalf("expected account 2")
- }
- }
- func TestOpenAISelectAccountForModelWithExclusions_NoAccounts(t *testing.T) {
- repo := stubOpenAIAccountRepo{accounts: []Account{}}
- cache := &stubGatewayCache{}
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- }
- acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, "", "", nil)
- if err == nil {
- t.Fatalf("expected error for no accounts")
- }
- if acc != nil {
- t.Fatalf("expected nil account")
- }
- if !strings.Contains(err.Error(), "no available OpenAI accounts") {
- t.Fatalf("unexpected error: %v", err)
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_NoCandidates(t *testing.T) {
- groupID := int64(1)
- resetAt := time.Now().Add(1 * time.Hour)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, RateLimitResetAt: &resetAt},
- },
- }
- cache := &stubGatewayCache{}
- concurrencyCache := stubConcurrencyCache{}
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
- if err == nil {
- t.Fatalf("expected error for no candidates")
- }
- if selection != nil {
- t.Fatalf("expected nil selection")
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_AllFullWaitPlan(t *testing.T) {
- groupID := int64(1)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- },
- }
- cache := &stubGatewayCache{}
- concurrencyCache := stubConcurrencyCache{
- loadMap: map[int64]*AccountLoadInfo{
- 1: {AccountID: 1, LoadRate: 100},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.WaitPlan == nil {
- t.Fatalf("expected wait plan")
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_LoadBatchErrorNoAcquire(t *testing.T) {
- groupID := int64(1)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- },
- }
- cache := &stubGatewayCache{}
- concurrencyCache := stubConcurrencyCache{
- loadBatchErr: errors.New("load batch failed"),
- acquireResults: map[int64]bool{1: false},
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.WaitPlan == nil {
- t.Fatalf("expected wait plan")
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_MissingLoadInfo(t *testing.T) {
- groupID := int64(1)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- },
- }
- cache := &stubGatewayCache{}
- concurrencyCache := stubConcurrencyCache{
- loadMap: map[int64]*AccountLoadInfo{
- 1: {AccountID: 1, LoadRate: 50},
- },
- skipDefaultLoad: true,
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
- t.Fatalf("expected account 2")
- }
- }
- func TestOpenAISelectAccountForModelWithExclusions_LeastRecentlyUsed(t *testing.T) {
- oldTime := time.Now().Add(-2 * time.Hour)
- newTime := time.Now().Add(-1 * time.Hour)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Priority: 1, LastUsedAt: &newTime},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Priority: 1, LastUsedAt: &oldTime},
- },
- }
- cache := &stubGatewayCache{}
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- }
- acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, "", "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
- }
- if acc == nil || acc.ID != 2 {
- t.Fatalf("expected account 2")
- }
- }
- func TestOpenAISelectAccountWithLoadAwareness_PreferNeverUsed(t *testing.T) {
- groupID := int64(1)
- lastUsed := time.Now().Add(-1 * time.Hour)
- repo := stubOpenAIAccountRepo{
- accounts: []Account{
- {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, LastUsedAt: &lastUsed},
- {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
- },
- }
- cache := &stubGatewayCache{}
- concurrencyCache := stubConcurrencyCache{
- loadMap: map[int64]*AccountLoadInfo{
- 1: {AccountID: 1, LoadRate: 10},
- 2: {AccountID: 2, LoadRate: 10},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: repo,
- cache: cache,
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
- if err != nil {
- t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
- }
- if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
- t.Fatalf("expected account 2")
- }
- }
- func TestOpenAIStreamingTimeout(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 1,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{},
- }
- start := time.Now()
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, start, "model", "model")
- _ = pw.Close()
- _ = pr.Close()
- if err == nil || !strings.Contains(err.Error(), "stream data interval timeout") {
- t.Fatalf("expected stream timeout error, got %v", err)
- }
- if !strings.Contains(rec.Body.String(), "\"type\":\"error\"") || !strings.Contains(rec.Body.String(), "stream_timeout") {
- t.Fatalf("expected OpenAI-compatible error SSE event, got %q", rec.Body.String())
- }
- }
- func TestOpenAIStreamingContextCanceledReturnsIncompleteErrorWithoutInjectingErrorEvent(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- ctx, cancel := context.WithCancel(context.Background())
- cancel()
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: cancelReadCloser{},
- Header: http.Header{},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
- if err == nil || !strings.Contains(err.Error(), "stream usage incomplete") {
- t.Fatalf("expected incomplete stream error, got %v", err)
- }
- if strings.Contains(rec.Body.String(), "event: error") || strings.Contains(rec.Body.String(), "stream_read_error") {
- t.Fatalf("expected no injected SSE error event, got %q", rec.Body.String())
- }
- }
- func TestOpenAIStreamingReadErrorBeforeOutputReturnsFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: errReadCloser{err: io.ErrUnexpectedEOF},
- Header: http.Header{"X-Request-Id": []string{"rid-disconnect"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
- require.False(t, c.Writer.Written())
- require.Empty(t, rec.Body.String())
- }
- func TestOpenAIStreamingPostOutputDisconnectQuarantinesSharedProxyWithoutSameStreamFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- proxyID := int64(4698)
- account := &Account{
- ID: 469801,
- Name: "oauth-on-shared-proxy",
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- ProxyID: &proxyID,
- }
- svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{
- MaxLineSize: defaultMaxLineSize,
- }}}
- // collapseInterval 0: the two loop iterations below record within the
- // production collapse window and must count as distinct failure events here.
- svc.openaiProxyStreamCircuit = newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
- failureThreshold: 2,
- failureWindow: time.Minute,
- quarantineTTL: 10 * time.Minute,
- maxEntries: 16,
- })
- for _, readErr := range []error{
- io.ErrUnexpectedEOF,
- errors.New("http2: client connection lost"),
- } {
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: &openAIStreamReadThenErrorCloser{
- reader: strings.NewReader(strings.Join([]string{
- "event: response.output_text.delta",
- `data: {"type":"response.output_text.delta","delta":"partial"}`,
- "",
- }, "\n")),
- err: readErr,
- },
- Header: http.Header{"X-Request-Id": []string{"rid-proxy-disconnect"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "gpt-5.6-sol", "gpt-5.6-sol")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr), "post-output disconnect must not fail over inside the same stream")
- require.Contains(t, rec.Body.String(), "partial")
- }
- scheduler := &defaultOpenAIAccountScheduler{service: svc}
- compatible, reason := scheduler.isAccountRequestCompatibleReason(context.Background(), account, OpenAIAccountScheduleRequest{})
- require.False(t, compatible, "the next request must exclude accounts sharing the quarantined proxy")
- require.Equal(t, "proxy_stream_quarantined", reason)
- }
- func TestOpenAIStreamingTerminalAndClientCancellationDoNotQuarantineProxy(t *testing.T) {
- gin.SetMode(gin.TestMode)
- proxyID := int64(4699)
- account := &Account{ID: 469901, Name: "oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, ProxyID: &proxyID}
- svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}}
- terminalRecorder := httptest.NewRecorder()
- terminalCtx, _ := gin.CreateTestContext(terminalRecorder)
- terminalCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
- terminalResp := &http.Response{
- StatusCode: http.StatusOK,
- Body: &openAIStreamReadThenErrorCloser{
- reader: strings.NewReader(strings.Join([]string{
- "event: response.completed",
- `data: {"type":"response.completed","response":{"status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":3,"total_tokens":8}}}`,
- "",
- }, "\n")),
- err: io.ErrUnexpectedEOF,
- },
- Header: http.Header{},
- }
- _, err := svc.handleStreamingResponse(terminalCtx.Request.Context(), terminalResp, terminalCtx, account, time.Now(), "model", "model")
- require.NoError(t, err)
- for range 2 {
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- ctx, cancel := context.WithCancel(context.Background())
- cancel()
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(ctx)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: &openAIStreamReadThenErrorCloser{
- reader: strings.NewReader("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"),
- err: context.Canceled,
- },
- Header: http.Header{},
- }
- _, err = svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
- require.Error(t, err)
- }
- scheduler := &defaultOpenAIAccountScheduler{service: svc}
- compatible, reason := scheduler.isAccountRequestCompatibleReason(context.Background(), account, OpenAIAccountScheduleRequest{})
- require.True(t, compatible)
- require.Empty(t, reason)
- }
- func TestOpenAIStreamingResponseFailedBeforeOutputReturnsFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.in_progress",
- `data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","error":{"message":"An error occurred while processing your request."}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-failed"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
- require.False(t, failoverErr.RetryableOnSameAccount)
- require.Contains(t, string(failoverErr.ResponseBody), "An error occurred while processing your request")
- require.False(t, c.Writer.Written())
- require.Empty(t, rec.Body.String())
- }
- func TestOpenAIStreamingResponseFailedBeforeOutputCapacityErrorReturnsFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.in_progress",
- `data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","error":{"message":"Selected model is at capacity. Please try a different model.","type":"invalid_request_error"}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-capacity-failed"}},
- }
- account := &Account{
- ID: 1,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Name: "pool-account",
- Credentials: map[string]any{
- "pool_mode": true,
- },
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
- require.True(t, failoverErr.RetryableOnSameAccount)
- require.Contains(t, string(failoverErr.ResponseBody), "Selected model is at capacity")
- require.False(t, c.Writer.Written())
- require.Empty(t, rec.Body.String())
- }
- func TestOpenAIStreamingResponseFailedBeforeOutputServerOverloadedCodeReturnsFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","response":{"id":"resp_1","error":{"code":"server_is_overloaded","message":"Please retry later."}}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-overloaded-failed"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusServiceUnavailable, failoverErr.StatusCode)
- require.Contains(t, string(failoverErr.ResponseBody), "Please retry later")
- // 容量降载是请求级信号:非池模式账号也要先在同账号重试,且不得据此临时封禁账号。
- // 否则单个被降载的请求会把整池账号逐个消耗掉,而降载因素在每个账号上都相同。
- require.True(t, failoverErr.RetryableOnSameAccount)
- require.True(t, failoverErr.RequestScopedTransient)
- require.False(t, c.Writer.Written())
- require.Empty(t, rec.Body.String())
- }
- func TestOpenAIStreamingResponseFailedBeforeOutputRateLimitUsesPoolRetryPolicy(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
- "",
- }, "\n"))),
- Header: http.Header{
- "X-Request-Id": []string{"rid-rate-limit-failed"},
- "Retry-After": []string{"1"},
- },
- }
- account := &Account{
- ID: 1,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Name: "pool-account",
- Credentials: map[string]any{
- "pool_mode": true,
- "pool_mode_retry_count": float64(1),
- "pool_mode_retry_status_codes": []any{float64(http.StatusTooManyRequests)},
- },
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
- require.True(t, failoverErr.RetryableOnSameAccount)
- require.Equal(t, "1", failoverErr.ResponseHeaders.Get("Retry-After"))
- require.Equal(t, "rate_limit_error", gjson.GetBytes(failoverErr.ResponseBody, "error.type").String())
- require.Contains(t, string(failoverErr.ResponseBody), "Concurrency limit exceeded")
- require.False(t, c.Writer.Written())
- require.Empty(t, rec.Body.String())
- opsVal, ok := c.Get(OpsUpstreamErrorsKey)
- require.True(t, ok)
- opsEvents, ok := opsVal.([]*OpsUpstreamErrorEvent)
- require.True(t, ok)
- require.NotEmpty(t, opsEvents)
- require.Equal(t, http.StatusTooManyRequests, opsEvents[len(opsEvents)-1].UpstreamStatusCode)
- }
- // 流内 rate limit 进入 OAuth 同账号重试窗口,但不立即写账号级限流/封禁状态:
- // HTTP 200 流的 x-codex-* 头不能让窗口内的账号提前失去调度资格。
- func TestOpenAIStreamingResponseFailedRateLimitDoesNotBlockAccountScheduling(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
- "",
- }, "\n"))),
- Header: http.Header{
- "X-Codex-Primary-Used-Percent": []string{"12"},
- "X-Codex-Primary-Reset-After-Seconds": []string{"604800"},
- "Retry-After": []string{"1"},
- },
- }
- account := &Account{
- ID: 11,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Name: "oauth-account",
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
- require.True(t, failoverErr.RetryableOnSameAccount)
- require.False(t, failoverErr.SameAccountRetryDeadline.IsZero())
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestOpenAIStreamingResponseFailedAfterOutputSanitizesVerboseResponseForClient(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- longInstructions := strings.Repeat("You are GPT-5.1 running in the Codex CLI. ", 20)
- failedPayload := fmt.Sprintf(
- `{"type":"response.failed","response":{"id":"resp_failed","object":"response","created_at":1782446336,"status":"failed","instructions":%q,"output":[{"type":"message","content":[{"type":"output_text","text":"large"}]}],"usage":{"input_tokens":123,"output_tokens":0},"error":{"code":"context_length_exceeded","message":"Your input exceeds the context window of this model. Please adjust your input and try again."}}}`,
- longInstructions,
- )
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_failed"}}`,
- "",
- "event: response.output_text.delta",
- `data: {"type":"response.output_text.delta","delta":"partial"}`,
- "",
- "event: response.failed",
- "data: " + failedPayload,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-failed-after-output"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.Error(t, err)
- body := rec.Body.String()
- require.Contains(t, body, "event: response.failed")
- require.Contains(t, body, "context_length_exceeded")
- require.Contains(t, body, `"type":"invalid_request_error"`)
- require.Contains(t, body, "Your input exceeds the context window")
- require.NotContains(t, body, "You are GPT-5.1 running in the Codex CLI")
- require.NotContains(t, body, `"instructions"`)
- require.NotContains(t, body, `"output"`)
- require.NotContains(t, body, `"usage"`)
- }
- func TestOpenAIStreamingContextWindowResponseFailedBeforeOutputPassesThrough(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","error":{"type":"upstream_error","message":"Your input exceeds the context window of this model. Please adjust your input and try again.","code":null}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-context-window-failed"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr))
- require.True(t, c.Writer.Written())
- require.Contains(t, rec.Body.String(), "response.failed")
- require.Contains(t, rec.Body.String(), `"type":"upstream_error"`)
- require.Contains(t, rec.Body.String(), "Your input exceeds the context window")
- }
- func TestOpenAIStreamingContextWindowResponseFailedBeforeOutputAppliesPassthroughRule(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- rule := newNonFailoverPassthroughRule(http.StatusBadRequest, "context_length_exceeded", http.StatusBadRequest, "")
- rule.Platforms = []string{PlatformOpenAI}
- rule.PassthroughBody = true
- rule.CustomMessage = nil
- ruleSvc := &ErrorPassthroughService{}
- ruleSvc.setLocalCache([]*model.ErrorPassthroughRule{rule})
- BindErrorPassthroughService(c, ruleSvc)
- upstreamMessage := "Your input exceeds the context window of this model. Please adjust your input and try again."
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","response":{"id":"resp_1","error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"` + upstreamMessage + `"}}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-context-window-passthrough-rule"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr))
- require.True(t, IsResponseCommitted(c))
- require.Equal(t, http.StatusBadRequest, rec.Code)
- body := rec.Body.String()
- require.Equal(t, "upstream_error", gjson.Get(body, "error.type").String())
- require.Equal(t, upstreamMessage, gjson.Get(body, "error.message").String())
- require.NotContains(t, body, "response.failed")
- require.NotContains(t, body, "Upstream request failed")
- // 命中透传规则也应记录 ops 上游错误事件(对齐 CC/Messages 与 antigravity 先例)。
- opsVal, opsRecorded := c.Get(OpsUpstreamErrorsKey)
- require.True(t, opsRecorded, "passthrough hit should record an ops upstream error event")
- opsEvents, _ := opsVal.([]*OpsUpstreamErrorEvent)
- require.NotEmpty(t, opsEvents)
- }
- func TestOpenAIStreamingPreambleOnlyMissingTerminalReturnsFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.in_progress",
- `data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-missing-terminal"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.False(t, c.Writer.Written())
- require.Empty(t, rec.Body.String())
- }
- func TestOpenAIStreamingPreambleKeepaliveUsesDownstreamIdle(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- // Keepalive is based on *downstream* idle time (last flush to client),
- // not upstream event cadence. Interval is seconds (config unit).
- StreamKeepaliveInterval: 1,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{},
- }
- go func() {
- defer func() { _ = pw.Close() }()
- // Emit preamble/progress quickly so clientOutputStarted is true, then
- // leave a real downstream idle gap longer than keepaliveInterval so the
- // ticker can write ":\n\n". Frequent upstream ticks used to refresh
- // lastDownstreamWriteAt and flake on loaded CI runners.
- _, _ = pw.Write([]byte("data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\n"))
- time.Sleep(50 * time.Millisecond)
- _, _ = pw.Write([]byte("data: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_1\"}}\n\n"))
- time.Sleep(1300 * time.Millisecond)
- _, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}}\n\n"))
- }()
- result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- _ = pr.Close()
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Contains(t, rec.Body.String(), ":\n\n")
- require.Contains(t, rec.Body.String(), "response.completed")
- }
- func TestOpenAIStreamingNormalizesTerminalOutputFromDeltas(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- `data: {"type":"response.created","response":{"id":"resp_sdk_parse"}}`,
- "",
- `data: {"type":"response.output_text.delta","delta":"pon"}`,
- "",
- `data: {"type":"response.output_text.delta","delta":"g"}`,
- "",
- `data: {"type":"response.completed","response":{"id":"resp_sdk_parse","status":"completed","output":null,"usage":{"input_tokens":1,"output_tokens":1}}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-sdk-parse"}},
- }
- result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.NoError(t, err)
- require.NotNil(t, result)
- terminalType, terminalPayload, ok := extractOpenAISSETerminalEvent(rec.Body.String())
- require.True(t, ok)
- require.Equal(t, "response.completed", terminalType)
- output := gjson.GetBytes(terminalPayload, "response.output")
- require.True(t, output.IsArray())
- require.Len(t, output.Array(), 1)
- require.Equal(t, "pong", gjson.GetBytes(terminalPayload, "response.output.0.content.0.text").String())
- }
- func TestOpenAIStreamingNormalizesTerminalOutputToEmptyArray(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- `data: {"type":"response.completed","response":{"id":"resp_empty","status":"completed","output":null,"usage":{"input_tokens":1,"output_tokens":0}}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-empty-output"}},
- }
- result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.NoError(t, err)
- require.NotNil(t, result)
- terminalType, terminalPayload, ok := extractOpenAISSETerminalEvent(rec.Body.String())
- require.True(t, ok)
- require.Equal(t, "response.completed", terminalType)
- output := gjson.GetBytes(terminalPayload, "response.output")
- require.True(t, output.IsArray())
- require.Len(t, output.Array(), 0)
- }
- func TestOpenAIStreamingPolicyResponseFailedBeforeOutputPassesThrough(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","error":{"type":"safety_error","message":"This request has been flagged for potentially high-risk cyber activity."}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-policy-failed"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr))
- require.True(t, c.Writer.Written())
- require.Contains(t, rec.Body.String(), "response.failed")
- require.Contains(t, rec.Body.String(), "high-risk cyber activity")
- }
- func TestOpenAIStreamingClientDisconnectDrainsUpstreamUsage(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- c.Writer = &failingGinWriter{ResponseWriter: c.Writer, failAfter: 0}
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{},
- }
- go func() {
- defer func() { _ = pw.Close() }()
- _, _ = pw.Write([]byte("data: {\"type\":\"response.in_progress\",\"response\":{}}\n\n"))
- _, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n"))
- }()
- result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
- _ = pr.Close()
- if err != nil {
- t.Fatalf("expected nil error, got %v", err)
- }
- if result == nil || result.usage == nil {
- t.Fatalf("expected usage result")
- }
- if result.usage.InputTokens != 3 || result.usage.OutputTokens != 5 || result.usage.CacheReadInputTokens != 1 {
- t.Fatalf("unexpected usage: %+v", *result.usage)
- }
- if strings.Contains(rec.Body.String(), "event: error") || strings.Contains(rec.Body.String(), "write_failed") {
- t.Fatalf("expected no injected SSE error event, got %q", rec.Body.String())
- }
- }
- func TestOpenAIStreamingMissingTerminalEventReturnsIncompleteError(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{},
- }
- go func() {
- defer func() { _ = pw.Close() }()
- _, _ = pw.Write([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\",\"output_index\":0}\n\n"))
- }()
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
- _ = pr.Close()
- if err == nil || !strings.Contains(err.Error(), "missing terminal event") {
- t.Fatalf("expected missing terminal event error, got %v", err)
- }
- }
- func TestOpenAIStreamingPassthroughMissingTerminalEventReturnsIncompleteError(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{},
- }
- go func() {
- defer func() { _ = pw.Close() }()
- _, _ = pw.Write([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\",\"output_index\":0}\n\n"))
- }()
- _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "", "")
- _ = pr.Close()
- if err == nil || !strings.Contains(err.Error(), "missing terminal event") {
- t.Fatalf("expected missing terminal event error, got %v", err)
- }
- }
- func TestOpenAIStreamingPassthroughPostOutputDisconnectQuarantinesSharedProxy(t *testing.T) {
- gin.SetMode(gin.TestMode)
- proxyID := int64(4698)
- account := &Account{ID: 469804, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, ProxyID: &proxyID}
- svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}}
- // collapseInterval 0: the loop below records within the production collapse
- // window and must count as distinct failure events here.
- svc.openaiProxyStreamCircuit = newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
- failureThreshold: 2,
- failureWindow: time.Minute,
- quarantineTTL: 10 * time.Minute,
- maxEntries: 16,
- })
- for _, readErr := range []error{io.ErrUnexpectedEOF, errors.New("http2: client connection lost")} {
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: &openAIStreamReadThenErrorCloser{
- reader: strings.NewReader("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"),
- err: readErr,
- },
- Header: http.Header{"X-Request-Id": []string{"rid-passthrough-proxy-disconnect"}},
- }
- _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr), "post-output disconnect must not fail over inside the same stream")
- require.Contains(t, rec.Body.String(), "partial")
- }
- require.True(t, svc.isOpenAIProxyStreamQuarantined(context.Background(), account))
- }
- func TestOpenAIStreamingPassthroughResponseFailedBeforeOutputReturnsFailover(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","error":{"message":"upstream processing failed"}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-passthrough-failed"}},
- }
- _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.ErrorAs(t, err, &failoverErr)
- require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
- require.Contains(t, string(failoverErr.ResponseBody), "upstream processing failed")
- require.False(t, c.Writer.Written())
- require.Empty(t, rec.Body.String())
- }
- func TestOpenAIStreamingPassthroughContextWindowResponseFailedBeforeOutputAppliesPassthroughRule(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- rule := newNonFailoverPassthroughRule(http.StatusBadRequest, "input exceeds the context window", http.StatusBadRequest, "")
- rule.Platforms = []string{PlatformOpenAI}
- rule.PassthroughBody = true
- rule.CustomMessage = nil
- ruleSvc := &ErrorPassthroughService{}
- ruleSvc.setLocalCache([]*model.ErrorPassthroughRule{rule})
- BindErrorPassthroughService(c, ruleSvc)
- upstreamMessage := "Your input exceeds the context window of this model. Please adjust your input and try again."
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","response":{"id":"resp_1","error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"` + upstreamMessage + `"}}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-pass-context-window-passthrough-rule"}},
- }
- _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr))
- require.True(t, IsResponseCommitted(c))
- require.Equal(t, http.StatusBadRequest, rec.Code)
- body := rec.Body.String()
- require.Equal(t, "upstream_error", gjson.Get(body, "error.type").String())
- require.Equal(t, upstreamMessage, gjson.Get(body, "error.message").String())
- require.NotContains(t, body, "response.failed")
- require.NotContains(t, body, "Upstream request failed")
- // 命中透传规则也应记录 ops 上游错误事件(对齐 CC/Messages 与 antigravity 先例)。
- opsVal, opsRecorded := c.Get(OpsUpstreamErrorsKey)
- require.True(t, opsRecorded, "passthrough hit should record an ops upstream error event")
- opsEvents, _ := opsVal.([]*OpsUpstreamErrorEvent)
- require.NotEmpty(t, opsEvents)
- }
- func TestOpenAIStreamingPassthroughContextWindowResponseFailedBeforeOutputWithoutRulePassesThrough(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","response":{"id":"resp_1","error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"Your input exceeds the context window of this model. Please adjust your input and try again."}}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-pass-context-window-no-rule"}},
- }
- _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr))
- body := rec.Body.String()
- require.Contains(t, body, "event: response.failed")
- require.Contains(t, body, "context_length_exceeded")
- require.Contains(t, body, "Your input exceeds the context window")
- }
- func TestOpenAIStreamingPassthroughResponseFailedAfterOutputSanitizesVerboseResponseForClient(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- longInstructions := strings.Repeat("You are GPT-5.1 running in the Codex CLI. ", 20)
- failedPayload := fmt.Sprintf(
- `{"type":"response.failed","response":{"id":"resp_pass_failed","object":"response","created_at":1782446336,"status":"failed","instructions":%q,"output":[{"type":"message","content":[{"type":"output_text","text":"large"}]}],"usage":{"input_tokens":123,"output_tokens":0},"error":{"code":"context_length_exceeded","message":"Your input exceeds the context window of this model. Please adjust your input and try again."}}}`,
- longInstructions,
- )
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"resp_pass_failed"}}`,
- "",
- "event: response.output_text.delta",
- `data: {"type":"response.output_text.delta","delta":"partial"}`,
- "",
- "event: response.failed",
- "data: " + failedPayload,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-pass-failed-after-output"}},
- }
- _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
- require.Error(t, err)
- body := rec.Body.String()
- require.Contains(t, body, "event: response.failed")
- require.Contains(t, body, "context_length_exceeded")
- require.Contains(t, body, `"type":"invalid_request_error"`)
- require.Contains(t, body, "Your input exceeds the context window")
- require.NotContains(t, body, "You are GPT-5.1 running in the Codex CLI")
- require.NotContains(t, body, `"instructions"`)
- require.NotContains(t, body, `"output"`)
- require.NotContains(t, body, `"usage"`)
- }
- func TestOpenAIStreamingPassthroughResponseDoneWithoutDoneMarkerStillSucceeds(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{},
- }
- go func() {
- defer func() { _ = pw.Close() }()
- _, _ = pw.Write([]byte("data: {\"type\":\"response.done\",\"response\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n"))
- }()
- result, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "", "")
- _ = pr.Close()
- require.NoError(t, err)
- require.NotNil(t, result)
- require.NotNil(t, result.usage)
- require.Equal(t, 2, result.usage.InputTokens)
- require.Equal(t, 3, result.usage.OutputTokens)
- require.Equal(t, 1, result.usage.CacheReadInputTokens)
- }
- func TestOpenAIStreamingPassthroughResponseIncompleteWithoutDoneMarkerStillSucceeds(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{},
- }
- go func() {
- defer func() { _ = pw.Close() }()
- _, _ = pw.Write([]byte("data: {\"type\":\"response.incomplete\",\"response\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n"))
- }()
- result, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "", "")
- _ = pr.Close()
- require.NoError(t, err)
- require.NotNil(t, result)
- require.NotNil(t, result.usage)
- require.Equal(t, 2, result.usage.InputTokens)
- require.Equal(t, 3, result.usage.OutputTokens)
- require.Equal(t, 1, result.usage.CacheReadInputTokens)
- }
- func TestOpenAIStreamingTooLong(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: 64 * 1024,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{},
- }
- go func() {
- defer func() { _ = pw.Close() }()
- // 写入超过 MaxLineSize 的单行数据,触发 ErrTooLong
- payload := "data: " + strings.Repeat("a", 128*1024) + "\n"
- _, _ = pw.Write([]byte(payload))
- }()
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 2}, time.Now(), "model", "model")
- _ = pr.Close()
- if !errors.Is(err, bufio.ErrTooLong) {
- t.Fatalf("expected ErrTooLong, got %v", err)
- }
- if !strings.Contains(rec.Body.String(), "\"type\":\"error\"") || !strings.Contains(rec.Body.String(), "response_too_large") {
- t.Fatalf("expected OpenAI-compatible error SSE event, got %q", rec.Body.String())
- }
- }
- func TestOpenAINonStreamingContentTypePassThrough(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Security: config.SecurityConfig{
- ResponseHeaders: config.ResponseHeaderConfig{Enabled: false},
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- body := []byte(`{"usage":{"input_tokens":1,"output_tokens":2,"input_tokens_details":{"cached_tokens":0}}}`)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(bytes.NewReader(body)),
- Header: http.Header{"Content-Type": []string{"application/vnd.test+json"}},
- }
- _, err := svc.handleNonStreamingResponse(c.Request.Context(), resp, c, &Account{}, "model", "model")
- if err != nil {
- t.Fatalf("handleNonStreamingResponse error: %v", err)
- }
- if !strings.Contains(rec.Header().Get("Content-Type"), "application/vnd.test+json") {
- t.Fatalf("expected Content-Type passthrough, got %q", rec.Header().Get("Content-Type"))
- }
- }
- func TestOpenAINonStreamingContentTypeDefault(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Security: config.SecurityConfig{
- ResponseHeaders: config.ResponseHeaderConfig{Enabled: false},
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- body := []byte(`{"usage":{"input_tokens":1,"output_tokens":2,"input_tokens_details":{"cached_tokens":0}}}`)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(bytes.NewReader(body)),
- Header: http.Header{},
- }
- _, err := svc.handleNonStreamingResponse(c.Request.Context(), resp, c, &Account{}, "model", "model")
- if err != nil {
- t.Fatalf("handleNonStreamingResponse error: %v", err)
- }
- if !strings.Contains(rec.Header().Get("Content-Type"), "application/json") {
- t.Fatalf("expected default Content-Type, got %q", rec.Header().Get("Content-Type"))
- }
- }
- func TestOpenAIStreamingHeadersOverride(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Security: config.SecurityConfig{
- ResponseHeaders: config.ResponseHeaderConfig{Enabled: false},
- },
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{
- "Cache-Control": []string{"upstream"},
- "X-Request-Id": []string{"req-123"},
- "Content-Type": []string{"application/custom"},
- },
- }
- go func() {
- defer func() { _ = pw.Close() }()
- _, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{}}\n\n"))
- }()
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
- _ = pr.Close()
- if err != nil {
- t.Fatalf("handleStreamingResponse error: %v", err)
- }
- if rec.Header().Get("Cache-Control") != "no-cache" {
- t.Fatalf("expected Cache-Control override, got %q", rec.Header().Get("Cache-Control"))
- }
- if rec.Header().Get("Content-Type") != "text/event-stream" {
- t.Fatalf("expected Content-Type override, got %q", rec.Header().Get("Content-Type"))
- }
- if rec.Header().Get("X-Request-Id") != "req-123" {
- t.Fatalf("expected X-Request-Id passthrough, got %q", rec.Header().Get("X-Request-Id"))
- }
- }
- func TestOpenAIStreamingReuseScannerBufferAndStillWorks(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Gateway: config.GatewayConfig{
- StreamDataIntervalTimeout: 0,
- StreamKeepaliveInterval: 0,
- MaxLineSize: defaultMaxLineSize,
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- pr, pw := io.Pipe()
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: pr,
- Header: http.Header{},
- }
- go func() {
- defer func() { _ = pw.Close() }()
- _, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"input_tokens_details\":{\"cached_tokens\":3}}}}\n\n"))
- }()
- result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
- _ = pr.Close()
- require.NoError(t, err)
- require.NotNil(t, result)
- require.NotNil(t, result.usage)
- require.Equal(t, 1, result.usage.InputTokens)
- require.Equal(t, 2, result.usage.OutputTokens)
- require.Equal(t, 3, result.usage.CacheReadInputTokens)
- }
- func TestOpenAIInvalidBaseURLWhenAllowlistDisabled(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{
- Security: config.SecurityConfig{
- URLAllowlist: config.URLAllowlistConfig{Enabled: false},
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- account := &Account{
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Credentials: map[string]any{"base_url": "://invalid-url"},
- }
- _, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte("{}"), "token", false, "", false)
- if err == nil {
- t.Fatalf("expected error for invalid base_url when allowlist disabled")
- }
- }
- func TestOpenAIValidateUpstreamBaseURLDisabledRequiresHTTPS(t *testing.T) {
- cfg := &config.Config{
- Security: config.SecurityConfig{
- URLAllowlist: config.URLAllowlistConfig{Enabled: false},
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- if _, err := svc.validateUpstreamBaseURL("http://not-https.example.com"); err == nil {
- t.Fatalf("expected http to be rejected when allow_insecure_http is false")
- }
- normalized, err := svc.validateUpstreamBaseURL("https://example.com")
- if err != nil {
- t.Fatalf("expected https to be allowed when allowlist disabled, got %v", err)
- }
- if normalized != "https://example.com" {
- t.Fatalf("expected raw url passthrough, got %q", normalized)
- }
- }
- func TestOpenAIValidateUpstreamBaseURLDisabledAllowsHTTP(t *testing.T) {
- cfg := &config.Config{
- Security: config.SecurityConfig{
- URLAllowlist: config.URLAllowlistConfig{
- Enabled: false,
- AllowInsecureHTTP: true,
- },
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- normalized, err := svc.validateUpstreamBaseURL("http://not-https.example.com")
- if err != nil {
- t.Fatalf("expected http allowed when allow_insecure_http is true, got %v", err)
- }
- if normalized != "http://not-https.example.com" {
- t.Fatalf("expected raw url passthrough, got %q", normalized)
- }
- }
- func TestOpenAIValidateUpstreamBaseURLEnabledEnforcesAllowlist(t *testing.T) {
- cfg := &config.Config{
- Security: config.SecurityConfig{
- URLAllowlist: config.URLAllowlistConfig{
- Enabled: true,
- UpstreamHosts: []string{"example.com"},
- },
- },
- }
- svc := &OpenAIGatewayService{cfg: cfg}
- if _, err := svc.validateUpstreamBaseURL("https://example.com"); err != nil {
- t.Fatalf("expected allowlisted host to pass, got %v", err)
- }
- if _, err := svc.validateUpstreamBaseURL("https://evil.com"); err == nil {
- t.Fatalf("expected non-allowlisted host to fail")
- }
- }
- func TestOpenAIUpdateCodexUsageSnapshotFromHeaders(t *testing.T) {
- repo := &snapshotUpdateAccountRepo{updateExtraCalls: make(chan map[string]any, 1)}
- svc := &OpenAIGatewayService{accountRepo: repo}
- headers := http.Header{}
- headers.Set("x-codex-primary-used-percent", "12")
- headers.Set("x-codex-secondary-used-percent", "34")
- headers.Set("x-codex-primary-window-minutes", "300")
- headers.Set("x-codex-secondary-window-minutes", "10080")
- headers.Set("x-codex-primary-reset-after-seconds", "600")
- headers.Set("x-codex-secondary-reset-after-seconds", "86400")
- svc.UpdateCodexUsageSnapshotFromHeaders(context.Background(), 123, headers)
- select {
- case updates := <-repo.updateExtraCalls:
- require.Equal(t, 12.0, updates["codex_5h_used_percent"])
- require.Equal(t, 34.0, updates["codex_7d_used_percent"])
- require.Equal(t, 600, updates["codex_5h_reset_after_seconds"])
- require.Equal(t, 86400, updates["codex_7d_reset_after_seconds"])
- case <-time.After(2 * time.Second):
- t.Fatal("expected UpdateExtra to be called")
- }
- }
- func TestOpenAIResponsesRequestPathSuffix(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- tests := []struct {
- name string
- path string
- want string
- }{
- {name: "exact v1 responses", path: "/v1/responses", want: ""},
- {name: "compact v1 responses", path: "/v1/responses/compact", want: "/compact"},
- {name: "compact alias responses", path: "/responses/compact/", want: "/compact"},
- {name: "nested suffix", path: "/openai/v1/responses/compact/detail", want: "/compact/detail"},
- {name: "unrelated path", path: "/v1/chat/completions", want: ""},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- c.Request = httptest.NewRequest(http.MethodPost, tt.path, nil)
- require.Equal(t, tt.want, openAIResponsesRequestPathSuffix(c))
- })
- }
- }
- func TestNormalizeOpenAICompactRequestBodyPreservesCurrentCodexPayloadFields(t *testing.T) {
- body := []byte(`{"model":"gpt-5.5","input":[{"type":"message","role":"user","content":"compact me"}],"instructions":"compact-test","tools":[{"type":"function","name":"shell"}],"parallel_tool_calls":true,"reasoning":{"effort":"high"},"text":{"verbosity":"low"},"previous_response_id":"resp_123","store":true,"stream":true,"prompt_cache_key":"cache_123"}`)
- normalized, changed, err := normalizeOpenAICompactRequestBody(body)
- require.NoError(t, err)
- require.True(t, changed)
- require.Equal(t, "gpt-5.5", gjson.GetBytes(normalized, "model").String())
- require.True(t, gjson.GetBytes(normalized, "tools").Exists())
- require.True(t, gjson.GetBytes(normalized, "parallel_tool_calls").Bool())
- require.Equal(t, "high", gjson.GetBytes(normalized, "reasoning.effort").String())
- require.Equal(t, "low", gjson.GetBytes(normalized, "text.verbosity").String())
- require.Equal(t, "resp_123", gjson.GetBytes(normalized, "previous_response_id").String())
- require.False(t, gjson.GetBytes(normalized, "store").Exists())
- require.False(t, gjson.GetBytes(normalized, "stream").Exists())
- require.False(t, gjson.GetBytes(normalized, "prompt_cache_key").Exists())
- }
- func TestNormalizeOpenAICompactRequestBodyDropsParallelToolCallsWithoutUsableTools(t *testing.T) {
- for _, body := range []string{
- `{"model":"gpt-5.5","input":[],"parallel_tool_calls":true}`,
- `{"model":"gpt-5.5","input":[],"tools":[],"parallel_tool_calls":false}`,
- `{"model":"gpt-5.5","input":[],"tools":null,"parallel_tool_calls":true}`,
- `{"model":"gpt-5.5","input":[],"tools":{},"parallel_tool_calls":true}`,
- } {
- normalized, changed, err := normalizeOpenAICompactRequestBody([]byte(body))
- require.NoError(t, err)
- require.True(t, changed)
- require.False(t, gjson.GetBytes(normalized, "parallel_tool_calls").Exists(), string(normalized))
- }
- withTools := []byte(`{"model":"gpt-5.5","tools":[{"type":"function","name":"lookup"}],"parallel_tool_calls":false}`)
- normalized, _, err := normalizeOpenAICompactRequestBody(withTools)
- require.NoError(t, err)
- require.True(t, gjson.GetBytes(normalized, "parallel_tool_calls").Exists())
- require.False(t, gjson.GetBytes(normalized, "parallel_tool_calls").Bool())
- }
- func TestOpenAIBuildUpstreamRequestOpenAIPassthroughPreservesCompactPath(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
- svc := &OpenAIGatewayService{}
- account := &Account{Type: AccountTypeOAuth}
- req, err := svc.buildUpstreamRequestOpenAIPassthrough(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token")
- require.NoError(t, err)
- require.Equal(t, chatgptCodexURL+"/compact", req.URL.String())
- require.Equal(t, "application/json", req.Header.Get("Accept"))
- require.Equal(t, codexCLIVersion, req.Header.Get("Version"))
- require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header")
- require.NotEmpty(t, req.Header.Get("Session_Id"))
- require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context()))
- }
- func TestOpenAIBuildUpstreamRequestOpenAIPassthroughPreservesExplicitAPIKeyBetaHeader(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
- c.Request.Header.Set("OpenAI-Beta", "api-key-specific-beta")
- svc := &OpenAIGatewayService{cfg: &config.Config{
- Security: config.SecurityConfig{
- URLAllowlist: config.URLAllowlistConfig{Enabled: false},
- },
- }}
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- req, err := svc.buildUpstreamRequestOpenAIPassthrough(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token")
- require.NoError(t, err)
- require.Equal(t, "api-key-specific-beta", req.Header.Get("OpenAI-Beta"), "OAuth-only backport must not alter API-key passthrough headers")
- }
- func TestOpenAIBuildUpstreamRequestCompactForcesJSONAcceptForOAuth(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
- svc := &OpenAIGatewayService{}
- account := &Account{
- Type: AccountTypeOAuth,
- Credentials: map[string]any{"chatgpt_account_id": "chatgpt-acc"},
- }
- req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", true)
- require.NoError(t, err)
- require.Equal(t, chatgptCodexURL+"/compact", req.URL.String())
- require.Equal(t, "application/json", req.Header.Get("Accept"))
- require.Equal(t, codexCLIVersion, req.Header.Get("Version"))
- require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header")
- require.NotEmpty(t, req.Header.Get("Session_Id"))
- require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context()))
- }
- func TestOpenAIBuildUpstreamRequestOAuthMessagesBridgeUsesSessionOnly(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- body := []byte(`{"model":"gpt-5.5","prompt_cache_key":"anthropic-metadata-session-1","input":[{"type":"message","role":"developer","content":[{"type":"input_text","text":"<sub2api-claude-code-todo-guard>"}]},{"type":"message","role":"user","content":"hello"}]}`)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
- c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
- c.Request.Header.Set("originator", "codex_cli_rs")
- svc := &OpenAIGatewayService{}
- account := &Account{
- Type: AccountTypeOAuth,
- Credentials: map[string]any{"chatgpt_account_id": "chatgpt-acc"},
- }
- req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, body, "token", true, "anthropic-metadata-session-1", false)
- require.NoError(t, err)
- require.NotEmpty(t, req.Header.Get("Session_Id"))
- require.Empty(t, req.Header.Get("Conversation_Id"))
- require.Empty(t, req.Header.Get("OpenAI-Beta"))
- require.Empty(t, req.Header.Get("originator"))
- }
- func TestOpenAIBuildUpstreamRequestPreservesCompactPathForAPIKeyBaseURL(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/responses/compact", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
- svc := &OpenAIGatewayService{cfg: &config.Config{
- Security: config.SecurityConfig{
- URLAllowlist: config.URLAllowlistConfig{Enabled: false},
- },
- }}
- account := &Account{
- Type: AccountTypeAPIKey,
- Platform: PlatformOpenAI,
- Credentials: map[string]any{"base_url": "https://example.com/v1"},
- }
- req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", false)
- require.NoError(t, err)
- require.Equal(t, "https://example.com/v1/responses/compact", req.URL.String())
- }
- func TestOpenAIBuildUpstreamRequestPreservesCodexIdentityHeaders(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
- c.Request.Header.Set("X-Codex-Window-ID", "window-http")
- c.Request.Header.Set("X-Codex-Installation-ID", "installation-http")
- c.Request.Header.Set("X-Test", "blocked")
- body := []byte(`{"model":"gpt-5","input":"hello"}`)
- svc := &OpenAIGatewayService{cfg: &config.Config{
- Security: config.SecurityConfig{
- URLAllowlist: config.URLAllowlistConfig{Enabled: false},
- },
- }}
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, body, "token", false, "", true)
- require.NoError(t, err)
- require.Equal(t, "window-http", req.Header.Get("X-Codex-Window-ID"))
- require.Equal(t, "installation-http", req.Header.Get("X-Codex-Installation-ID"))
- require.Empty(t, req.Header.Get("X-Test"))
- require.True(t, openai.EvaluateEngineFingerprint(req.Header, body, openai.DefaultEngineFingerprintSignals))
- }
- func TestOpenAIBuildUpstreamRequestOAuthOfficialClientOriginatorCompatibility(t *testing.T) {
- gin.SetMode(gin.TestMode)
- // 强制统一出口:客户端自报的 originator / User-Agent 都不参与上游身份构造,
- // 一律改写为网关规范身份,天然满足 originator 与 UA 首段配套的上游校验(issue #3901)。
- tests := []struct {
- name string
- userAgent string
- originator string
- }{
- {name: "official desktop ua", userAgent: "Codex Desktop/1.2.3"},
- {
- name: "mismatched originator",
- userAgent: "codex_vscode/0.140.2 (Mac OS X 14.0; arm64) vscode (codex_vscode; 0.140.2)",
- originator: "codex_cli_rs",
- },
- {
- name: "tui identity",
- userAgent: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
- originator: "codex-tui",
- },
- {name: "official originator without ua", originator: "codex_vscode"},
- {name: "third-party ua", userAgent: "luna/1.2.0"},
- }
- 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/responses", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
- if tt.userAgent != "" {
- c.Request.Header.Set("User-Agent", tt.userAgent)
- }
- if tt.originator != "" {
- c.Request.Header.Set("originator", tt.originator)
- }
- svc := &OpenAIGatewayService{}
- account := &Account{
- Type: AccountTypeOAuth,
- Credentials: map[string]any{"chatgpt_account_id": "chatgpt-acc"},
- }
- isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator"))
- req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", isCodexCLI)
- require.NoError(t, err)
- require.Equal(t, openai.CodexDefaultOriginator, req.Header.Get("originator"))
- require.Equal(t, codexCLIUserAgent, req.Header.Get("User-Agent"))
- require.Equal(t, codexCLIVersion, req.Header.Get("version"))
- })
- }
- }
- // ==================== P1-08 修复:model 替换性能优化测试 ====================
- // ==================== P1-08 修复:model 替换性能优化测试 =============
- func TestReplaceModelInSSELine(t *testing.T) {
- svc := &OpenAIGatewayService{}
- tests := []struct {
- name string
- line string
- from string
- to string
- expected string
- }{
- {
- name: "顶层 model 字段替换",
- line: `data: {"id":"chatcmpl-123","model":"gpt-4o","choices":[]}`,
- from: "gpt-4o",
- to: "my-custom-model",
- expected: `data: {"id":"chatcmpl-123","model":"my-custom-model","choices":[]}`,
- },
- {
- name: "嵌套 response.model 替换",
- line: `data: {"type":"response","response":{"id":"resp-1","model":"gpt-4o","output":[]}}`,
- from: "gpt-4o",
- to: "my-model",
- expected: `data: {"type":"response","response":{"id":"resp-1","model":"my-model","output":[]}}`,
- },
- {
- name: "model 不匹配时不替换",
- line: `data: {"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
- from: "gpt-4o",
- to: "my-model",
- expected: `data: {"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
- },
- {
- name: "无 model 字段时不替换",
- line: `data: {"id":"chatcmpl-123","choices":[]}`,
- from: "gpt-4o",
- to: "my-model",
- expected: `data: {"id":"chatcmpl-123","choices":[]}`,
- },
- {
- name: "空 data 行",
- line: `data: `,
- from: "gpt-4o",
- to: "my-model",
- expected: `data: `,
- },
- {
- name: "[DONE] 行",
- line: `data: [DONE]`,
- from: "gpt-4o",
- to: "my-model",
- expected: `data: [DONE]`,
- },
- {
- name: "非 data: 前缀行",
- line: `event: message`,
- from: "gpt-4o",
- to: "my-model",
- expected: `event: message`,
- },
- {
- name: "非法 JSON 不替换",
- line: `data: {invalid json}`,
- from: "gpt-4o",
- to: "my-model",
- expected: `data: {invalid json}`,
- },
- {
- name: "无空格 data: 格式",
- line: `data:{"id":"x","model":"gpt-4o"}`,
- from: "gpt-4o",
- to: "my-model",
- expected: `data: {"id":"x","model":"my-model"}`,
- },
- {
- name: "model 名含特殊字符",
- line: `data: {"model":"org/model-v2.1-beta"}`,
- from: "org/model-v2.1-beta",
- to: "custom/alias",
- expected: `data: {"model":"custom/alias"}`,
- },
- {
- name: "空行",
- line: "",
- from: "gpt-4o",
- to: "my-model",
- expected: "",
- },
- {
- name: "保持其他字段不变",
- line: `data: {"id":"abc","object":"chat.completion.chunk","model":"gpt-4o","created":1234567890,"choices":[{"index":0,"delta":{"content":"hi"}}]}`,
- from: "gpt-4o",
- to: "alias",
- expected: `data: {"id":"abc","object":"chat.completion.chunk","model":"alias","created":1234567890,"choices":[{"index":0,"delta":{"content":"hi"}}]}`,
- },
- {
- name: "顶层优先于嵌套:同时存在两个 model",
- line: `data: {"model":"gpt-4o","response":{"model":"gpt-4o"}}`,
- from: "gpt-4o",
- to: "replaced",
- expected: `data: {"model":"replaced","response":{"model":"gpt-4o"}}`,
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- got := svc.replaceModelInSSELine(tt.line, tt.from, tt.to)
- require.Equal(t, tt.expected, got)
- })
- }
- }
- func TestReplaceModelInSSEBody(t *testing.T) {
- svc := &OpenAIGatewayService{}
- tests := []struct {
- name string
- body string
- from string
- to string
- expected string
- }{
- {
- name: "多行 SSE body 替换",
- body: "data: {\"model\":\"gpt-4o\",\"choices\":[]}\n\ndata: {\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n",
- from: "gpt-4o",
- to: "alias",
- expected: "data: {\"model\":\"alias\",\"choices\":[]}\n\ndata: {\"model\":\"alias\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n",
- },
- {
- name: "无需替换的 body",
- body: "data: {\"model\":\"gpt-3.5-turbo\"}\n\ndata: [DONE]\n",
- from: "gpt-4o",
- to: "alias",
- expected: "data: {\"model\":\"gpt-3.5-turbo\"}\n\ndata: [DONE]\n",
- },
- {
- name: "混合 event 和 data 行",
- body: "event: message\ndata: {\"model\":\"gpt-4o\"}\n\n",
- from: "gpt-4o",
- to: "alias",
- expected: "event: message\ndata: {\"model\":\"alias\"}\n\n",
- },
- {
- name: "空 body",
- body: "",
- from: "gpt-4o",
- to: "alias",
- expected: "",
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- got := svc.replaceModelInSSEBody(tt.body, tt.from, tt.to)
- require.Equal(t, tt.expected, got)
- })
- }
- }
- func TestReplaceModelInResponseBody(t *testing.T) {
- svc := &OpenAIGatewayService{}
- tests := []struct {
- name string
- body string
- from string
- to string
- expected string
- }{
- {
- name: "替换顶层 model",
- body: `{"id":"chatcmpl-123","model":"gpt-4o","choices":[]}`,
- from: "gpt-4o",
- to: "alias",
- expected: `{"id":"chatcmpl-123","model":"alias","choices":[]}`,
- },
- {
- name: "model 不匹配不替换",
- body: `{"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
- from: "gpt-4o",
- to: "alias",
- expected: `{"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
- },
- {
- name: "无 model 字段不替换",
- body: `{"id":"chatcmpl-123","choices":[]}`,
- from: "gpt-4o",
- to: "alias",
- expected: `{"id":"chatcmpl-123","choices":[]}`,
- },
- {
- name: "非法 JSON 返回原值",
- body: `not json`,
- from: "gpt-4o",
- to: "alias",
- expected: `not json`,
- },
- {
- name: "空 body 返回原值",
- body: ``,
- from: "gpt-4o",
- to: "alias",
- expected: ``,
- },
- {
- name: "保持嵌套结构不变",
- body: `{"model":"gpt-4o","usage":{"prompt_tokens":10,"completion_tokens":20},"choices":[{"message":{"role":"assistant","content":"hello"}}]}`,
- from: "gpt-4o",
- to: "alias",
- expected: `{"model":"alias","usage":{"prompt_tokens":10,"completion_tokens":20},"choices":[{"message":{"role":"assistant","content":"hello"}}]}`,
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- got := svc.replaceModelInResponseBody([]byte(tt.body), tt.from, tt.to)
- require.Equal(t, tt.expected, string(got))
- })
- }
- }
- func TestExtractOpenAISSEDataLine(t *testing.T) {
- tests := []struct {
- name string
- line string
- wantData string
- wantOK bool
- }{
- {name: "标准格式", line: `data: {"type":"x"}`, wantData: `{"type":"x"}`, wantOK: true},
- {name: "无空格格式", line: `data:{"type":"x"}`, wantData: `{"type":"x"}`, wantOK: true},
- {name: "纯空数据", line: `data: `, wantData: ``, wantOK: true},
- {name: "非 data 行", line: `event: message`, wantData: ``, wantOK: false},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- got, ok := extractOpenAISSEDataLine(tt.line)
- require.Equal(t, tt.wantOK, ok)
- require.Equal(t, tt.wantData, got)
- })
- }
- }
- func TestParseSSEUsage_SelectiveParsing(t *testing.T) {
- svc := &OpenAIGatewayService{}
- usage := &OpenAIUsage{InputTokens: 9, OutputTokens: 8, CacheReadInputTokens: 7}
- // 非终态事件中的显式 usage 作为兼容 fallback,非零字段会被合并。
- svc.parseSSEUsage(`{"type":"response.in_progress","response":{"usage":{"input_tokens":1,"output_tokens":2}}}`, usage)
- require.Equal(t, 1, usage.InputTokens)
- require.Equal(t, 2, usage.OutputTokens)
- require.Equal(t, 7, usage.CacheReadInputTokens)
- // completed 事件,应提取 usage
- svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"input_tokens":3,"output_tokens":5,"input_tokens_details":{"cached_tokens":2}}}}`, usage)
- require.Equal(t, 3, usage.InputTokens)
- require.Equal(t, 5, usage.OutputTokens)
- require.Equal(t, 2, usage.CacheReadInputTokens)
- // done 事件同样可能携带最终 usage
- svc.parseSSEUsage(`{"type":"response.done","response":{"usage":{"input_tokens":13,"output_tokens":15,"input_tokens_details":{"cached_tokens":4}}}}`, usage)
- require.Equal(t, 13, usage.InputTokens)
- require.Equal(t, 15, usage.OutputTokens)
- require.Equal(t, 4, usage.CacheReadInputTokens)
- // failed 事件在部分上游路径也会携带已消耗 usage,应与 WS/passthrough 保持一致
- svc.parseSSEUsage(`{"type":"response.failed","response":{"usage":{"input_tokens":17,"output_tokens":19,"input_tokens_details":{"cached_tokens":6}}}}`, usage)
- require.Equal(t, 17, usage.InputTokens)
- require.Equal(t, 19, usage.OutputTokens)
- require.Equal(t, 6, usage.CacheReadInputTokens)
- svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"prompt_tokens":21,"completion_tokens":8,"prompt_tokens_details":{"cached_tokens":6}}}}`, usage)
- require.Equal(t, 21, usage.InputTokens)
- require.Equal(t, 8, usage.OutputTokens)
- require.Equal(t, 6, usage.CacheReadInputTokens)
- }
- func TestParseSSEUsage_NonTerminalUsageMergesNonZeroFields(t *testing.T) {
- svc := &OpenAIGatewayService{}
- usage := &OpenAIUsage{}
- svc.parseSSEUsage(`{"type":"response.in_progress","usage":{"input_tokens":17,"output_tokens":1,"input_tokens_details":{"cached_tokens":4}}}`, usage)
- svc.parseSSEUsage(`{"type":"response.output_text.done","usage":{"input_tokens":0,"output_tokens":5,"input_tokens_details":{"cached_tokens":0,"cache_write_tokens":3}}}`, usage)
- require.Equal(t, 17, usage.InputTokens)
- require.Equal(t, 5, usage.OutputTokens)
- require.Equal(t, 4, usage.CacheReadInputTokens)
- require.Equal(t, 3, usage.CacheCreationInputTokens)
- }
- func TestParseSSEUsage_TerminalUsageReplacesFallback(t *testing.T) {
- svc := &OpenAIGatewayService{}
- usage := &OpenAIUsage{}
- svc.parseSSEUsage(`{"type":"response.output_text.done","usage":{"input_tokens":17,"output_tokens":5,"input_tokens_details":{"cached_tokens":4}}}`, usage)
- svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"input_tokens":19,"output_tokens":7}}}`, usage)
- require.Equal(t, 19, usage.InputTokens)
- require.Equal(t, 7, usage.OutputTokens)
- require.Zero(t, usage.CacheReadInputTokens)
- }
- func TestParseSSEUsage_TerminalWithoutUsageKeepsFallback(t *testing.T) {
- svc := &OpenAIGatewayService{}
- usage := &OpenAIUsage{}
- svc.parseSSEUsage(`{"type":"response.in_progress","usage":{"input_tokens":17,"output_tokens":5}}`, usage)
- svc.parseSSEUsage(`{"type":"response.completed","response":{"id":"resp_1"}}`, usage)
- svc.parseSSEUsage(" [DONE]\n", usage)
- require.Equal(t, 17, usage.InputTokens)
- require.Equal(t, 5, usage.OutputTokens)
- }
- func TestExtractOpenAIUsageFromJSONBytes_AcceptsResponseAndChatUsageShapes(t *testing.T) {
- usage, ok := extractOpenAIUsageFromJSONBytes([]byte(`{"id":"resp_1","usage":{"input_tokens":9,"output_tokens":5,"input_tokens_details":{"cached_tokens":2,"cache_write_tokens":4}}}`))
- require.True(t, ok)
- require.Equal(t, 9, usage.InputTokens)
- require.Equal(t, 5, usage.OutputTokens)
- require.Equal(t, 2, usage.CacheReadInputTokens)
- require.Equal(t, 4, usage.CacheCreationInputTokens)
- usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"type":"response.completed","response":{"usage":{"prompt_tokens":13,"completion_tokens":7,"prompt_tokens_details":{"cached_tokens":4,"cache_creation_tokens":3}}}}`))
- require.True(t, ok)
- require.Equal(t, 13, usage.InputTokens)
- require.Equal(t, 7, usage.OutputTokens)
- require.Equal(t, 4, usage.CacheReadInputTokens)
- require.Equal(t, 3, usage.CacheCreationInputTokens)
- usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":11,"output_tokens":2,"cache_write_input_tokens":6}}`))
- require.True(t, ok)
- require.Equal(t, 6, usage.CacheCreationInputTokens)
- usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_creation_input_tokens":19,"input_tokens_details":{"cache_write_tokens":7}}}`))
- require.True(t, ok)
- require.Equal(t, 7, usage.CacheCreationInputTokens, "官方嵌套字段应优先于兼容顶层别名")
- usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_creation_input_tokens":19,"input_tokens_details":{"cache_write_tokens":0}}}`))
- require.True(t, ok)
- require.Zero(t, usage.CacheCreationInputTokens, "官方嵌套字段显式为零时仍应优先于兼容顶层别名")
- usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_read_input_tokens":19,"input_tokens_details":{"cached_tokens":0}}}`))
- require.True(t, ok)
- require.Zero(t, usage.CacheReadInputTokens, "官方嵌套缓存读取字段显式为零时仍应优先于兼容顶层别名")
- // xAI reports reasoning_tokens outside visible output_tokens. Only the
- // arithmetic-consistent shape is independent; OpenAI's canonical shape
- // already includes reasoning in completion/output_tokens.
- usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":10000,"output_tokens":500,"total_tokens":10800,"output_tokens_details":{"reasoning_tokens":300}}}`))
- require.True(t, ok)
- require.Equal(t, 800, usage.OutputTokens)
- usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":10000,"output_tokens":500,"total_tokens":10500,"output_tokens_details":{"reasoning_tokens":300}}}`))
- require.True(t, ok)
- require.Equal(t, 500, usage.OutputTokens)
- }
- func TestExtractOpenAIUsageFromJSONBytes_IncludesGrokReasoningTokens(t *testing.T) {
- usage, ok := extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"prompt_tokens":32,"completion_tokens":9,"total_tokens":135,"completion_tokens_details":{"reasoning_tokens":94}}}`))
- require.True(t, ok)
- require.Equal(t, 103, usage.OutputTokens, "Grok Chat usage bills visible completion plus reasoning tokens")
- usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":32,"output_tokens":103,"total_tokens":135,"output_tokens_details":{"reasoning_tokens":94}}}`))
- require.True(t, ok)
- require.Equal(t, 103, usage.OutputTokens, "Responses output_tokens already includes reasoning when total confirms it")
- usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":32,"output_tokens":9,"total_tokens":135,"output_tokens_details":{"reasoning_tokens":94}}}`))
- require.True(t, ok)
- require.Equal(t, 103, usage.OutputTokens, "Responses detail-only shape is normalized when total exposes the full output")
- }
- func TestExtractCodexFinalResponse_SampleReplay(t *testing.T) {
- body := strings.Join([]string{
- `event: message`,
- `data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
- `data: {"type":"response.completed","response":{"id":"resp_1","model":"gpt-4o","usage":{"input_tokens":11,"output_tokens":22,"input_tokens_details":{"cached_tokens":3}}}}`,
- `data: [DONE]`,
- }, "\n")
- finalResp, ok := extractCodexFinalResponse(body)
- require.True(t, ok)
- require.Contains(t, string(finalResp), `"id":"resp_1"`)
- require.Contains(t, string(finalResp), `"input_tokens":11`)
- }
- func TestHandleSSEToJSON_CompletedEventReturnsJSON(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- svc := &OpenAIGatewayService{cfg: &config.Config{}}
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"text/event-stream"}},
- }
- body := []byte(strings.Join([]string{
- `data: {"type":"response.in_progress","response":{"id":"resp_2"}}`,
- `data: {"type":"response.completed","response":{"id":"resp_2","model":"gpt-4o","usage":{"input_tokens":7,"output_tokens":9,"input_tokens_details":{"cached_tokens":1}}}}`,
- `data: [DONE]`,
- }, "\n"))
- usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-4o", "gpt-4o")
- require.NoError(t, err)
- require.NotNil(t, usage)
- require.Equal(t, 7, usage.InputTokens)
- require.Equal(t, 9, usage.OutputTokens)
- require.Equal(t, 1, usage.CacheReadInputTokens)
- // Header 可能由上游 Content-Type 透传;关键是 body 已转换为最终 JSON 响应。
- require.NotContains(t, rec.Body.String(), "event:")
- require.Contains(t, rec.Body.String(), `"id":"resp_2"`)
- require.NotContains(t, rec.Body.String(), "data:")
- }
- func TestHandleNonStreamingResponse_APIKeyFallsBackToSSEBodyWhenContentTypeIsWrong(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
- svc := &OpenAIGatewayService{cfg: &config.Config{}}
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"application/json"}},
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- `data: {"type":"response.output_text.delta","delta":"hel"}`,
- `data: {"type":"response.output_text.delta","delta":"lo"}`,
- `data: {"type":"response.completed","response":{"id":"resp_api_key_sse","object":"response","model":"gpt-5.4","status":"completed","output":[],"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}}`,
- `data: [DONE]`,
- }, "\n"))),
- }
- account := &Account{ID: 1, Type: AccountTypeAPIKey}
- result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.4", "gpt-5.4")
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, 3, result.InputTokens)
- require.Equal(t, 2, result.OutputTokens)
- require.NotContains(t, rec.Body.String(), "data:")
- require.Equal(t, "resp_api_key_sse", gjson.Get(rec.Body.String(), "id").String())
- require.Equal(t, "hello", gjson.Get(rec.Body.String(), "output.0.content.0.text").String())
- }
- func TestHandleNonStreamingResponse_OAuthJSONBodyWithDataEventTextKeepsJSONUsage(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", nil)
- svc := &OpenAIGatewayService{cfg: &config.Config{}}
- // Plain JSON compact response whose output text happens to contain the
- // literal substrings "data:" and "event:" (e.g. echoing shell/log output).
- // This must NOT be misdetected as SSE framing: it has a top-level usage
- // object and no upstream text/event-stream Content-Type.
- jsonBody := `{"id":"resp_oauth_compact","object":"response","model":"gpt-5.4","status":"completed",` +
- `"output":[{"type":"message","content":[{"type":"output_text",` +
- `"text":"processing data: 1,2,3 then event: click finished"}]}],` +
- `"usage":{"input_tokens":11,"output_tokens":22,"total_tokens":33}}`
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"application/json"}},
- Body: io.NopCloser(strings.NewReader(jsonBody)),
- }
- account := &Account{ID: 146, Type: AccountTypeOAuth}
- result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.4", "gpt-5.4")
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, 11, result.InputTokens)
- require.Equal(t, 22, result.OutputTokens)
- // Response must remain the original JSON body (not routed through the SSE
- // path, which would rewrite/lose the body or usage).
- require.Equal(t, "application/json", rec.Header().Get("Content-Type"))
- require.Equal(t, "resp_oauth_compact", gjson.Get(rec.Body.String(), "id").String())
- require.Equal(t, int64(33), gjson.Get(rec.Body.String(), "usage.total_tokens").Int())
- require.Contains(t, rec.Body.String(), "processing data: 1,2,3 then event: click finished")
- }
- func TestHandleNonStreamingResponse_ObservesUpstreamModelBeforeClientRewrite(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
- svc := &OpenAIGatewayService{cfg: &config.Config{}}
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"application/json"}},
- Body: io.NopCloser(strings.NewReader(
- `{"id":"resp_model_audit","object":"response","model":"gpt-5.5","status":"completed","output":[],"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}`,
- )),
- }
- account := &Account{ID: 1, Type: AccountTypeAPIKey}
- result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.6-sol", "gpt-5.5")
- require.NoError(t, err)
- require.NotNil(t, result)
- // 客户端仍看到自己请求的模型名,审计观察器则保留改写前的上游声明。
- require.Equal(t, "gpt-5.6-sol", gjson.Get(rec.Body.String(), "model").String())
- require.Equal(t, "gpt-5.5", observedUpstreamResponseModel(c))
- require.False(t, observedUpstreamResponseModelConflict(c))
- }
- func TestHandleSSEToJSON_ReconstructsImageGenerationOutputItemDone(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- svc := &OpenAIGatewayService{cfg: &config.Config{}}
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"text/event-stream"}},
- }
- body := []byte(strings.Join([]string{
- `data: {"type":"response.output_item.done","item":{"id":"ig_123","type":"image_generation_call","status":"generating","result":"aGVsbG8=","revised_prompt":"draw a cat","output_format":"png"}}`,
- `data: {"type":"response.completed","response":{"id":"resp_img","model":"gpt-5.4","output":[],"usage":{"input_tokens":7,"output_tokens":9,"output_tokens_details":{"image_tokens":4}}}}`,
- `data: [DONE]`,
- }, "\n"))
- usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-5.4", "gpt-5.4")
- require.NoError(t, err)
- require.NotNil(t, usage)
- require.Equal(t, 4, usage.ImageOutputTokens)
- require.NotContains(t, rec.Body.String(), "data:")
- require.Equal(t, "image_generation_call", gjson.Get(rec.Body.String(), "output.0.type").String())
- require.Equal(t, "completed", gjson.Get(rec.Body.String(), "output.0.status").String())
- require.Equal(t, "aGVsbG8=", gjson.Get(rec.Body.String(), "output.0.result").String())
- require.Equal(t, "draw a cat", gjson.Get(rec.Body.String(), "output.0.revised_prompt").String())
- }
- func TestHandleSSEToJSON_NoFinalResponseKeepsSSEBody(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- svc := &OpenAIGatewayService{cfg: &config.Config{}}
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"text/event-stream"}},
- }
- body := []byte(strings.Join([]string{
- `data: {"type":"response.in_progress","response":{"id":"resp_3"}}`,
- `data: [DONE]`,
- }, "\n"))
- usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-4o", "gpt-4o")
- require.NoError(t, err)
- require.NotNil(t, usage)
- require.Equal(t, 0, usage.InputTokens)
- require.Contains(t, rec.Header().Get("Content-Type"), "text/event-stream")
- require.Contains(t, rec.Body.String(), `data: {"type":"response.in_progress"`)
- }
- // 无账号时没有可换的对象:newOpenAIStreamFailoverError 要拿 account 记录 ops 归属与
- // 账号健康,故这一支保持原有的协议错误行为。带真实账号的同一报文改为换号,
- // 由 TestNonStreamingSSEToJSON_UnclassifiedFailedEventFailsOver 钉死(issue #5281)。
- func TestHandleSSEToJSON_ResponseFailedWithoutAccountReturnsProtocolError(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- svc := &OpenAIGatewayService{cfg: &config.Config{}}
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"text/event-stream"}},
- }
- body := []byte(strings.Join([]string{
- `data: {"type":"response.failed","error":{"message":"upstream rejected request"}}`,
- `data: [DONE]`,
- }, "\n"))
- usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-4o", "gpt-4o")
- require.Nil(t, usage)
- require.Error(t, err)
- require.Equal(t, http.StatusBadGateway, rec.Code)
- require.Contains(t, rec.Body.String(), "upstream rejected request")
- require.Contains(t, rec.Header().Get("Content-Type"), "application/json")
- }
- func TestOpenAICompatSSEFrameParserResetsEventTypeAtFrameBoundary(t *testing.T) {
- var parser openAICompatSSEFrameParser
- frame, ok := parser.AddLine("event: response.created")
- require.False(t, ok)
- require.Empty(t, frame)
- frame, ok = parser.AddLine(`data: {"response":{"id":"resp_1"}}`)
- require.False(t, ok)
- require.Empty(t, frame)
- frame, ok = parser.AddLine("")
- require.True(t, ok)
- require.Equal(t, "response.created", frame.EventType)
- require.JSONEq(t, `{"response":{"id":"resp_1"}}`, frame.Data)
- frame, ok = parser.AddLine(`data: {"delta":"ok"}`)
- require.False(t, ok)
- require.Empty(t, frame.EventType)
- frame, ok = parser.AddLine("")
- require.True(t, ok)
- require.Empty(t, frame.EventType)
- require.JSONEq(t, `{"delta":"ok"}`, frame.Data)
- }
- func TestStreamingPassthroughCyberPolicyMarksAndPassesThrough(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"r1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","response":{"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid-cyber"}},
- }
- _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, time.Now(), "m", "m")
- require.Error(t, err)
- var failoverErr *UpstreamFailoverError
- require.False(t, errors.As(err, &failoverErr), "cyber must NOT failover")
- require.Contains(t, rec.Body.String(), "cyber_policy", "response.failed passed through to client")
- mark := GetOpsCyberPolicy(c)
- require.NotNil(t, mark)
- require.Equal(t, "flagged for cyber policy", mark.Message)
- }
- func TestHandleStreamingResponseCyberPolicyMarks(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}
- svc := &OpenAIGatewayService{cfg: cfg}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Body: io.NopCloser(strings.NewReader(strings.Join([]string{
- "event: response.created",
- `data: {"type":"response.created","response":{"id":"r1"}}`,
- "",
- "event: response.failed",
- `data: {"type":"response.failed","error":{"code":"cyber_policy","message":"flagged"}}`,
- "",
- }, "\n"))),
- Header: http.Header{"X-Request-Id": []string{"rid"}},
- }
- _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, time.Now(), "m", "m")
- require.Error(t, err)
- var fo *UpstreamFailoverError
- require.False(t, errors.As(err, &fo))
- require.NotNil(t, GetOpsCyberPolicy(c))
- }
- func TestHandleErrorResponseCyberPolicyPassthrough(t *testing.T) {
- gin.SetMode(gin.TestMode)
- svc := &OpenAIGatewayService{cfg: &config.Config{}}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- cyberBody := `{"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}`
- resp := &http.Response{
- StatusCode: http.StatusBadRequest,
- Header: http.Header{"Content-Type": []string{"application/json"}, "X-Request-Id": []string{"rid"}},
- Body: io.NopCloser(strings.NewReader(cyberBody)),
- }
- _, err := svc.handleErrorResponse(context.Background(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, nil)
- require.Error(t, err)
- require.Equal(t, http.StatusBadRequest, rec.Code, "passthrough upstream 400, not rewrapped 502")
- require.Contains(t, rec.Body.String(), "cyber_policy", "client sees original cyber body")
- require.NotContains(t, rec.Body.String(), "Upstream request failed", "must not 502-rewrap")
- mark := GetOpsCyberPolicy(c)
- require.NotNil(t, mark)
- require.Equal(t, http.StatusBadRequest, mark.UpstreamStatus)
- }
- func TestHandleCompatErrorResponseCyberPolicyEarlyReturn(t *testing.T) {
- gin.SetMode(gin.TestMode)
- svc := &OpenAIGatewayService{cfg: &config.Config{}}
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
- cyberBody := `{"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}`
- resp := &http.Response{
- StatusCode: http.StatusBadRequest,
- Header: http.Header{"Content-Type": []string{"application/json"}},
- Body: io.NopCloser(strings.NewReader(cyberBody)),
- }
- var gotStatus int
- var gotType, gotMsg string
- writeError := func(_ *gin.Context, statusCode int, errType, message string) {
- gotStatus, gotType, gotMsg = statusCode, errType, message
- }
- // cyber 命中应早返回(写兼容错误 + 不冷却账号),而非落到通用 "Upstream request failed"。
- _, err := svc.handleCompatErrorResponse(resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, writeError)
- require.Error(t, err)
- require.Equal(t, http.StatusBadRequest, gotStatus)
- require.Equal(t, "invalid_request_error", gotType)
- require.Contains(t, gotMsg, "flagged for cyber policy")
- require.NotContains(t, gotMsg, "Upstream request failed")
- require.NotNil(t, GetOpsCyberPolicy(c))
- }
|