openai_gateway_service_test.go 145 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546354735483549355035513552355335543555355635573558355935603561356235633564356535663567356835693570357135723573357435753576357735783579358035813582358335843585358635873588358935903591359235933594359535963597359835993600360136023603360436053606360736083609361036113612361336143615361636173618361936203621362236233624362536263627362836293630363136323633363436353636363736383639364036413642364336443645364636473648364936503651365236533654365536563657365836593660366136623663366436653666366736683669367036713672367336743675367636773678367936803681368236833684368536863687368836893690369136923693369436953696369736983699370037013702370337043705370637073708370937103711371237133714371537163717371837193720372137223723372437253726372737283729373037313732373337343735373637373738373937403741374237433744374537463747374837493750375137523753375437553756375737583759376037613762376337643765376637673768376937703771377237733774377537763777377837793780378137823783378437853786378737883789379037913792379337943795379637973798379938003801380238033804380538063807380838093810381138123813381438153816381738183819382038213822382338243825382638273828382938303831383238333834383538363837383838393840384138423843384438453846384738483849385038513852385338543855385638573858385938603861386238633864386538663867386838693870387138723873387438753876387738783879388038813882388338843885388638873888388938903891389238933894389538963897389838993900390139023903390439053906390739083909391039113912391339143915391639173918391939203921392239233924392539263927392839293930393139323933393439353936393739383939394039413942394339443945394639473948394939503951395239533954395539563957395839593960396139623963396439653966
  1. package service
  2. import (
  3. "bufio"
  4. "bytes"
  5. "context"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "net/http"
  10. "net/http/httptest"
  11. "strings"
  12. "testing"
  13. "time"
  14. "github.com/Wei-Shaw/sub2api/internal/config"
  15. "github.com/Wei-Shaw/sub2api/internal/model"
  16. "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
  17. "github.com/cespare/xxhash/v2"
  18. "github.com/gin-gonic/gin"
  19. "github.com/stretchr/testify/require"
  20. "github.com/tidwall/gjson"
  21. )
  22. // 编译期接口断言
  23. var _ AccountRepository = (*stubOpenAIAccountRepo)(nil)
  24. var _ GatewayCache = (*stubGatewayCache)(nil)
  25. type stubOpenAIAccountRepo struct {
  26. AccountRepository
  27. accounts []Account
  28. }
  29. type tempUnschedulableOpenAIAccountRepo struct {
  30. stubOpenAIAccountRepo
  31. modelRateLimitAccountID int64
  32. modelRateLimitKey string
  33. }
  34. func (r *tempUnschedulableOpenAIAccountRepo) SetModelRateLimit(_ context.Context, accountID int64, modelKey string, _ time.Time, _ ...string) error {
  35. r.modelRateLimitAccountID = accountID
  36. r.modelRateLimitKey = modelKey
  37. return nil
  38. }
  39. type snapshotUpdateAccountRepo struct {
  40. stubOpenAIAccountRepo
  41. updateExtraCalls chan map[string]any
  42. }
  43. func (r *snapshotUpdateAccountRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
  44. if r.updateExtraCalls != nil {
  45. copied := make(map[string]any, len(updates))
  46. for k, v := range updates {
  47. copied[k] = v
  48. }
  49. r.updateExtraCalls <- copied
  50. }
  51. return nil
  52. }
  53. func (r stubOpenAIAccountRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
  54. for i := range r.accounts {
  55. if r.accounts[i].ID == id {
  56. return &r.accounts[i], nil
  57. }
  58. }
  59. return nil, errors.New("account not found")
  60. }
  61. func (r stubOpenAIAccountRepo) GetByIDs(ctx context.Context, ids []int64) ([]*Account, error) {
  62. if len(ids) == 0 {
  63. return []*Account{}, nil
  64. }
  65. index := make(map[int64]*Account, len(r.accounts))
  66. for i := range r.accounts {
  67. account := &r.accounts[i]
  68. index[account.ID] = account
  69. }
  70. out := make([]*Account, 0, len(ids))
  71. seen := make(map[int64]struct{}, len(ids))
  72. for _, id := range ids {
  73. if _, ok := seen[id]; ok {
  74. continue
  75. }
  76. seen[id] = struct{}{}
  77. if account, ok := index[id]; ok {
  78. out = append(out, account)
  79. }
  80. }
  81. return out, nil
  82. }
  83. func (r stubOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
  84. var result []Account
  85. for _, acc := range r.accounts {
  86. if acc.Platform == platform {
  87. result = append(result, acc)
  88. }
  89. }
  90. return result, nil
  91. }
  92. func (r stubOpenAIAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
  93. var result []Account
  94. for _, acc := range r.accounts {
  95. if acc.Platform == platform {
  96. result = append(result, acc)
  97. }
  98. }
  99. return result, nil
  100. }
  101. func (r stubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
  102. return r.ListSchedulableByPlatform(ctx, platform)
  103. }
  104. func TestOpenAIGatewayService_ForwardAsAnthropic_CapacityShedReturnsRequestScopedFailoverWithoutCommit(t *testing.T) {
  105. gin.SetMode(gin.TestMode)
  106. body := []byte(`{"model":"gpt-5.4","max_tokens":32,"messages":[{"role":"user","content":"hello"}],"stream":false}`)
  107. rec := httptest.NewRecorder()
  108. c, _ := gin.CreateTestContext(rec)
  109. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
  110. c.Request.Header.Set("Content-Type", "application/json")
  111. upstreamBody := []byte(`{"error":{"message":"Our servers are currently overloaded. Please try again later."}}`)
  112. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  113. {
  114. StatusCode: http.StatusBadRequest,
  115. Header: http.Header{"Content-Type": []string{"application/json"}},
  116. Body: io.NopCloser(bytes.NewReader(upstreamBody)),
  117. },
  118. {
  119. StatusCode: http.StatusOK,
  120. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  121. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  122. `data: {"type":"response.output_text.delta","delta":"ok"}`,
  123. "",
  124. `data: {"type":"response.completed","response":{"id":"resp_second","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1}}}`,
  125. "",
  126. }, "\n"))),
  127. },
  128. }}
  129. repo := &tempUnschedulableOpenAIAccountRepo{}
  130. rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  131. svc := &OpenAIGatewayService{
  132. cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{
  133. Enabled: false, AllowInsecureHTTP: true,
  134. }}},
  135. httpUpstream: upstream,
  136. rateLimitService: rateLimitService,
  137. }
  138. account := &Account{
  139. ID: 5099, Name: "temporary-unschedulable", Platform: PlatformOpenAI,
  140. Type: AccountTypeAPIKey, Concurrency: 1,
  141. Credentials: map[string]any{
  142. "api_key": "sk-test",
  143. "base_url": "http://upstream.example",
  144. "model_provider": "env-openai",
  145. "temp_unschedulable_enabled": true,
  146. "temp_unschedulable_rules": []any{map[string]any{
  147. "error_code": float64(http.StatusBadRequest),
  148. "keywords": []any{"our servers are currently overloaded", "please try again later"},
  149. "duration_minutes": float64(1),
  150. }},
  151. },
  152. }
  153. _, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
  154. var failoverErr *UpstreamFailoverError
  155. require.ErrorAs(t, err, &failoverErr)
  156. require.Equal(t, http.StatusBadRequest, failoverErr.StatusCode)
  157. require.True(t, failoverErr.ShouldRetryNextAccount())
  158. require.True(t, failoverErr.RetryableOnSameAccount)
  159. require.True(t, failoverErr.RequestScopedTransient)
  160. require.Zero(t, repo.modelRateLimitAccountID, "request-scoped capacity shedding must not change account health")
  161. require.Empty(t, repo.modelRateLimitKey)
  162. require.False(t, IsResponseCommitted(c))
  163. require.Equal(t, http.StatusOK, rec.Code)
  164. require.Empty(t, rec.Body.String())
  165. secondRec := httptest.NewRecorder()
  166. secondContext, _ := gin.CreateTestContext(secondRec)
  167. secondContext.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
  168. secondContext.Request.Header.Set("Content-Type", "application/json")
  169. secondAccount := *account
  170. secondAccount.ID = 5100
  171. secondAccount.Name = "healthy-failover-account"
  172. result, secondErr := svc.ForwardAsAnthropic(context.Background(), secondContext, &secondAccount, body, "", "")
  173. require.NoError(t, secondErr)
  174. require.NotNil(t, result)
  175. require.Equal(t, "resp_second", result.ResponseID)
  176. require.NotEmpty(t, secondRec.Body.String())
  177. }
  178. func TestFailoverOpenAIUpstreamHTTPError_NilContextSkipsTempUnschedulablePolicy(t *testing.T) {
  179. repo := &tempUnschedulableOpenAIAccountRepo{}
  180. svc := &OpenAIGatewayService{
  181. rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil),
  182. }
  183. account := &Account{
  184. ID: 5099, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
  185. Credentials: map[string]any{
  186. "temp_unschedulable_enabled": true,
  187. "temp_unschedulable_rules": []any{map[string]any{
  188. "error_code": float64(http.StatusBadRequest),
  189. "keywords": []any{"custom temporary outage"},
  190. "duration_minutes": float64(1),
  191. }},
  192. },
  193. }
  194. body := []byte(`{"error":{"message":"Custom temporary outage."}}`)
  195. resp := &http.Response{StatusCode: http.StatusBadRequest, Header: http.Header{}}
  196. got := svc.failoverOpenAIUpstreamHTTPError(
  197. context.Background(), nil, account, resp, body,
  198. "Custom temporary outage.", "gpt-5.4",
  199. )
  200. require.Nil(t, got)
  201. require.Zero(t, repo.modelRateLimitAccountID)
  202. require.Empty(t, repo.modelRateLimitKey)
  203. }
  204. type groupAwareStubOpenAIAccountRepo struct {
  205. stubOpenAIAccountRepo
  206. }
  207. func (r groupAwareStubOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
  208. var result []Account
  209. for _, acc := range r.accounts {
  210. if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, &groupID) {
  211. result = append(result, acc)
  212. }
  213. }
  214. return result, nil
  215. }
  216. func (r groupAwareStubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
  217. var result []Account
  218. for _, acc := range r.accounts {
  219. if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, nil) {
  220. result = append(result, acc)
  221. }
  222. }
  223. return result, nil
  224. }
  225. type stubConcurrencyCache struct {
  226. ConcurrencyCache
  227. loadBatchErr error
  228. loadMap map[int64]*AccountLoadInfo
  229. acquireResults map[int64]bool
  230. waitCounts map[int64]int
  231. skipDefaultLoad bool
  232. }
  233. type cancelReadCloser struct{}
  234. func (c cancelReadCloser) Read(p []byte) (int, error) { return 0, context.Canceled }
  235. func (c cancelReadCloser) Close() error { return nil }
  236. type errReadCloser struct {
  237. err error
  238. }
  239. func (r errReadCloser) Read([]byte) (int, error) { return 0, r.err }
  240. func (r errReadCloser) Close() error { return nil }
  241. type openAIStreamReadThenErrorCloser struct {
  242. reader *strings.Reader
  243. err error
  244. }
  245. func (r *openAIStreamReadThenErrorCloser) Read(p []byte) (int, error) {
  246. if r.reader != nil && r.reader.Len() > 0 {
  247. return r.reader.Read(p)
  248. }
  249. return 0, r.err
  250. }
  251. func (r *openAIStreamReadThenErrorCloser) Close() error { return nil }
  252. type failingGinWriter struct {
  253. gin.ResponseWriter
  254. failAfter int
  255. writes int
  256. }
  257. func (w *failingGinWriter) Write(p []byte) (int, error) {
  258. if w.writes >= w.failAfter {
  259. return 0, errors.New("write failed")
  260. }
  261. w.writes++
  262. return w.ResponseWriter.Write(p)
  263. }
  264. func (c stubConcurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
  265. if c.acquireResults != nil {
  266. if result, ok := c.acquireResults[accountID]; ok {
  267. return result, nil
  268. }
  269. }
  270. return true, nil
  271. }
  272. func (c stubConcurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int64, requestID string) error {
  273. return nil
  274. }
  275. func (c stubConcurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []AccountWithConcurrency) (map[int64]*AccountLoadInfo, error) {
  276. if c.loadBatchErr != nil {
  277. return nil, c.loadBatchErr
  278. }
  279. out := make(map[int64]*AccountLoadInfo, len(accounts))
  280. if c.skipDefaultLoad && c.loadMap != nil {
  281. for _, acc := range accounts {
  282. if load, ok := c.loadMap[acc.ID]; ok {
  283. out[acc.ID] = load
  284. }
  285. }
  286. return out, nil
  287. }
  288. for _, acc := range accounts {
  289. if c.loadMap != nil {
  290. if load, ok := c.loadMap[acc.ID]; ok {
  291. out[acc.ID] = load
  292. continue
  293. }
  294. }
  295. out[acc.ID] = &AccountLoadInfo{AccountID: acc.ID, LoadRate: 0}
  296. }
  297. return out, nil
  298. }
  299. func TestOpenAIGatewayService_GenerateSessionHash_Priority(t *testing.T) {
  300. gin.SetMode(gin.TestMode)
  301. rec := httptest.NewRecorder()
  302. c, _ := gin.CreateTestContext(rec)
  303. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  304. svc := &OpenAIGatewayService{}
  305. bodyWithKey := []byte(`{"prompt_cache_key":"ses_aaa"}`)
  306. // 1) session_id header wins
  307. c.Request.Header.Set("session_id", "sess-123")
  308. c.Request.Header.Set("conversation_id", "conv-456")
  309. h1 := svc.GenerateSessionHash(c, bodyWithKey)
  310. if h1 == "" {
  311. t.Fatalf("expected non-empty hash")
  312. }
  313. // 2) conversation_id used when session_id absent
  314. c.Request.Header.Del("session_id")
  315. h2 := svc.GenerateSessionHash(c, bodyWithKey)
  316. if h2 == "" {
  317. t.Fatalf("expected non-empty hash")
  318. }
  319. if h1 == h2 {
  320. t.Fatalf("expected different hashes for different keys")
  321. }
  322. // 3) prompt_cache_key used when both headers absent
  323. c.Request.Header.Del("conversation_id")
  324. h3 := svc.GenerateSessionHash(c, bodyWithKey)
  325. if h3 == "" {
  326. t.Fatalf("expected non-empty hash")
  327. }
  328. if h2 == h3 {
  329. t.Fatalf("expected different hashes for different keys")
  330. }
  331. // 4) empty when no signals
  332. h4 := svc.GenerateSessionHash(c, []byte(`{}`))
  333. if h4 != "" {
  334. t.Fatalf("expected empty hash when no signals")
  335. }
  336. }
  337. func TestOpenAIGatewayService_ClientSessionHeaderPriority(t *testing.T) {
  338. gin.SetMode(gin.TestMode)
  339. recorder := httptest.NewRecorder()
  340. c, _ := gin.CreateTestContext(recorder)
  341. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  342. c.Set("api_key", &APIKey{ID: 901, Group: &Group{Platform: PlatformGrok}})
  343. headers := []struct {
  344. name string
  345. value string
  346. }{
  347. {name: "session-id", value: "codex-session"},
  348. {name: "session_id", value: "generic-session"},
  349. {name: "conversation_id", value: "generic-conversation"},
  350. {name: openCodeSessionAffinityHeader, value: "opencode-affinity"},
  351. {name: openCodeSessionIDHeader, value: "opencode-session-id"},
  352. {name: openCodeNativeSessionHeader, value: "opencode-native-session"},
  353. {name: codeBuddyConversationHeader, value: "codebuddy-conversation"},
  354. {name: grokConversationIDHeader, value: "grok-conversation"},
  355. }
  356. for _, header := range headers {
  357. c.Request.Header.Set(header.name, header.value)
  358. }
  359. svc := &OpenAIGatewayService{}
  360. body := []byte(`{"prompt_cache_key":"body-session"}`)
  361. for _, header := range headers {
  362. require.Equal(t, header.value, svc.ExtractSessionID(c, body), header.name)
  363. require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String(header.value)), svc.GenerateExplicitSessionHash(c, body), header.name)
  364. if header.name != grokConversationIDHeader {
  365. require.Equal(t, header.value, explicitOpenAISessionID(c, body), header.name)
  366. }
  367. c.Request.Header.Del(header.name)
  368. }
  369. require.Equal(t, "body-session", svc.ExtractSessionID(c, body))
  370. }
  371. func TestOpenAIGatewayService_CodexSessionIDKeepsReconnectHashStable(t *testing.T) {
  372. gin.SetMode(gin.TestMode)
  373. recorder := httptest.NewRecorder()
  374. c, _ := gin.CreateTestContext(recorder)
  375. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  376. c.Request.Header.Set("session-id", "codex-reconnect-session")
  377. svc := &OpenAIGatewayService{}
  378. warmup := []byte(`{
  379. "type":"response.create",
  380. "model":"gpt-5.6-sol",
  381. "generate":false,
  382. "tools":[{"type":"custom","name":"exec"}],
  383. "input":[{"role":"user","content":"warmup"}]
  384. }`)
  385. business := []byte(`{
  386. "type":"response.create",
  387. "model":"gpt-5.6-sol",
  388. "input":[{"role":"user","content":"install codex"}]
  389. }`)
  390. require.Equal(t, svc.GenerateSessionHash(c, warmup), svc.GenerateSessionHash(c, business))
  391. require.Equal(t, "codex-reconnect-session", svc.ExtractSessionID(c, business))
  392. }
  393. func TestOpenAIGatewayService_ClientSessionHeadersIgnorePerRequestIDs(t *testing.T) {
  394. gin.SetMode(gin.TestMode)
  395. recorder := httptest.NewRecorder()
  396. c, _ := gin.CreateTestContext(recorder)
  397. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  398. for name, value := range map[string]string{
  399. "X-Conversation-Request-ID": "request-rotates-every-turn",
  400. "X-Conversation-Message-ID": "message-rotates-every-turn",
  401. "X-Request-ID": "generic-request-id",
  402. } {
  403. c.Request.Header.Set(name, value)
  404. }
  405. svc := &OpenAIGatewayService{}
  406. require.Empty(t, explicitOpenAIHeaderSessionID(c))
  407. require.Empty(t, svc.ExtractSessionID(c, nil))
  408. require.Empty(t, svc.GenerateExplicitSessionHash(c, nil))
  409. }
  410. func TestOpenAIGatewayService_GenerateSessionHash_UsesXXHash64(t *testing.T) {
  411. gin.SetMode(gin.TestMode)
  412. rec := httptest.NewRecorder()
  413. c, _ := gin.CreateTestContext(rec)
  414. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  415. c.Request.Header.Set("session_id", "sess-fixed-value")
  416. svc := &OpenAIGatewayService{}
  417. got := svc.GenerateSessionHash(c, nil)
  418. want := fmt.Sprintf("%016x", xxhash.Sum64String("sess-fixed-value"))
  419. require.Equal(t, want, got)
  420. }
  421. func TestOpenAIGatewayService_GenerateSessionHash_AttachesLegacyHashToContext(t *testing.T) {
  422. gin.SetMode(gin.TestMode)
  423. rec := httptest.NewRecorder()
  424. c, _ := gin.CreateTestContext(rec)
  425. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  426. c.Request.Header.Set("session_id", "sess-legacy-check")
  427. svc := &OpenAIGatewayService{}
  428. sessionHash := svc.GenerateSessionHash(c, nil)
  429. require.NotEmpty(t, sessionHash)
  430. require.NotNil(t, c.Request)
  431. require.NotNil(t, c.Request.Context())
  432. require.NotEmpty(t, openAILegacySessionHashFromContext(c.Request.Context()))
  433. }
  434. func TestExtractOpenAIResponseIDFromJSONBytes(t *testing.T) {
  435. require.Equal(t, "resp_json", extractOpenAIResponseIDFromJSONBytes([]byte(`{"id":"resp_json"}`)))
  436. require.Equal(t, "resp_sse", extractOpenAIResponseIDFromJSONBytes([]byte(`{"type":"response.completed","response":{"id":"resp_sse"}}`)))
  437. require.Empty(t, extractOpenAIResponseIDFromJSONBytes([]byte(`{"response":{}}`)))
  438. require.Empty(t, extractOpenAIResponseIDFromJSONBytes([]byte(`not-json`)))
  439. }
  440. // 复现 #4386:gpt-image-2 /v1/images/edits 的 usage 携带 input_tokens_details.image_tokens,
  441. // 提取器须将图片输入 token 单独填入 ImageInputTokens(此前被丢弃并入 InputTokens 按文本价计费)。
  442. func TestExtractOpenAIUsage_CapturesImageInputTokens(t *testing.T) {
  443. 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}}`)
  444. usage, ok := extractOpenAIUsageFromJSONBytes(body)
  445. require.True(t, ok)
  446. require.Equal(t, 371, usage.InputTokens)
  447. require.Equal(t, 352, usage.ImageInputTokens)
  448. require.Equal(t, 439, usage.OutputTokens)
  449. require.Equal(t, 439, usage.ImageOutputTokens)
  450. // prompt_tokens_details 回退路径(部分上游用 prompt_tokens 口径)。
  451. promptStyle := []byte(`{"usage":{"prompt_tokens":100,"prompt_tokens_details":{"image_tokens":80}}}`)
  452. pu, ok := extractOpenAIUsageFromJSONBytes(promptStyle)
  453. require.True(t, ok)
  454. require.Equal(t, 100, pu.InputTokens)
  455. require.Equal(t, 80, pu.ImageInputTokens)
  456. // 纯文本请求:无 image_tokens 时 ImageInputTokens 为 0,行为不变。
  457. textOnly := []byte(`{"usage":{"input_tokens":50,"output_tokens":10}}`)
  458. tu, ok := extractOpenAIUsageFromJSONBytes(textOnly)
  459. require.True(t, ok)
  460. require.Zero(t, tu.ImageInputTokens)
  461. }
  462. func TestExtractOpenAIUsage_ReadsClineDataEnvelope(t *testing.T) {
  463. 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}`)
  464. usage, ok := extractOpenAIUsageFromJSONBytes(body)
  465. require.True(t, ok)
  466. require.Equal(t, 8, usage.InputTokens)
  467. require.Equal(t, 27, usage.OutputTokens)
  468. require.Equal(t, 4, usage.CacheReadInputTokens)
  469. }
  470. func TestExtractOpenAIUsage_ReadsWrappedResponsesDataEnvelope(t *testing.T) {
  471. body := []byte(`{"data":{"response":{"usage":{"input_tokens":11,"output_tokens":5,"total_tokens":16,"input_tokens_details":{"cached_tokens":2}}}}}`)
  472. usage, ok := extractOpenAIUsageFromJSONBytes(body)
  473. require.True(t, ok)
  474. require.Equal(t, 11, usage.InputTokens)
  475. require.Equal(t, 5, usage.OutputTokens)
  476. require.Equal(t, 2, usage.CacheReadInputTokens)
  477. }
  478. func TestExtractOpenAIUsage_PreservesResponseUsagePriority(t *testing.T) {
  479. body := []byte(`{"data":{"usage":{"prompt_tokens":100,"completion_tokens":50}},"response":{"usage":{"input_tokens":11,"output_tokens":5}}}`)
  480. usage, ok := extractOpenAIUsageFromJSONBytes(body)
  481. require.True(t, ok)
  482. require.Equal(t, 11, usage.InputTokens)
  483. require.Equal(t, 5, usage.OutputTokens)
  484. }
  485. func TestOpenAIGatewayService_BindHTTPResponseAccount(t *testing.T) {
  486. gin.SetMode(gin.TestMode)
  487. rec := httptest.NewRecorder()
  488. c, _ := gin.CreateTestContext(rec)
  489. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  490. groupID := int64(4201)
  491. c.Set("api_key", &APIKey{ID: 501, GroupID: &groupID})
  492. SetOpenAIHTTPResponseOwner(c, 601, 501)
  493. svc := &OpenAIGatewayService{}
  494. account := &Account{ID: 37001, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  495. svc.bindHTTPResponseAccount(context.Background(), c, account, "resp_http_001")
  496. got, err := svc.getOpenAIWSStateStore().GetResponseAccount(context.Background(), groupID, "resp_http_001")
  497. require.NoError(t, err)
  498. require.Equal(t, account.ID, got)
  499. owned, err := svc.ValidateOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_http_001", 601, 501)
  500. require.NoError(t, err)
  501. require.True(t, owned)
  502. owned, err = svc.ValidateOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_http_001", 601, 502)
  503. require.NoError(t, err)
  504. require.True(t, owned, "API keys owned by the same downstream user remain interoperable")
  505. owned, err = svc.ValidateOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_http_001", 602, 501)
  506. require.NoError(t, err)
  507. require.False(t, owned)
  508. owned, err = svc.ValidateOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_unknown", 601, 501)
  509. require.NoError(t, err)
  510. require.False(t, owned)
  511. }
  512. func TestOpenAIGatewayService_GenerateExplicitSessionHash_SkipsContentFallback(t *testing.T) {
  513. gin.SetMode(gin.TestMode)
  514. svc := &OpenAIGatewayService{}
  515. body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`)
  516. t.Run("stateless image body stays unstuck", func(t *testing.T) {
  517. rec := httptest.NewRecorder()
  518. c, _ := gin.CreateTestContext(rec)
  519. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
  520. require.Empty(t, svc.GenerateExplicitSessionHash(c, body))
  521. require.Empty(t, openAILegacySessionHashFromContext(c.Request.Context()))
  522. })
  523. t.Run("prompt_cache_key is explicit", func(t *testing.T) {
  524. rec := httptest.NewRecorder()
  525. c, _ := gin.CreateTestContext(rec)
  526. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
  527. got := svc.GenerateExplicitSessionHash(c, []byte(`{"model":"gpt-image-2","prompt_cache_key":"image-session"}`))
  528. require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String("image-session")), got)
  529. require.NotEmpty(t, openAILegacySessionHashFromContext(c.Request.Context()))
  530. })
  531. t.Run("header overrides body", func(t *testing.T) {
  532. rec := httptest.NewRecorder()
  533. c, _ := gin.CreateTestContext(rec)
  534. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
  535. c.Request.Header.Set("session_id", "header-session")
  536. got := svc.GenerateExplicitSessionHash(c, []byte(`{"prompt_cache_key":"body-session"}`))
  537. require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String("header-session")), got)
  538. })
  539. }
  540. func TestOpenAIGatewayService_GenerateSessionHashWithFallback(t *testing.T) {
  541. gin.SetMode(gin.TestMode)
  542. rec := httptest.NewRecorder()
  543. c, _ := gin.CreateTestContext(rec)
  544. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  545. svc := &OpenAIGatewayService{}
  546. seed := "openai_ws_ingress:9:100:200"
  547. got := svc.GenerateSessionHashWithFallback(c, []byte(`{}`), seed)
  548. want := fmt.Sprintf("%016x", xxhash.Sum64String(seed))
  549. require.Equal(t, want, got)
  550. require.NotEmpty(t, openAILegacySessionHashFromContext(c.Request.Context()))
  551. empty := svc.GenerateSessionHashWithFallback(c, []byte(`{}`), " ")
  552. require.Equal(t, "", empty)
  553. }
  554. func TestOpenAIGatewayService_GenerateSessionHash_ContentFallback(t *testing.T) {
  555. gin.SetMode(gin.TestMode)
  556. rec := httptest.NewRecorder()
  557. c, _ := gin.CreateTestContext(rec)
  558. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil)
  559. svc := &OpenAIGatewayService{}
  560. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"You are helpful."},{"role":"user","content":"Hello"}]}`)
  561. hash := svc.GenerateSessionHash(c, body)
  562. require.NotEmpty(t, hash, "content-based fallback should produce a hash")
  563. hash2 := svc.GenerateSessionHash(c, body)
  564. require.Equal(t, hash, hash2, "same content should produce same hash")
  565. 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?"}]}`)
  566. hashExtended := svc.GenerateSessionHash(c, bodyExtended)
  567. require.Equal(t, hash, hashExtended, "hash should be stable across later turns")
  568. bodyDifferent := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"Different question"}]}`)
  569. hashDifferent := svc.GenerateSessionHash(c, bodyDifferent)
  570. require.NotEqual(t, hash, hashDifferent, "different content should produce different hash")
  571. }
  572. func TestOpenAIGatewayService_GenerateSessionHash_ExplicitSignalWinsOverContent(t *testing.T) {
  573. gin.SetMode(gin.TestMode)
  574. rec := httptest.NewRecorder()
  575. c, _ := gin.CreateTestContext(rec)
  576. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil)
  577. svc := &OpenAIGatewayService{}
  578. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"Hello"}]}`)
  579. contentHash := svc.GenerateSessionHash(c, body)
  580. require.NotEmpty(t, contentHash)
  581. c.Request.Header.Set("session_id", "explicit-session")
  582. explicitHash := svc.GenerateSessionHash(c, body)
  583. require.NotEmpty(t, explicitHash)
  584. require.NotEqual(t, contentHash, explicitHash, "explicit session_id should override content fallback")
  585. }
  586. func TestOpenAIGatewayService_GenerateSessionHash_EmptyBodyStillEmpty(t *testing.T) {
  587. gin.SetMode(gin.TestMode)
  588. rec := httptest.NewRecorder()
  589. c, _ := gin.CreateTestContext(rec)
  590. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil)
  591. svc := &OpenAIGatewayService{}
  592. require.Empty(t, svc.GenerateSessionHash(c, []byte(`{}`)))
  593. require.Empty(t, svc.GenerateSessionHash(c, nil))
  594. }
  595. func (c stubConcurrencyCache) GetAccountWaitingCount(ctx context.Context, accountID int64) (int, error) {
  596. if c.waitCounts != nil {
  597. if count, ok := c.waitCounts[accountID]; ok {
  598. return count, nil
  599. }
  600. }
  601. return 0, nil
  602. }
  603. type stubGatewayCache struct {
  604. sessionBindings map[string]int64
  605. deletedSessions map[string]int
  606. }
  607. func (c *stubGatewayCache) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) {
  608. if id, ok := c.sessionBindings[sessionHash]; ok {
  609. return id, nil
  610. }
  611. return 0, errors.New("not found")
  612. }
  613. func (c *stubGatewayCache) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error {
  614. if c.sessionBindings == nil {
  615. c.sessionBindings = make(map[string]int64)
  616. }
  617. c.sessionBindings[sessionHash] = accountID
  618. return nil
  619. }
  620. func (c *stubGatewayCache) RefreshSessionTTL(ctx context.Context, groupID int64, sessionHash string, ttl time.Duration) error {
  621. return nil
  622. }
  623. func (c *stubGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error {
  624. if c.sessionBindings == nil {
  625. return nil
  626. }
  627. if c.deletedSessions == nil {
  628. c.deletedSessions = make(map[string]int)
  629. }
  630. c.deletedSessions[sessionHash]++
  631. delete(c.sessionBindings, sessionHash)
  632. return nil
  633. }
  634. func (c *stubGatewayCache) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
  635. return nil
  636. }
  637. func (c *stubGatewayCache) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
  638. return nil, nil
  639. }
  640. func (c *stubGatewayCache) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
  641. return true, nil
  642. }
  643. func (c *stubGatewayCache) ReleaseGrokVideoBilled(_ context.Context, _ string) error {
  644. return nil
  645. }
  646. func (c *stubGatewayCache) SetReasoningContent(_ context.Context, _ string, _ string, _ time.Duration) error {
  647. return nil
  648. }
  649. func (c *stubGatewayCache) GetReasoningContent(_ context.Context, _ string) (string, error) {
  650. return "", ErrReasoningContentNotFound
  651. }
  652. func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) {
  653. now := time.Now()
  654. resetAt := now.Add(10 * time.Minute)
  655. groupID := int64(1)
  656. rateLimited := Account{
  657. ID: 1,
  658. Platform: PlatformOpenAI,
  659. Type: AccountTypeAPIKey,
  660. Status: StatusActive,
  661. Schedulable: true,
  662. Concurrency: 1,
  663. Priority: 0,
  664. RateLimitResetAt: &resetAt,
  665. }
  666. available := Account{
  667. ID: 2,
  668. Platform: PlatformOpenAI,
  669. Type: AccountTypeAPIKey,
  670. Status: StatusActive,
  671. Schedulable: true,
  672. Concurrency: 1,
  673. Priority: 1,
  674. }
  675. svc := &OpenAIGatewayService{
  676. accountRepo: stubOpenAIAccountRepo{accounts: []Account{rateLimited, available}},
  677. concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
  678. }
  679. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.2", nil)
  680. if err != nil {
  681. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  682. }
  683. if selection == nil || selection.Account == nil {
  684. t.Fatalf("expected selection with account")
  685. }
  686. if selection.Account.ID != available.ID {
  687. t.Fatalf("expected account %d, got %d", available.ID, selection.Account.ID)
  688. }
  689. if selection.ReleaseFunc != nil {
  690. selection.ReleaseFunc()
  691. }
  692. }
  693. func TestOpenAISelectAccountWithLoadAwareness_ImageRateLimitSkipsOnlyImageRequests(t *testing.T) {
  694. future := time.Now().Add(10 * time.Minute).Format(time.RFC3339)
  695. groupID := int64(1)
  696. imageLimited := Account{
  697. ID: 1,
  698. Platform: PlatformOpenAI,
  699. Type: AccountTypeAPIKey,
  700. Status: StatusActive,
  701. Schedulable: true,
  702. Concurrency: 1,
  703. Priority: 0,
  704. Extra: map[string]any{
  705. modelRateLimitsKey: map[string]any{
  706. openAIImageGenerationRateLimitKey: map[string]any{
  707. "rate_limit_reset_at": future,
  708. },
  709. },
  710. },
  711. }
  712. available := Account{
  713. ID: 2,
  714. Platform: PlatformOpenAI,
  715. Type: AccountTypeAPIKey,
  716. Status: StatusActive,
  717. Schedulable: true,
  718. Concurrency: 1,
  719. Priority: 1,
  720. }
  721. svc := &OpenAIGatewayService{
  722. accountRepo: stubOpenAIAccountRepo{accounts: []Account{imageLimited, available}},
  723. concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
  724. }
  725. imageSelection, err := svc.SelectAccountWithLoadAwareness(WithOpenAIImageGenerationIntent(context.Background()), &groupID, "", "gpt-5.4", nil)
  726. require.NoError(t, err)
  727. require.NotNil(t, imageSelection)
  728. require.Equal(t, available.ID, imageSelection.Account.ID)
  729. if imageSelection.ReleaseFunc != nil {
  730. imageSelection.ReleaseFunc()
  731. }
  732. textSelection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.4", nil)
  733. require.NoError(t, err)
  734. require.NotNil(t, textSelection)
  735. require.Equal(t, imageLimited.ID, textSelection.Account.ID)
  736. if textSelection.ReleaseFunc != nil {
  737. textSelection.ReleaseFunc()
  738. }
  739. }
  740. func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulableWhenNoConcurrencyService(t *testing.T) {
  741. now := time.Now()
  742. resetAt := now.Add(10 * time.Minute)
  743. groupID := int64(1)
  744. rateLimited := Account{
  745. ID: 1,
  746. Platform: PlatformOpenAI,
  747. Type: AccountTypeAPIKey,
  748. Status: StatusActive,
  749. Schedulable: true,
  750. Concurrency: 1,
  751. Priority: 0,
  752. RateLimitResetAt: &resetAt,
  753. }
  754. available := Account{
  755. ID: 2,
  756. Platform: PlatformOpenAI,
  757. Type: AccountTypeAPIKey,
  758. Status: StatusActive,
  759. Schedulable: true,
  760. Concurrency: 1,
  761. Priority: 1,
  762. }
  763. svc := &OpenAIGatewayService{
  764. accountRepo: stubOpenAIAccountRepo{accounts: []Account{rateLimited, available}},
  765. // concurrencyService is nil, forcing the non-load-batch selection path.
  766. }
  767. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.2", nil)
  768. if err != nil {
  769. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  770. }
  771. if selection == nil || selection.Account == nil {
  772. t.Fatalf("expected selection with account")
  773. }
  774. if selection.Account.ID != available.ID {
  775. t.Fatalf("expected account %d, got %d", available.ID, selection.Account.ID)
  776. }
  777. if selection.ReleaseFunc != nil {
  778. selection.ReleaseFunc()
  779. }
  780. }
  781. func TestOpenAISelectAccountForModelWithExclusions_StickyUnschedulableClearsSession(t *testing.T) {
  782. sessionHash := "session-1"
  783. repo := stubOpenAIAccountRepo{
  784. accounts: []Account{
  785. {ID: 1, Platform: PlatformOpenAI, Status: StatusDisabled, Schedulable: true, Concurrency: 1},
  786. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
  787. },
  788. }
  789. cache := &stubGatewayCache{
  790. sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
  791. }
  792. svc := &OpenAIGatewayService{
  793. accountRepo: repo,
  794. cache: cache,
  795. }
  796. acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", nil)
  797. if err != nil {
  798. t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
  799. }
  800. if acc == nil || acc.ID != 2 {
  801. t.Fatalf("expected account 2, got %+v", acc)
  802. }
  803. if cache.deletedSessions["openai:"+sessionHash] != 1 {
  804. t.Fatalf("expected sticky session to be deleted")
  805. }
  806. if cache.sessionBindings["openai:"+sessionHash] != 2 {
  807. t.Fatalf("expected sticky session to bind to account 2")
  808. }
  809. }
  810. func TestOpenAISelectAccountForModelWithExclusions_StickyOutsideGroupClearsSession(t *testing.T) {
  811. sessionHash := "session-outside-group"
  812. groupID := int64(1001)
  813. repo := groupAwareStubOpenAIAccountRepo{
  814. stubOpenAIAccountRepo{
  815. accounts: []Account{
  816. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
  817. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, AccountGroups: []AccountGroup{{GroupID: groupID}}},
  818. },
  819. },
  820. }
  821. cache := &stubGatewayCache{
  822. sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
  823. }
  824. svc := &OpenAIGatewayService{
  825. accountRepo: repo,
  826. cache: cache,
  827. }
  828. acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), &groupID, sessionHash, "gpt-4", nil)
  829. if err != nil {
  830. t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
  831. }
  832. if acc == nil || acc.ID != 2 {
  833. t.Fatalf("expected account 2, got %+v", acc)
  834. }
  835. if cache.deletedSessions["openai:"+sessionHash] != 1 {
  836. t.Fatalf("expected sticky session to be deleted")
  837. }
  838. if cache.sessionBindings["openai:"+sessionHash] != 2 {
  839. t.Fatalf("expected sticky session to bind to account 2")
  840. }
  841. }
  842. func TestOpenAISelectAccountWithLoadAwareness_StickyUnschedulableClearsSession(t *testing.T) {
  843. sessionHash := "session-2"
  844. groupID := int64(1)
  845. repo := stubOpenAIAccountRepo{
  846. accounts: []Account{
  847. {ID: 1, Platform: PlatformOpenAI, Status: StatusDisabled, Schedulable: true, Concurrency: 1},
  848. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
  849. },
  850. }
  851. cache := &stubGatewayCache{
  852. sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
  853. }
  854. svc := &OpenAIGatewayService{
  855. accountRepo: repo,
  856. cache: cache,
  857. concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
  858. }
  859. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
  860. if err != nil {
  861. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  862. }
  863. if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
  864. t.Fatalf("expected account 2, got %+v", selection)
  865. }
  866. if cache.deletedSessions["openai:"+sessionHash] != 1 {
  867. t.Fatalf("expected sticky session to be deleted")
  868. }
  869. if cache.sessionBindings["openai:"+sessionHash] != 2 {
  870. t.Fatalf("expected sticky session to bind to account 2")
  871. }
  872. if selection.ReleaseFunc != nil {
  873. selection.ReleaseFunc()
  874. }
  875. }
  876. func TestOpenAISelectAccountWithLoadAwareness_StickyOutsideGroupClearsSession(t *testing.T) {
  877. sessionHash := "session-load-outside-group"
  878. groupID := int64(1002)
  879. repo := groupAwareStubOpenAIAccountRepo{
  880. stubOpenAIAccountRepo{
  881. accounts: []Account{
  882. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
  883. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, AccountGroups: []AccountGroup{{GroupID: groupID}}},
  884. },
  885. },
  886. }
  887. cache := &stubGatewayCache{
  888. sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
  889. }
  890. svc := &OpenAIGatewayService{
  891. accountRepo: repo,
  892. cache: cache,
  893. concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
  894. }
  895. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
  896. if err != nil {
  897. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  898. }
  899. if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
  900. t.Fatalf("expected account 2, got %+v", selection)
  901. }
  902. if cache.deletedSessions["openai:"+sessionHash] != 1 {
  903. t.Fatalf("expected sticky session to be deleted")
  904. }
  905. if cache.sessionBindings["openai:"+sessionHash] != 2 {
  906. t.Fatalf("expected sticky session to bind to account 2")
  907. }
  908. if selection.ReleaseFunc != nil {
  909. selection.ReleaseFunc()
  910. }
  911. }
  912. func TestOpenAISelectAccountForModelWithExclusions_NoModelSupport(t *testing.T) {
  913. repo := stubOpenAIAccountRepo{
  914. accounts: []Account{
  915. {
  916. ID: 1,
  917. Platform: PlatformOpenAI,
  918. Status: StatusActive,
  919. Schedulable: true,
  920. Credentials: map[string]any{"model_mapping": map[string]any{"gpt-3.5-turbo": "gpt-3.5-turbo"}},
  921. },
  922. },
  923. }
  924. cache := &stubGatewayCache{}
  925. svc := &OpenAIGatewayService{
  926. accountRepo: repo,
  927. cache: cache,
  928. }
  929. acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, "", "gpt-4", nil)
  930. if err == nil {
  931. t.Fatalf("expected error for unsupported model")
  932. }
  933. if acc != nil {
  934. t.Fatalf("expected nil account for unsupported model")
  935. }
  936. if !strings.Contains(err.Error(), "supporting model") {
  937. t.Fatalf("unexpected error: %v", err)
  938. }
  939. }
  940. func TestOpenAISelectAccountWithLoadAwareness_LoadBatchErrorFallback(t *testing.T) {
  941. groupID := int64(1)
  942. repo := stubOpenAIAccountRepo{
  943. accounts: []Account{
  944. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 2},
  945. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  946. },
  947. }
  948. cache := &stubGatewayCache{}
  949. concurrencyCache := stubConcurrencyCache{
  950. loadBatchErr: errors.New("load batch failed"),
  951. }
  952. svc := &OpenAIGatewayService{
  953. accountRepo: repo,
  954. cache: cache,
  955. concurrencyService: NewConcurrencyService(concurrencyCache),
  956. }
  957. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "fallback", "gpt-4", nil)
  958. if err != nil {
  959. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  960. }
  961. if selection == nil || selection.Account == nil {
  962. t.Fatalf("expected selection")
  963. }
  964. if selection.Account.ID != 2 {
  965. t.Fatalf("expected account 2, got %d", selection.Account.ID)
  966. }
  967. if cache.sessionBindings["openai:fallback"] != 2 {
  968. t.Fatalf("expected sticky session updated")
  969. }
  970. if selection.ReleaseFunc != nil {
  971. selection.ReleaseFunc()
  972. }
  973. }
  974. func TestOpenAISelectAccountWithLoadAwareness_NoSlotFallbackWait(t *testing.T) {
  975. groupID := int64(1)
  976. repo := stubOpenAIAccountRepo{
  977. accounts: []Account{
  978. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  979. },
  980. }
  981. cache := &stubGatewayCache{}
  982. concurrencyCache := stubConcurrencyCache{
  983. acquireResults: map[int64]bool{1: false},
  984. loadMap: map[int64]*AccountLoadInfo{
  985. 1: {AccountID: 1, LoadRate: 10},
  986. },
  987. }
  988. svc := &OpenAIGatewayService{
  989. accountRepo: repo,
  990. cache: cache,
  991. concurrencyService: NewConcurrencyService(concurrencyCache),
  992. }
  993. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
  994. if err != nil {
  995. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  996. }
  997. if selection == nil || selection.WaitPlan == nil {
  998. t.Fatalf("expected wait plan fallback")
  999. }
  1000. if selection.Account == nil || selection.Account.ID != 1 {
  1001. t.Fatalf("expected account 1")
  1002. }
  1003. }
  1004. func TestOpenAISelectAccountForModelWithExclusions_SetsStickyBinding(t *testing.T) {
  1005. sessionHash := "bind"
  1006. repo := stubOpenAIAccountRepo{
  1007. accounts: []Account{
  1008. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1009. },
  1010. }
  1011. cache := &stubGatewayCache{}
  1012. svc := &OpenAIGatewayService{
  1013. accountRepo: repo,
  1014. cache: cache,
  1015. }
  1016. acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", nil)
  1017. if err != nil {
  1018. t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
  1019. }
  1020. if acc == nil || acc.ID != 1 {
  1021. t.Fatalf("expected account 1")
  1022. }
  1023. if cache.sessionBindings["openai:"+sessionHash] != 1 {
  1024. t.Fatalf("expected sticky session binding")
  1025. }
  1026. }
  1027. func TestOpenAISelectAccountWithLoadAwareness_StickyWaitPlan(t *testing.T) {
  1028. sessionHash := "sticky-wait"
  1029. groupID := int64(1)
  1030. repo := stubOpenAIAccountRepo{
  1031. accounts: []Account{
  1032. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1033. },
  1034. }
  1035. cache := &stubGatewayCache{
  1036. sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
  1037. }
  1038. concurrencyCache := stubConcurrencyCache{
  1039. acquireResults: map[int64]bool{1: false},
  1040. waitCounts: map[int64]int{1: 0},
  1041. }
  1042. svc := &OpenAIGatewayService{
  1043. accountRepo: repo,
  1044. cache: cache,
  1045. concurrencyService: NewConcurrencyService(concurrencyCache),
  1046. }
  1047. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
  1048. if err != nil {
  1049. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  1050. }
  1051. if selection == nil || selection.WaitPlan == nil {
  1052. t.Fatalf("expected sticky wait plan")
  1053. }
  1054. if selection.Account == nil || selection.Account.ID != 1 {
  1055. t.Fatalf("expected account 1")
  1056. }
  1057. }
  1058. func TestOpenAISelectAccountWithLoadAwareness_StickyCapacitySpilloverKeepsBinding(t *testing.T) {
  1059. sessionHash := "sticky-spillover"
  1060. groupID := int64(1)
  1061. repo := stubOpenAIAccountRepo{
  1062. accounts: []Account{
  1063. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 6, Priority: 1, GroupIDs: []int64{groupID}},
  1064. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 6, Priority: 1, GroupIDs: []int64{groupID}},
  1065. },
  1066. }
  1067. cache := &stubGatewayCache{
  1068. sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
  1069. }
  1070. concurrencyCache := stubConcurrencyCache{
  1071. acquireResults: map[int64]bool{1: false, 2: true},
  1072. waitCounts: map[int64]int{1: 1},
  1073. loadMap: map[int64]*AccountLoadInfo{
  1074. 1: {AccountID: 1, LoadRate: 100},
  1075. 2: {AccountID: 2, LoadRate: 10},
  1076. },
  1077. }
  1078. cfg := &config.Config{RunMode: config.RunModeStandard}
  1079. cfg.Gateway.Scheduling.LoadBatchEnabled = true
  1080. cfg.Gateway.Scheduling.StickySessionMaxWaiting = 1
  1081. svc := &OpenAIGatewayService{
  1082. accountRepo: repo,
  1083. cache: cache,
  1084. cfg: cfg,
  1085. concurrencyService: NewConcurrencyService(concurrencyCache),
  1086. }
  1087. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
  1088. require.NoError(t, err)
  1089. require.NotNil(t, selection)
  1090. require.NotNil(t, selection.Account)
  1091. require.Equal(t, int64(2), selection.Account.ID, "capacity spillover should use the other account for this request")
  1092. require.True(t, selection.Acquired)
  1093. require.Equal(t, int64(1), cache.sessionBindings["openai:"+sessionHash], "capacity spillover must not migrate the durable sticky binding")
  1094. if selection.ReleaseFunc != nil {
  1095. selection.ReleaseFunc()
  1096. }
  1097. }
  1098. func TestOpenAISelectAccountWithLoadAwareness_PrefersLowerLoad(t *testing.T) {
  1099. groupID := int64(1)
  1100. repo := stubOpenAIAccountRepo{
  1101. accounts: []Account{
  1102. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1103. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1104. },
  1105. }
  1106. cache := &stubGatewayCache{}
  1107. concurrencyCache := stubConcurrencyCache{
  1108. loadMap: map[int64]*AccountLoadInfo{
  1109. 1: {AccountID: 1, LoadRate: 80},
  1110. 2: {AccountID: 2, LoadRate: 10},
  1111. },
  1112. }
  1113. svc := &OpenAIGatewayService{
  1114. accountRepo: repo,
  1115. cache: cache,
  1116. concurrencyService: NewConcurrencyService(concurrencyCache),
  1117. }
  1118. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "load", "gpt-4", nil)
  1119. if err != nil {
  1120. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  1121. }
  1122. if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
  1123. t.Fatalf("expected account 2")
  1124. }
  1125. if cache.sessionBindings["openai:load"] != 2 {
  1126. t.Fatalf("expected sticky session updated")
  1127. }
  1128. }
  1129. func TestOpenAISelectAccountForModelWithExclusions_StickyExcludedFallback(t *testing.T) {
  1130. sessionHash := "excluded"
  1131. repo := stubOpenAIAccountRepo{
  1132. accounts: []Account{
  1133. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1134. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 2},
  1135. },
  1136. }
  1137. cache := &stubGatewayCache{
  1138. sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
  1139. }
  1140. svc := &OpenAIGatewayService{
  1141. accountRepo: repo,
  1142. cache: cache,
  1143. }
  1144. excluded := map[int64]struct{}{1: {}}
  1145. acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", excluded)
  1146. if err != nil {
  1147. t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
  1148. }
  1149. if acc == nil || acc.ID != 2 {
  1150. t.Fatalf("expected account 2")
  1151. }
  1152. }
  1153. func TestOpenAISelectAccountForModelWithExclusions_StickyNonOpenAI(t *testing.T) {
  1154. sessionHash := "non-openai"
  1155. repo := stubOpenAIAccountRepo{
  1156. accounts: []Account{
  1157. {ID: 1, Platform: PlatformAnthropic, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1158. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 2},
  1159. },
  1160. }
  1161. cache := &stubGatewayCache{
  1162. sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
  1163. }
  1164. svc := &OpenAIGatewayService{
  1165. accountRepo: repo,
  1166. cache: cache,
  1167. }
  1168. acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", nil)
  1169. if err != nil {
  1170. t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
  1171. }
  1172. if acc == nil || acc.ID != 2 {
  1173. t.Fatalf("expected account 2")
  1174. }
  1175. }
  1176. func TestOpenAISelectAccountForModelWithExclusions_NoAccounts(t *testing.T) {
  1177. repo := stubOpenAIAccountRepo{accounts: []Account{}}
  1178. cache := &stubGatewayCache{}
  1179. svc := &OpenAIGatewayService{
  1180. accountRepo: repo,
  1181. cache: cache,
  1182. }
  1183. acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, "", "", nil)
  1184. if err == nil {
  1185. t.Fatalf("expected error for no accounts")
  1186. }
  1187. if acc != nil {
  1188. t.Fatalf("expected nil account")
  1189. }
  1190. if !strings.Contains(err.Error(), "no available OpenAI accounts") {
  1191. t.Fatalf("unexpected error: %v", err)
  1192. }
  1193. }
  1194. func TestOpenAISelectAccountWithLoadAwareness_NoCandidates(t *testing.T) {
  1195. groupID := int64(1)
  1196. resetAt := time.Now().Add(1 * time.Hour)
  1197. repo := stubOpenAIAccountRepo{
  1198. accounts: []Account{
  1199. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, RateLimitResetAt: &resetAt},
  1200. },
  1201. }
  1202. cache := &stubGatewayCache{}
  1203. concurrencyCache := stubConcurrencyCache{}
  1204. svc := &OpenAIGatewayService{
  1205. accountRepo: repo,
  1206. cache: cache,
  1207. concurrencyService: NewConcurrencyService(concurrencyCache),
  1208. }
  1209. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
  1210. if err == nil {
  1211. t.Fatalf("expected error for no candidates")
  1212. }
  1213. if selection != nil {
  1214. t.Fatalf("expected nil selection")
  1215. }
  1216. }
  1217. func TestOpenAISelectAccountWithLoadAwareness_AllFullWaitPlan(t *testing.T) {
  1218. groupID := int64(1)
  1219. repo := stubOpenAIAccountRepo{
  1220. accounts: []Account{
  1221. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1222. },
  1223. }
  1224. cache := &stubGatewayCache{}
  1225. concurrencyCache := stubConcurrencyCache{
  1226. loadMap: map[int64]*AccountLoadInfo{
  1227. 1: {AccountID: 1, LoadRate: 100},
  1228. },
  1229. }
  1230. svc := &OpenAIGatewayService{
  1231. accountRepo: repo,
  1232. cache: cache,
  1233. concurrencyService: NewConcurrencyService(concurrencyCache),
  1234. }
  1235. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
  1236. if err != nil {
  1237. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  1238. }
  1239. if selection == nil || selection.WaitPlan == nil {
  1240. t.Fatalf("expected wait plan")
  1241. }
  1242. }
  1243. func TestOpenAISelectAccountWithLoadAwareness_LoadBatchErrorNoAcquire(t *testing.T) {
  1244. groupID := int64(1)
  1245. repo := stubOpenAIAccountRepo{
  1246. accounts: []Account{
  1247. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1248. },
  1249. }
  1250. cache := &stubGatewayCache{}
  1251. concurrencyCache := stubConcurrencyCache{
  1252. loadBatchErr: errors.New("load batch failed"),
  1253. acquireResults: map[int64]bool{1: false},
  1254. }
  1255. svc := &OpenAIGatewayService{
  1256. accountRepo: repo,
  1257. cache: cache,
  1258. concurrencyService: NewConcurrencyService(concurrencyCache),
  1259. }
  1260. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
  1261. if err != nil {
  1262. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  1263. }
  1264. if selection == nil || selection.WaitPlan == nil {
  1265. t.Fatalf("expected wait plan")
  1266. }
  1267. }
  1268. func TestOpenAISelectAccountWithLoadAwareness_MissingLoadInfo(t *testing.T) {
  1269. groupID := int64(1)
  1270. repo := stubOpenAIAccountRepo{
  1271. accounts: []Account{
  1272. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1273. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1274. },
  1275. }
  1276. cache := &stubGatewayCache{}
  1277. concurrencyCache := stubConcurrencyCache{
  1278. loadMap: map[int64]*AccountLoadInfo{
  1279. 1: {AccountID: 1, LoadRate: 50},
  1280. },
  1281. skipDefaultLoad: true,
  1282. }
  1283. svc := &OpenAIGatewayService{
  1284. accountRepo: repo,
  1285. cache: cache,
  1286. concurrencyService: NewConcurrencyService(concurrencyCache),
  1287. }
  1288. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
  1289. if err != nil {
  1290. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  1291. }
  1292. if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
  1293. t.Fatalf("expected account 2")
  1294. }
  1295. }
  1296. func TestOpenAISelectAccountForModelWithExclusions_LeastRecentlyUsed(t *testing.T) {
  1297. oldTime := time.Now().Add(-2 * time.Hour)
  1298. newTime := time.Now().Add(-1 * time.Hour)
  1299. repo := stubOpenAIAccountRepo{
  1300. accounts: []Account{
  1301. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Priority: 1, LastUsedAt: &newTime},
  1302. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Priority: 1, LastUsedAt: &oldTime},
  1303. },
  1304. }
  1305. cache := &stubGatewayCache{}
  1306. svc := &OpenAIGatewayService{
  1307. accountRepo: repo,
  1308. cache: cache,
  1309. }
  1310. acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, "", "gpt-4", nil)
  1311. if err != nil {
  1312. t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
  1313. }
  1314. if acc == nil || acc.ID != 2 {
  1315. t.Fatalf("expected account 2")
  1316. }
  1317. }
  1318. func TestOpenAISelectAccountWithLoadAwareness_PreferNeverUsed(t *testing.T) {
  1319. groupID := int64(1)
  1320. lastUsed := time.Now().Add(-1 * time.Hour)
  1321. repo := stubOpenAIAccountRepo{
  1322. accounts: []Account{
  1323. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, LastUsedAt: &lastUsed},
  1324. {ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
  1325. },
  1326. }
  1327. cache := &stubGatewayCache{}
  1328. concurrencyCache := stubConcurrencyCache{
  1329. loadMap: map[int64]*AccountLoadInfo{
  1330. 1: {AccountID: 1, LoadRate: 10},
  1331. 2: {AccountID: 2, LoadRate: 10},
  1332. },
  1333. }
  1334. svc := &OpenAIGatewayService{
  1335. accountRepo: repo,
  1336. cache: cache,
  1337. concurrencyService: NewConcurrencyService(concurrencyCache),
  1338. }
  1339. selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
  1340. if err != nil {
  1341. t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
  1342. }
  1343. if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
  1344. t.Fatalf("expected account 2")
  1345. }
  1346. }
  1347. func TestOpenAIStreamingTimeout(t *testing.T) {
  1348. gin.SetMode(gin.TestMode)
  1349. cfg := &config.Config{
  1350. Gateway: config.GatewayConfig{
  1351. StreamDataIntervalTimeout: 1,
  1352. StreamKeepaliveInterval: 0,
  1353. MaxLineSize: defaultMaxLineSize,
  1354. },
  1355. }
  1356. svc := &OpenAIGatewayService{cfg: cfg}
  1357. rec := httptest.NewRecorder()
  1358. c, _ := gin.CreateTestContext(rec)
  1359. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1360. pr, pw := io.Pipe()
  1361. resp := &http.Response{
  1362. StatusCode: http.StatusOK,
  1363. Body: pr,
  1364. Header: http.Header{},
  1365. }
  1366. start := time.Now()
  1367. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, start, "model", "model")
  1368. _ = pw.Close()
  1369. _ = pr.Close()
  1370. if err == nil || !strings.Contains(err.Error(), "stream data interval timeout") {
  1371. t.Fatalf("expected stream timeout error, got %v", err)
  1372. }
  1373. if !strings.Contains(rec.Body.String(), "\"type\":\"error\"") || !strings.Contains(rec.Body.String(), "stream_timeout") {
  1374. t.Fatalf("expected OpenAI-compatible error SSE event, got %q", rec.Body.String())
  1375. }
  1376. }
  1377. func TestOpenAIStreamingContextCanceledReturnsIncompleteErrorWithoutInjectingErrorEvent(t *testing.T) {
  1378. gin.SetMode(gin.TestMode)
  1379. cfg := &config.Config{
  1380. Gateway: config.GatewayConfig{
  1381. StreamDataIntervalTimeout: 0,
  1382. StreamKeepaliveInterval: 0,
  1383. MaxLineSize: defaultMaxLineSize,
  1384. },
  1385. }
  1386. svc := &OpenAIGatewayService{cfg: cfg}
  1387. rec := httptest.NewRecorder()
  1388. c, _ := gin.CreateTestContext(rec)
  1389. ctx, cancel := context.WithCancel(context.Background())
  1390. cancel()
  1391. c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
  1392. resp := &http.Response{
  1393. StatusCode: http.StatusOK,
  1394. Body: cancelReadCloser{},
  1395. Header: http.Header{},
  1396. }
  1397. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
  1398. if err == nil || !strings.Contains(err.Error(), "stream usage incomplete") {
  1399. t.Fatalf("expected incomplete stream error, got %v", err)
  1400. }
  1401. if strings.Contains(rec.Body.String(), "event: error") || strings.Contains(rec.Body.String(), "stream_read_error") {
  1402. t.Fatalf("expected no injected SSE error event, got %q", rec.Body.String())
  1403. }
  1404. }
  1405. func TestOpenAIStreamingReadErrorBeforeOutputReturnsFailover(t *testing.T) {
  1406. gin.SetMode(gin.TestMode)
  1407. cfg := &config.Config{
  1408. Gateway: config.GatewayConfig{
  1409. StreamDataIntervalTimeout: 0,
  1410. StreamKeepaliveInterval: 0,
  1411. MaxLineSize: defaultMaxLineSize,
  1412. },
  1413. }
  1414. svc := &OpenAIGatewayService{cfg: cfg}
  1415. rec := httptest.NewRecorder()
  1416. c, _ := gin.CreateTestContext(rec)
  1417. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1418. resp := &http.Response{
  1419. StatusCode: http.StatusOK,
  1420. Body: errReadCloser{err: io.ErrUnexpectedEOF},
  1421. Header: http.Header{"X-Request-Id": []string{"rid-disconnect"}},
  1422. }
  1423. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  1424. require.Error(t, err)
  1425. var failoverErr *UpstreamFailoverError
  1426. require.ErrorAs(t, err, &failoverErr)
  1427. require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
  1428. require.False(t, c.Writer.Written())
  1429. require.Empty(t, rec.Body.String())
  1430. }
  1431. func TestOpenAIStreamingPostOutputDisconnectQuarantinesSharedProxyWithoutSameStreamFailover(t *testing.T) {
  1432. gin.SetMode(gin.TestMode)
  1433. proxyID := int64(4698)
  1434. account := &Account{
  1435. ID: 469801,
  1436. Name: "oauth-on-shared-proxy",
  1437. Platform: PlatformOpenAI,
  1438. Type: AccountTypeOAuth,
  1439. ProxyID: &proxyID,
  1440. }
  1441. svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{
  1442. MaxLineSize: defaultMaxLineSize,
  1443. }}}
  1444. // collapseInterval 0: the two loop iterations below record within the
  1445. // production collapse window and must count as distinct failure events here.
  1446. svc.openaiProxyStreamCircuit = newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
  1447. failureThreshold: 2,
  1448. failureWindow: time.Minute,
  1449. quarantineTTL: 10 * time.Minute,
  1450. maxEntries: 16,
  1451. })
  1452. for _, readErr := range []error{
  1453. io.ErrUnexpectedEOF,
  1454. errors.New("http2: client connection lost"),
  1455. } {
  1456. rec := httptest.NewRecorder()
  1457. c, _ := gin.CreateTestContext(rec)
  1458. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  1459. resp := &http.Response{
  1460. StatusCode: http.StatusOK,
  1461. Body: &openAIStreamReadThenErrorCloser{
  1462. reader: strings.NewReader(strings.Join([]string{
  1463. "event: response.output_text.delta",
  1464. `data: {"type":"response.output_text.delta","delta":"partial"}`,
  1465. "",
  1466. }, "\n")),
  1467. err: readErr,
  1468. },
  1469. Header: http.Header{"X-Request-Id": []string{"rid-proxy-disconnect"}},
  1470. }
  1471. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "gpt-5.6-sol", "gpt-5.6-sol")
  1472. require.Error(t, err)
  1473. var failoverErr *UpstreamFailoverError
  1474. require.False(t, errors.As(err, &failoverErr), "post-output disconnect must not fail over inside the same stream")
  1475. require.Contains(t, rec.Body.String(), "partial")
  1476. }
  1477. scheduler := &defaultOpenAIAccountScheduler{service: svc}
  1478. compatible, reason := scheduler.isAccountRequestCompatibleReason(context.Background(), account, OpenAIAccountScheduleRequest{})
  1479. require.False(t, compatible, "the next request must exclude accounts sharing the quarantined proxy")
  1480. require.Equal(t, "proxy_stream_quarantined", reason)
  1481. }
  1482. func TestOpenAIStreamingTerminalAndClientCancellationDoNotQuarantineProxy(t *testing.T) {
  1483. gin.SetMode(gin.TestMode)
  1484. proxyID := int64(4699)
  1485. account := &Account{ID: 469901, Name: "oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, ProxyID: &proxyID}
  1486. svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}}
  1487. terminalRecorder := httptest.NewRecorder()
  1488. terminalCtx, _ := gin.CreateTestContext(terminalRecorder)
  1489. terminalCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  1490. terminalResp := &http.Response{
  1491. StatusCode: http.StatusOK,
  1492. Body: &openAIStreamReadThenErrorCloser{
  1493. reader: strings.NewReader(strings.Join([]string{
  1494. "event: response.completed",
  1495. `data: {"type":"response.completed","response":{"status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":3,"total_tokens":8}}}`,
  1496. "",
  1497. }, "\n")),
  1498. err: io.ErrUnexpectedEOF,
  1499. },
  1500. Header: http.Header{},
  1501. }
  1502. _, err := svc.handleStreamingResponse(terminalCtx.Request.Context(), terminalResp, terminalCtx, account, time.Now(), "model", "model")
  1503. require.NoError(t, err)
  1504. for range 2 {
  1505. rec := httptest.NewRecorder()
  1506. c, _ := gin.CreateTestContext(rec)
  1507. ctx, cancel := context.WithCancel(context.Background())
  1508. cancel()
  1509. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(ctx)
  1510. resp := &http.Response{
  1511. StatusCode: http.StatusOK,
  1512. Body: &openAIStreamReadThenErrorCloser{
  1513. reader: strings.NewReader("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"),
  1514. err: context.Canceled,
  1515. },
  1516. Header: http.Header{},
  1517. }
  1518. _, err = svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
  1519. require.Error(t, err)
  1520. }
  1521. scheduler := &defaultOpenAIAccountScheduler{service: svc}
  1522. compatible, reason := scheduler.isAccountRequestCompatibleReason(context.Background(), account, OpenAIAccountScheduleRequest{})
  1523. require.True(t, compatible)
  1524. require.Empty(t, reason)
  1525. }
  1526. func TestOpenAIStreamingResponseFailedBeforeOutputReturnsFailover(t *testing.T) {
  1527. gin.SetMode(gin.TestMode)
  1528. cfg := &config.Config{
  1529. Gateway: config.GatewayConfig{
  1530. StreamDataIntervalTimeout: 0,
  1531. StreamKeepaliveInterval: 0,
  1532. MaxLineSize: defaultMaxLineSize,
  1533. },
  1534. }
  1535. svc := &OpenAIGatewayService{cfg: cfg}
  1536. rec := httptest.NewRecorder()
  1537. c, _ := gin.CreateTestContext(rec)
  1538. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1539. resp := &http.Response{
  1540. StatusCode: http.StatusOK,
  1541. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1542. "event: response.created",
  1543. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  1544. "",
  1545. "event: response.in_progress",
  1546. `data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
  1547. "",
  1548. "event: response.failed",
  1549. `data: {"type":"response.failed","error":{"message":"An error occurred while processing your request."}}`,
  1550. "",
  1551. }, "\n"))),
  1552. Header: http.Header{"X-Request-Id": []string{"rid-failed"}},
  1553. }
  1554. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  1555. require.Error(t, err)
  1556. var failoverErr *UpstreamFailoverError
  1557. require.ErrorAs(t, err, &failoverErr)
  1558. require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
  1559. require.False(t, failoverErr.RetryableOnSameAccount)
  1560. require.Contains(t, string(failoverErr.ResponseBody), "An error occurred while processing your request")
  1561. require.False(t, c.Writer.Written())
  1562. require.Empty(t, rec.Body.String())
  1563. }
  1564. func TestOpenAIStreamingResponseFailedBeforeOutputCapacityErrorReturnsFailover(t *testing.T) {
  1565. gin.SetMode(gin.TestMode)
  1566. cfg := &config.Config{
  1567. Gateway: config.GatewayConfig{
  1568. StreamDataIntervalTimeout: 0,
  1569. StreamKeepaliveInterval: 0,
  1570. MaxLineSize: defaultMaxLineSize,
  1571. },
  1572. }
  1573. svc := &OpenAIGatewayService{cfg: cfg}
  1574. rec := httptest.NewRecorder()
  1575. c, _ := gin.CreateTestContext(rec)
  1576. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1577. resp := &http.Response{
  1578. StatusCode: http.StatusOK,
  1579. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1580. "event: response.created",
  1581. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  1582. "",
  1583. "event: response.in_progress",
  1584. `data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
  1585. "",
  1586. "event: response.failed",
  1587. `data: {"type":"response.failed","error":{"message":"Selected model is at capacity. Please try a different model.","type":"invalid_request_error"}}`,
  1588. "",
  1589. }, "\n"))),
  1590. Header: http.Header{"X-Request-Id": []string{"rid-capacity-failed"}},
  1591. }
  1592. account := &Account{
  1593. ID: 1,
  1594. Platform: PlatformOpenAI,
  1595. Type: AccountTypeAPIKey,
  1596. Name: "pool-account",
  1597. Credentials: map[string]any{
  1598. "pool_mode": true,
  1599. },
  1600. }
  1601. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
  1602. require.Error(t, err)
  1603. var failoverErr *UpstreamFailoverError
  1604. require.ErrorAs(t, err, &failoverErr)
  1605. require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
  1606. require.True(t, failoverErr.RetryableOnSameAccount)
  1607. require.Contains(t, string(failoverErr.ResponseBody), "Selected model is at capacity")
  1608. require.False(t, c.Writer.Written())
  1609. require.Empty(t, rec.Body.String())
  1610. }
  1611. func TestOpenAIStreamingResponseFailedBeforeOutputServerOverloadedCodeReturnsFailover(t *testing.T) {
  1612. gin.SetMode(gin.TestMode)
  1613. cfg := &config.Config{
  1614. Gateway: config.GatewayConfig{
  1615. StreamDataIntervalTimeout: 0,
  1616. StreamKeepaliveInterval: 0,
  1617. MaxLineSize: defaultMaxLineSize,
  1618. },
  1619. }
  1620. svc := &OpenAIGatewayService{cfg: cfg}
  1621. rec := httptest.NewRecorder()
  1622. c, _ := gin.CreateTestContext(rec)
  1623. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1624. resp := &http.Response{
  1625. StatusCode: http.StatusOK,
  1626. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1627. "event: response.created",
  1628. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  1629. "",
  1630. "event: response.failed",
  1631. `data: {"type":"response.failed","response":{"id":"resp_1","error":{"code":"server_is_overloaded","message":"Please retry later."}}}`,
  1632. "",
  1633. }, "\n"))),
  1634. Header: http.Header{"X-Request-Id": []string{"rid-overloaded-failed"}},
  1635. }
  1636. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  1637. require.Error(t, err)
  1638. var failoverErr *UpstreamFailoverError
  1639. require.ErrorAs(t, err, &failoverErr)
  1640. require.Equal(t, http.StatusServiceUnavailable, failoverErr.StatusCode)
  1641. require.Contains(t, string(failoverErr.ResponseBody), "Please retry later")
  1642. // 容量降载是请求级信号:非池模式账号也要先在同账号重试,且不得据此临时封禁账号。
  1643. // 否则单个被降载的请求会把整池账号逐个消耗掉,而降载因素在每个账号上都相同。
  1644. require.True(t, failoverErr.RetryableOnSameAccount)
  1645. require.True(t, failoverErr.RequestScopedTransient)
  1646. require.False(t, c.Writer.Written())
  1647. require.Empty(t, rec.Body.String())
  1648. }
  1649. func TestOpenAIStreamingResponseFailedBeforeOutputRateLimitUsesPoolRetryPolicy(t *testing.T) {
  1650. gin.SetMode(gin.TestMode)
  1651. cfg := &config.Config{
  1652. Gateway: config.GatewayConfig{
  1653. StreamDataIntervalTimeout: 0,
  1654. StreamKeepaliveInterval: 0,
  1655. MaxLineSize: defaultMaxLineSize,
  1656. },
  1657. }
  1658. svc := &OpenAIGatewayService{cfg: cfg}
  1659. rec := httptest.NewRecorder()
  1660. c, _ := gin.CreateTestContext(rec)
  1661. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1662. resp := &http.Response{
  1663. StatusCode: http.StatusOK,
  1664. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1665. "event: response.created",
  1666. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  1667. "",
  1668. "event: response.failed",
  1669. `data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
  1670. "",
  1671. }, "\n"))),
  1672. Header: http.Header{
  1673. "X-Request-Id": []string{"rid-rate-limit-failed"},
  1674. "Retry-After": []string{"1"},
  1675. },
  1676. }
  1677. account := &Account{
  1678. ID: 1,
  1679. Platform: PlatformOpenAI,
  1680. Type: AccountTypeAPIKey,
  1681. Name: "pool-account",
  1682. Credentials: map[string]any{
  1683. "pool_mode": true,
  1684. "pool_mode_retry_count": float64(1),
  1685. "pool_mode_retry_status_codes": []any{float64(http.StatusTooManyRequests)},
  1686. },
  1687. }
  1688. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
  1689. require.Error(t, err)
  1690. var failoverErr *UpstreamFailoverError
  1691. require.ErrorAs(t, err, &failoverErr)
  1692. require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
  1693. require.True(t, failoverErr.RetryableOnSameAccount)
  1694. require.Equal(t, "1", failoverErr.ResponseHeaders.Get("Retry-After"))
  1695. require.Equal(t, "rate_limit_error", gjson.GetBytes(failoverErr.ResponseBody, "error.type").String())
  1696. require.Contains(t, string(failoverErr.ResponseBody), "Concurrency limit exceeded")
  1697. require.False(t, c.Writer.Written())
  1698. require.Empty(t, rec.Body.String())
  1699. opsVal, ok := c.Get(OpsUpstreamErrorsKey)
  1700. require.True(t, ok)
  1701. opsEvents, ok := opsVal.([]*OpsUpstreamErrorEvent)
  1702. require.True(t, ok)
  1703. require.NotEmpty(t, opsEvents)
  1704. require.Equal(t, http.StatusTooManyRequests, opsEvents[len(opsEvents)-1].UpstreamStatusCode)
  1705. }
  1706. // 流内 rate limit 进入 OAuth 同账号重试窗口,但不立即写账号级限流/封禁状态:
  1707. // HTTP 200 流的 x-codex-* 头不能让窗口内的账号提前失去调度资格。
  1708. func TestOpenAIStreamingResponseFailedRateLimitDoesNotBlockAccountScheduling(t *testing.T) {
  1709. gin.SetMode(gin.TestMode)
  1710. cfg := &config.Config{
  1711. Gateway: config.GatewayConfig{
  1712. StreamDataIntervalTimeout: 0,
  1713. StreamKeepaliveInterval: 0,
  1714. MaxLineSize: defaultMaxLineSize,
  1715. },
  1716. }
  1717. svc := &OpenAIGatewayService{cfg: cfg}
  1718. rec := httptest.NewRecorder()
  1719. c, _ := gin.CreateTestContext(rec)
  1720. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1721. resp := &http.Response{
  1722. StatusCode: http.StatusOK,
  1723. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1724. "event: response.created",
  1725. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  1726. "",
  1727. "event: response.failed",
  1728. `data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
  1729. "",
  1730. }, "\n"))),
  1731. Header: http.Header{
  1732. "X-Codex-Primary-Used-Percent": []string{"12"},
  1733. "X-Codex-Primary-Reset-After-Seconds": []string{"604800"},
  1734. "Retry-After": []string{"1"},
  1735. },
  1736. }
  1737. account := &Account{
  1738. ID: 11,
  1739. Platform: PlatformOpenAI,
  1740. Type: AccountTypeOAuth,
  1741. Name: "oauth-account",
  1742. }
  1743. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
  1744. require.Error(t, err)
  1745. var failoverErr *UpstreamFailoverError
  1746. require.ErrorAs(t, err, &failoverErr)
  1747. require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
  1748. require.True(t, failoverErr.RetryableOnSameAccount)
  1749. require.False(t, failoverErr.SameAccountRetryDeadline.IsZero())
  1750. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1751. }
  1752. func TestOpenAIStreamingResponseFailedAfterOutputSanitizesVerboseResponseForClient(t *testing.T) {
  1753. gin.SetMode(gin.TestMode)
  1754. cfg := &config.Config{
  1755. Gateway: config.GatewayConfig{
  1756. StreamDataIntervalTimeout: 0,
  1757. StreamKeepaliveInterval: 0,
  1758. MaxLineSize: defaultMaxLineSize,
  1759. },
  1760. }
  1761. svc := &OpenAIGatewayService{cfg: cfg}
  1762. rec := httptest.NewRecorder()
  1763. c, _ := gin.CreateTestContext(rec)
  1764. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1765. longInstructions := strings.Repeat("You are GPT-5.1 running in the Codex CLI. ", 20)
  1766. failedPayload := fmt.Sprintf(
  1767. `{"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."}}}`,
  1768. longInstructions,
  1769. )
  1770. resp := &http.Response{
  1771. StatusCode: http.StatusOK,
  1772. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1773. "event: response.created",
  1774. `data: {"type":"response.created","response":{"id":"resp_failed"}}`,
  1775. "",
  1776. "event: response.output_text.delta",
  1777. `data: {"type":"response.output_text.delta","delta":"partial"}`,
  1778. "",
  1779. "event: response.failed",
  1780. "data: " + failedPayload,
  1781. "",
  1782. }, "\n"))),
  1783. Header: http.Header{"X-Request-Id": []string{"rid-failed-after-output"}},
  1784. }
  1785. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  1786. require.Error(t, err)
  1787. body := rec.Body.String()
  1788. require.Contains(t, body, "event: response.failed")
  1789. require.Contains(t, body, "context_length_exceeded")
  1790. require.Contains(t, body, `"type":"invalid_request_error"`)
  1791. require.Contains(t, body, "Your input exceeds the context window")
  1792. require.NotContains(t, body, "You are GPT-5.1 running in the Codex CLI")
  1793. require.NotContains(t, body, `"instructions"`)
  1794. require.NotContains(t, body, `"output"`)
  1795. require.NotContains(t, body, `"usage"`)
  1796. }
  1797. func TestOpenAIStreamingContextWindowResponseFailedBeforeOutputPassesThrough(t *testing.T) {
  1798. gin.SetMode(gin.TestMode)
  1799. cfg := &config.Config{
  1800. Gateway: config.GatewayConfig{
  1801. StreamDataIntervalTimeout: 0,
  1802. StreamKeepaliveInterval: 0,
  1803. MaxLineSize: defaultMaxLineSize,
  1804. },
  1805. }
  1806. svc := &OpenAIGatewayService{cfg: cfg}
  1807. rec := httptest.NewRecorder()
  1808. c, _ := gin.CreateTestContext(rec)
  1809. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1810. resp := &http.Response{
  1811. StatusCode: http.StatusOK,
  1812. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1813. "event: response.created",
  1814. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  1815. "",
  1816. "event: response.failed",
  1817. `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}}`,
  1818. "",
  1819. }, "\n"))),
  1820. Header: http.Header{"X-Request-Id": []string{"rid-context-window-failed"}},
  1821. }
  1822. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  1823. require.Error(t, err)
  1824. var failoverErr *UpstreamFailoverError
  1825. require.False(t, errors.As(err, &failoverErr))
  1826. require.True(t, c.Writer.Written())
  1827. require.Contains(t, rec.Body.String(), "response.failed")
  1828. require.Contains(t, rec.Body.String(), `"type":"upstream_error"`)
  1829. require.Contains(t, rec.Body.String(), "Your input exceeds the context window")
  1830. }
  1831. func TestOpenAIStreamingContextWindowResponseFailedBeforeOutputAppliesPassthroughRule(t *testing.T) {
  1832. gin.SetMode(gin.TestMode)
  1833. cfg := &config.Config{
  1834. Gateway: config.GatewayConfig{
  1835. StreamDataIntervalTimeout: 0,
  1836. StreamKeepaliveInterval: 0,
  1837. MaxLineSize: defaultMaxLineSize,
  1838. },
  1839. }
  1840. svc := &OpenAIGatewayService{cfg: cfg}
  1841. rec := httptest.NewRecorder()
  1842. c, _ := gin.CreateTestContext(rec)
  1843. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1844. rule := newNonFailoverPassthroughRule(http.StatusBadRequest, "context_length_exceeded", http.StatusBadRequest, "")
  1845. rule.Platforms = []string{PlatformOpenAI}
  1846. rule.PassthroughBody = true
  1847. rule.CustomMessage = nil
  1848. ruleSvc := &ErrorPassthroughService{}
  1849. ruleSvc.setLocalCache([]*model.ErrorPassthroughRule{rule})
  1850. BindErrorPassthroughService(c, ruleSvc)
  1851. upstreamMessage := "Your input exceeds the context window of this model. Please adjust your input and try again."
  1852. resp := &http.Response{
  1853. StatusCode: http.StatusOK,
  1854. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1855. "event: response.created",
  1856. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  1857. "",
  1858. "event: response.failed",
  1859. `data: {"type":"response.failed","response":{"id":"resp_1","error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"` + upstreamMessage + `"}}}`,
  1860. "",
  1861. }, "\n"))),
  1862. Header: http.Header{"X-Request-Id": []string{"rid-context-window-passthrough-rule"}},
  1863. }
  1864. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  1865. require.Error(t, err)
  1866. var failoverErr *UpstreamFailoverError
  1867. require.False(t, errors.As(err, &failoverErr))
  1868. require.True(t, IsResponseCommitted(c))
  1869. require.Equal(t, http.StatusBadRequest, rec.Code)
  1870. body := rec.Body.String()
  1871. require.Equal(t, "upstream_error", gjson.Get(body, "error.type").String())
  1872. require.Equal(t, upstreamMessage, gjson.Get(body, "error.message").String())
  1873. require.NotContains(t, body, "response.failed")
  1874. require.NotContains(t, body, "Upstream request failed")
  1875. // 命中透传规则也应记录 ops 上游错误事件(对齐 CC/Messages 与 antigravity 先例)。
  1876. opsVal, opsRecorded := c.Get(OpsUpstreamErrorsKey)
  1877. require.True(t, opsRecorded, "passthrough hit should record an ops upstream error event")
  1878. opsEvents, _ := opsVal.([]*OpsUpstreamErrorEvent)
  1879. require.NotEmpty(t, opsEvents)
  1880. }
  1881. func TestOpenAIStreamingPreambleOnlyMissingTerminalReturnsFailover(t *testing.T) {
  1882. gin.SetMode(gin.TestMode)
  1883. cfg := &config.Config{
  1884. Gateway: config.GatewayConfig{
  1885. StreamDataIntervalTimeout: 0,
  1886. StreamKeepaliveInterval: 0,
  1887. MaxLineSize: defaultMaxLineSize,
  1888. },
  1889. }
  1890. svc := &OpenAIGatewayService{cfg: cfg}
  1891. rec := httptest.NewRecorder()
  1892. c, _ := gin.CreateTestContext(rec)
  1893. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1894. resp := &http.Response{
  1895. StatusCode: http.StatusOK,
  1896. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1897. "event: response.created",
  1898. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  1899. "",
  1900. "event: response.in_progress",
  1901. `data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
  1902. "",
  1903. }, "\n"))),
  1904. Header: http.Header{"X-Request-Id": []string{"rid-missing-terminal"}},
  1905. }
  1906. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  1907. require.Error(t, err)
  1908. var failoverErr *UpstreamFailoverError
  1909. require.ErrorAs(t, err, &failoverErr)
  1910. require.False(t, c.Writer.Written())
  1911. require.Empty(t, rec.Body.String())
  1912. }
  1913. func TestOpenAIStreamingPreambleKeepaliveUsesDownstreamIdle(t *testing.T) {
  1914. gin.SetMode(gin.TestMode)
  1915. cfg := &config.Config{
  1916. Gateway: config.GatewayConfig{
  1917. StreamDataIntervalTimeout: 0,
  1918. // Keepalive is based on *downstream* idle time (last flush to client),
  1919. // not upstream event cadence. Interval is seconds (config unit).
  1920. StreamKeepaliveInterval: 1,
  1921. MaxLineSize: defaultMaxLineSize,
  1922. },
  1923. }
  1924. svc := &OpenAIGatewayService{cfg: cfg}
  1925. rec := httptest.NewRecorder()
  1926. c, _ := gin.CreateTestContext(rec)
  1927. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1928. pr, pw := io.Pipe()
  1929. resp := &http.Response{
  1930. StatusCode: http.StatusOK,
  1931. Body: pr,
  1932. Header: http.Header{},
  1933. }
  1934. go func() {
  1935. defer func() { _ = pw.Close() }()
  1936. // Emit preamble/progress quickly so clientOutputStarted is true, then
  1937. // leave a real downstream idle gap longer than keepaliveInterval so the
  1938. // ticker can write ":\n\n". Frequent upstream ticks used to refresh
  1939. // lastDownstreamWriteAt and flake on loaded CI runners.
  1940. _, _ = pw.Write([]byte("data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\n"))
  1941. time.Sleep(50 * time.Millisecond)
  1942. _, _ = pw.Write([]byte("data: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_1\"}}\n\n"))
  1943. time.Sleep(1300 * time.Millisecond)
  1944. _, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}}\n\n"))
  1945. }()
  1946. result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  1947. _ = pr.Close()
  1948. require.NoError(t, err)
  1949. require.NotNil(t, result)
  1950. require.Contains(t, rec.Body.String(), ":\n\n")
  1951. require.Contains(t, rec.Body.String(), "response.completed")
  1952. }
  1953. func TestOpenAIStreamingNormalizesTerminalOutputFromDeltas(t *testing.T) {
  1954. gin.SetMode(gin.TestMode)
  1955. cfg := &config.Config{
  1956. Gateway: config.GatewayConfig{
  1957. StreamDataIntervalTimeout: 0,
  1958. StreamKeepaliveInterval: 0,
  1959. MaxLineSize: defaultMaxLineSize,
  1960. },
  1961. }
  1962. svc := &OpenAIGatewayService{cfg: cfg}
  1963. rec := httptest.NewRecorder()
  1964. c, _ := gin.CreateTestContext(rec)
  1965. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  1966. resp := &http.Response{
  1967. StatusCode: http.StatusOK,
  1968. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1969. `data: {"type":"response.created","response":{"id":"resp_sdk_parse"}}`,
  1970. "",
  1971. `data: {"type":"response.output_text.delta","delta":"pon"}`,
  1972. "",
  1973. `data: {"type":"response.output_text.delta","delta":"g"}`,
  1974. "",
  1975. `data: {"type":"response.completed","response":{"id":"resp_sdk_parse","status":"completed","output":null,"usage":{"input_tokens":1,"output_tokens":1}}}`,
  1976. "",
  1977. }, "\n"))),
  1978. Header: http.Header{"X-Request-Id": []string{"rid-sdk-parse"}},
  1979. }
  1980. result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  1981. require.NoError(t, err)
  1982. require.NotNil(t, result)
  1983. terminalType, terminalPayload, ok := extractOpenAISSETerminalEvent(rec.Body.String())
  1984. require.True(t, ok)
  1985. require.Equal(t, "response.completed", terminalType)
  1986. output := gjson.GetBytes(terminalPayload, "response.output")
  1987. require.True(t, output.IsArray())
  1988. require.Len(t, output.Array(), 1)
  1989. require.Equal(t, "pong", gjson.GetBytes(terminalPayload, "response.output.0.content.0.text").String())
  1990. }
  1991. func TestOpenAIStreamingNormalizesTerminalOutputToEmptyArray(t *testing.T) {
  1992. gin.SetMode(gin.TestMode)
  1993. cfg := &config.Config{
  1994. Gateway: config.GatewayConfig{
  1995. StreamDataIntervalTimeout: 0,
  1996. StreamKeepaliveInterval: 0,
  1997. MaxLineSize: defaultMaxLineSize,
  1998. },
  1999. }
  2000. svc := &OpenAIGatewayService{cfg: cfg}
  2001. rec := httptest.NewRecorder()
  2002. c, _ := gin.CreateTestContext(rec)
  2003. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2004. resp := &http.Response{
  2005. StatusCode: http.StatusOK,
  2006. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  2007. `data: {"type":"response.completed","response":{"id":"resp_empty","status":"completed","output":null,"usage":{"input_tokens":1,"output_tokens":0}}}`,
  2008. "",
  2009. }, "\n"))),
  2010. Header: http.Header{"X-Request-Id": []string{"rid-empty-output"}},
  2011. }
  2012. result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  2013. require.NoError(t, err)
  2014. require.NotNil(t, result)
  2015. terminalType, terminalPayload, ok := extractOpenAISSETerminalEvent(rec.Body.String())
  2016. require.True(t, ok)
  2017. require.Equal(t, "response.completed", terminalType)
  2018. output := gjson.GetBytes(terminalPayload, "response.output")
  2019. require.True(t, output.IsArray())
  2020. require.Len(t, output.Array(), 0)
  2021. }
  2022. func TestOpenAIStreamingPolicyResponseFailedBeforeOutputPassesThrough(t *testing.T) {
  2023. gin.SetMode(gin.TestMode)
  2024. cfg := &config.Config{
  2025. Gateway: config.GatewayConfig{
  2026. StreamDataIntervalTimeout: 0,
  2027. StreamKeepaliveInterval: 0,
  2028. MaxLineSize: defaultMaxLineSize,
  2029. },
  2030. }
  2031. svc := &OpenAIGatewayService{cfg: cfg}
  2032. rec := httptest.NewRecorder()
  2033. c, _ := gin.CreateTestContext(rec)
  2034. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2035. resp := &http.Response{
  2036. StatusCode: http.StatusOK,
  2037. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  2038. "event: response.created",
  2039. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  2040. "",
  2041. "event: response.failed",
  2042. `data: {"type":"response.failed","error":{"type":"safety_error","message":"This request has been flagged for potentially high-risk cyber activity."}}`,
  2043. "",
  2044. }, "\n"))),
  2045. Header: http.Header{"X-Request-Id": []string{"rid-policy-failed"}},
  2046. }
  2047. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
  2048. require.Error(t, err)
  2049. var failoverErr *UpstreamFailoverError
  2050. require.False(t, errors.As(err, &failoverErr))
  2051. require.True(t, c.Writer.Written())
  2052. require.Contains(t, rec.Body.String(), "response.failed")
  2053. require.Contains(t, rec.Body.String(), "high-risk cyber activity")
  2054. }
  2055. func TestOpenAIStreamingClientDisconnectDrainsUpstreamUsage(t *testing.T) {
  2056. gin.SetMode(gin.TestMode)
  2057. cfg := &config.Config{
  2058. Gateway: config.GatewayConfig{
  2059. StreamDataIntervalTimeout: 0,
  2060. StreamKeepaliveInterval: 0,
  2061. MaxLineSize: defaultMaxLineSize,
  2062. },
  2063. }
  2064. svc := &OpenAIGatewayService{cfg: cfg}
  2065. rec := httptest.NewRecorder()
  2066. c, _ := gin.CreateTestContext(rec)
  2067. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2068. c.Writer = &failingGinWriter{ResponseWriter: c.Writer, failAfter: 0}
  2069. pr, pw := io.Pipe()
  2070. resp := &http.Response{
  2071. StatusCode: http.StatusOK,
  2072. Body: pr,
  2073. Header: http.Header{},
  2074. }
  2075. go func() {
  2076. defer func() { _ = pw.Close() }()
  2077. _, _ = pw.Write([]byte("data: {\"type\":\"response.in_progress\",\"response\":{}}\n\n"))
  2078. _, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n"))
  2079. }()
  2080. result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
  2081. _ = pr.Close()
  2082. if err != nil {
  2083. t.Fatalf("expected nil error, got %v", err)
  2084. }
  2085. if result == nil || result.usage == nil {
  2086. t.Fatalf("expected usage result")
  2087. }
  2088. if result.usage.InputTokens != 3 || result.usage.OutputTokens != 5 || result.usage.CacheReadInputTokens != 1 {
  2089. t.Fatalf("unexpected usage: %+v", *result.usage)
  2090. }
  2091. if strings.Contains(rec.Body.String(), "event: error") || strings.Contains(rec.Body.String(), "write_failed") {
  2092. t.Fatalf("expected no injected SSE error event, got %q", rec.Body.String())
  2093. }
  2094. }
  2095. func TestOpenAIStreamingMissingTerminalEventReturnsIncompleteError(t *testing.T) {
  2096. gin.SetMode(gin.TestMode)
  2097. cfg := &config.Config{
  2098. Gateway: config.GatewayConfig{
  2099. StreamDataIntervalTimeout: 0,
  2100. StreamKeepaliveInterval: 0,
  2101. MaxLineSize: defaultMaxLineSize,
  2102. },
  2103. }
  2104. svc := &OpenAIGatewayService{cfg: cfg}
  2105. rec := httptest.NewRecorder()
  2106. c, _ := gin.CreateTestContext(rec)
  2107. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2108. pr, pw := io.Pipe()
  2109. resp := &http.Response{
  2110. StatusCode: http.StatusOK,
  2111. Body: pr,
  2112. Header: http.Header{},
  2113. }
  2114. go func() {
  2115. defer func() { _ = pw.Close() }()
  2116. _, _ = pw.Write([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\",\"output_index\":0}\n\n"))
  2117. }()
  2118. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
  2119. _ = pr.Close()
  2120. if err == nil || !strings.Contains(err.Error(), "missing terminal event") {
  2121. t.Fatalf("expected missing terminal event error, got %v", err)
  2122. }
  2123. }
  2124. func TestOpenAIStreamingPassthroughMissingTerminalEventReturnsIncompleteError(t *testing.T) {
  2125. gin.SetMode(gin.TestMode)
  2126. cfg := &config.Config{
  2127. Gateway: config.GatewayConfig{
  2128. MaxLineSize: defaultMaxLineSize,
  2129. },
  2130. }
  2131. svc := &OpenAIGatewayService{cfg: cfg}
  2132. rec := httptest.NewRecorder()
  2133. c, _ := gin.CreateTestContext(rec)
  2134. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2135. pr, pw := io.Pipe()
  2136. resp := &http.Response{
  2137. StatusCode: http.StatusOK,
  2138. Body: pr,
  2139. Header: http.Header{},
  2140. }
  2141. go func() {
  2142. defer func() { _ = pw.Close() }()
  2143. _, _ = pw.Write([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\",\"output_index\":0}\n\n"))
  2144. }()
  2145. _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "", "")
  2146. _ = pr.Close()
  2147. if err == nil || !strings.Contains(err.Error(), "missing terminal event") {
  2148. t.Fatalf("expected missing terminal event error, got %v", err)
  2149. }
  2150. }
  2151. func TestOpenAIStreamingPassthroughPostOutputDisconnectQuarantinesSharedProxy(t *testing.T) {
  2152. gin.SetMode(gin.TestMode)
  2153. proxyID := int64(4698)
  2154. account := &Account{ID: 469804, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, ProxyID: &proxyID}
  2155. svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}}
  2156. // collapseInterval 0: the loop below records within the production collapse
  2157. // window and must count as distinct failure events here.
  2158. svc.openaiProxyStreamCircuit = newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
  2159. failureThreshold: 2,
  2160. failureWindow: time.Minute,
  2161. quarantineTTL: 10 * time.Minute,
  2162. maxEntries: 16,
  2163. })
  2164. for _, readErr := range []error{io.ErrUnexpectedEOF, errors.New("http2: client connection lost")} {
  2165. rec := httptest.NewRecorder()
  2166. c, _ := gin.CreateTestContext(rec)
  2167. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  2168. resp := &http.Response{
  2169. StatusCode: http.StatusOK,
  2170. Body: &openAIStreamReadThenErrorCloser{
  2171. reader: strings.NewReader("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"),
  2172. err: readErr,
  2173. },
  2174. Header: http.Header{"X-Request-Id": []string{"rid-passthrough-proxy-disconnect"}},
  2175. }
  2176. _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
  2177. require.Error(t, err)
  2178. var failoverErr *UpstreamFailoverError
  2179. require.False(t, errors.As(err, &failoverErr), "post-output disconnect must not fail over inside the same stream")
  2180. require.Contains(t, rec.Body.String(), "partial")
  2181. }
  2182. require.True(t, svc.isOpenAIProxyStreamQuarantined(context.Background(), account))
  2183. }
  2184. func TestOpenAIStreamingPassthroughResponseFailedBeforeOutputReturnsFailover(t *testing.T) {
  2185. gin.SetMode(gin.TestMode)
  2186. cfg := &config.Config{
  2187. Gateway: config.GatewayConfig{
  2188. MaxLineSize: defaultMaxLineSize,
  2189. },
  2190. }
  2191. svc := &OpenAIGatewayService{cfg: cfg}
  2192. rec := httptest.NewRecorder()
  2193. c, _ := gin.CreateTestContext(rec)
  2194. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2195. resp := &http.Response{
  2196. StatusCode: http.StatusOK,
  2197. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  2198. "event: response.created",
  2199. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  2200. "",
  2201. "event: response.failed",
  2202. `data: {"type":"response.failed","error":{"message":"upstream processing failed"}}`,
  2203. "",
  2204. }, "\n"))),
  2205. Header: http.Header{"X-Request-Id": []string{"rid-passthrough-failed"}},
  2206. }
  2207. _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
  2208. require.Error(t, err)
  2209. var failoverErr *UpstreamFailoverError
  2210. require.ErrorAs(t, err, &failoverErr)
  2211. require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
  2212. require.Contains(t, string(failoverErr.ResponseBody), "upstream processing failed")
  2213. require.False(t, c.Writer.Written())
  2214. require.Empty(t, rec.Body.String())
  2215. }
  2216. func TestOpenAIStreamingPassthroughContextWindowResponseFailedBeforeOutputAppliesPassthroughRule(t *testing.T) {
  2217. gin.SetMode(gin.TestMode)
  2218. cfg := &config.Config{
  2219. Gateway: config.GatewayConfig{
  2220. MaxLineSize: defaultMaxLineSize,
  2221. },
  2222. }
  2223. svc := &OpenAIGatewayService{cfg: cfg}
  2224. rec := httptest.NewRecorder()
  2225. c, _ := gin.CreateTestContext(rec)
  2226. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2227. rule := newNonFailoverPassthroughRule(http.StatusBadRequest, "input exceeds the context window", http.StatusBadRequest, "")
  2228. rule.Platforms = []string{PlatformOpenAI}
  2229. rule.PassthroughBody = true
  2230. rule.CustomMessage = nil
  2231. ruleSvc := &ErrorPassthroughService{}
  2232. ruleSvc.setLocalCache([]*model.ErrorPassthroughRule{rule})
  2233. BindErrorPassthroughService(c, ruleSvc)
  2234. upstreamMessage := "Your input exceeds the context window of this model. Please adjust your input and try again."
  2235. resp := &http.Response{
  2236. StatusCode: http.StatusOK,
  2237. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  2238. "event: response.created",
  2239. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  2240. "",
  2241. "event: response.failed",
  2242. `data: {"type":"response.failed","response":{"id":"resp_1","error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"` + upstreamMessage + `"}}}`,
  2243. "",
  2244. }, "\n"))),
  2245. Header: http.Header{"X-Request-Id": []string{"rid-pass-context-window-passthrough-rule"}},
  2246. }
  2247. _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
  2248. require.Error(t, err)
  2249. var failoverErr *UpstreamFailoverError
  2250. require.False(t, errors.As(err, &failoverErr))
  2251. require.True(t, IsResponseCommitted(c))
  2252. require.Equal(t, http.StatusBadRequest, rec.Code)
  2253. body := rec.Body.String()
  2254. require.Equal(t, "upstream_error", gjson.Get(body, "error.type").String())
  2255. require.Equal(t, upstreamMessage, gjson.Get(body, "error.message").String())
  2256. require.NotContains(t, body, "response.failed")
  2257. require.NotContains(t, body, "Upstream request failed")
  2258. // 命中透传规则也应记录 ops 上游错误事件(对齐 CC/Messages 与 antigravity 先例)。
  2259. opsVal, opsRecorded := c.Get(OpsUpstreamErrorsKey)
  2260. require.True(t, opsRecorded, "passthrough hit should record an ops upstream error event")
  2261. opsEvents, _ := opsVal.([]*OpsUpstreamErrorEvent)
  2262. require.NotEmpty(t, opsEvents)
  2263. }
  2264. func TestOpenAIStreamingPassthroughContextWindowResponseFailedBeforeOutputWithoutRulePassesThrough(t *testing.T) {
  2265. gin.SetMode(gin.TestMode)
  2266. cfg := &config.Config{
  2267. Gateway: config.GatewayConfig{
  2268. MaxLineSize: defaultMaxLineSize,
  2269. },
  2270. }
  2271. svc := &OpenAIGatewayService{cfg: cfg}
  2272. rec := httptest.NewRecorder()
  2273. c, _ := gin.CreateTestContext(rec)
  2274. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2275. resp := &http.Response{
  2276. StatusCode: http.StatusOK,
  2277. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  2278. "event: response.created",
  2279. `data: {"type":"response.created","response":{"id":"resp_1"}}`,
  2280. "",
  2281. "event: response.failed",
  2282. `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."}}}`,
  2283. "",
  2284. }, "\n"))),
  2285. Header: http.Header{"X-Request-Id": []string{"rid-pass-context-window-no-rule"}},
  2286. }
  2287. _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
  2288. require.Error(t, err)
  2289. var failoverErr *UpstreamFailoverError
  2290. require.False(t, errors.As(err, &failoverErr))
  2291. body := rec.Body.String()
  2292. require.Contains(t, body, "event: response.failed")
  2293. require.Contains(t, body, "context_length_exceeded")
  2294. require.Contains(t, body, "Your input exceeds the context window")
  2295. }
  2296. func TestOpenAIStreamingPassthroughResponseFailedAfterOutputSanitizesVerboseResponseForClient(t *testing.T) {
  2297. gin.SetMode(gin.TestMode)
  2298. cfg := &config.Config{
  2299. Gateway: config.GatewayConfig{
  2300. MaxLineSize: defaultMaxLineSize,
  2301. },
  2302. }
  2303. svc := &OpenAIGatewayService{cfg: cfg}
  2304. rec := httptest.NewRecorder()
  2305. c, _ := gin.CreateTestContext(rec)
  2306. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2307. longInstructions := strings.Repeat("You are GPT-5.1 running in the Codex CLI. ", 20)
  2308. failedPayload := fmt.Sprintf(
  2309. `{"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."}}}`,
  2310. longInstructions,
  2311. )
  2312. resp := &http.Response{
  2313. StatusCode: http.StatusOK,
  2314. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  2315. "event: response.created",
  2316. `data: {"type":"response.created","response":{"id":"resp_pass_failed"}}`,
  2317. "",
  2318. "event: response.output_text.delta",
  2319. `data: {"type":"response.output_text.delta","delta":"partial"}`,
  2320. "",
  2321. "event: response.failed",
  2322. "data: " + failedPayload,
  2323. "",
  2324. }, "\n"))),
  2325. Header: http.Header{"X-Request-Id": []string{"rid-pass-failed-after-output"}},
  2326. }
  2327. _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
  2328. require.Error(t, err)
  2329. body := rec.Body.String()
  2330. require.Contains(t, body, "event: response.failed")
  2331. require.Contains(t, body, "context_length_exceeded")
  2332. require.Contains(t, body, `"type":"invalid_request_error"`)
  2333. require.Contains(t, body, "Your input exceeds the context window")
  2334. require.NotContains(t, body, "You are GPT-5.1 running in the Codex CLI")
  2335. require.NotContains(t, body, `"instructions"`)
  2336. require.NotContains(t, body, `"output"`)
  2337. require.NotContains(t, body, `"usage"`)
  2338. }
  2339. func TestOpenAIStreamingPassthroughResponseDoneWithoutDoneMarkerStillSucceeds(t *testing.T) {
  2340. gin.SetMode(gin.TestMode)
  2341. cfg := &config.Config{
  2342. Gateway: config.GatewayConfig{
  2343. MaxLineSize: defaultMaxLineSize,
  2344. },
  2345. }
  2346. svc := &OpenAIGatewayService{cfg: cfg}
  2347. rec := httptest.NewRecorder()
  2348. c, _ := gin.CreateTestContext(rec)
  2349. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2350. pr, pw := io.Pipe()
  2351. resp := &http.Response{
  2352. StatusCode: http.StatusOK,
  2353. Body: pr,
  2354. Header: http.Header{},
  2355. }
  2356. go func() {
  2357. defer func() { _ = pw.Close() }()
  2358. _, _ = pw.Write([]byte("data: {\"type\":\"response.done\",\"response\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n"))
  2359. }()
  2360. result, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "", "")
  2361. _ = pr.Close()
  2362. require.NoError(t, err)
  2363. require.NotNil(t, result)
  2364. require.NotNil(t, result.usage)
  2365. require.Equal(t, 2, result.usage.InputTokens)
  2366. require.Equal(t, 3, result.usage.OutputTokens)
  2367. require.Equal(t, 1, result.usage.CacheReadInputTokens)
  2368. }
  2369. func TestOpenAIStreamingPassthroughResponseIncompleteWithoutDoneMarkerStillSucceeds(t *testing.T) {
  2370. gin.SetMode(gin.TestMode)
  2371. cfg := &config.Config{
  2372. Gateway: config.GatewayConfig{
  2373. MaxLineSize: defaultMaxLineSize,
  2374. },
  2375. }
  2376. svc := &OpenAIGatewayService{cfg: cfg}
  2377. rec := httptest.NewRecorder()
  2378. c, _ := gin.CreateTestContext(rec)
  2379. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2380. pr, pw := io.Pipe()
  2381. resp := &http.Response{
  2382. StatusCode: http.StatusOK,
  2383. Body: pr,
  2384. Header: http.Header{},
  2385. }
  2386. go func() {
  2387. defer func() { _ = pw.Close() }()
  2388. _, _ = pw.Write([]byte("data: {\"type\":\"response.incomplete\",\"response\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n"))
  2389. }()
  2390. result, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "", "")
  2391. _ = pr.Close()
  2392. require.NoError(t, err)
  2393. require.NotNil(t, result)
  2394. require.NotNil(t, result.usage)
  2395. require.Equal(t, 2, result.usage.InputTokens)
  2396. require.Equal(t, 3, result.usage.OutputTokens)
  2397. require.Equal(t, 1, result.usage.CacheReadInputTokens)
  2398. }
  2399. func TestOpenAIStreamingTooLong(t *testing.T) {
  2400. gin.SetMode(gin.TestMode)
  2401. cfg := &config.Config{
  2402. Gateway: config.GatewayConfig{
  2403. StreamDataIntervalTimeout: 0,
  2404. StreamKeepaliveInterval: 0,
  2405. MaxLineSize: 64 * 1024,
  2406. },
  2407. }
  2408. svc := &OpenAIGatewayService{cfg: cfg}
  2409. rec := httptest.NewRecorder()
  2410. c, _ := gin.CreateTestContext(rec)
  2411. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2412. pr, pw := io.Pipe()
  2413. resp := &http.Response{
  2414. StatusCode: http.StatusOK,
  2415. Body: pr,
  2416. Header: http.Header{},
  2417. }
  2418. go func() {
  2419. defer func() { _ = pw.Close() }()
  2420. // 写入超过 MaxLineSize 的单行数据,触发 ErrTooLong
  2421. payload := "data: " + strings.Repeat("a", 128*1024) + "\n"
  2422. _, _ = pw.Write([]byte(payload))
  2423. }()
  2424. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 2}, time.Now(), "model", "model")
  2425. _ = pr.Close()
  2426. if !errors.Is(err, bufio.ErrTooLong) {
  2427. t.Fatalf("expected ErrTooLong, got %v", err)
  2428. }
  2429. if !strings.Contains(rec.Body.String(), "\"type\":\"error\"") || !strings.Contains(rec.Body.String(), "response_too_large") {
  2430. t.Fatalf("expected OpenAI-compatible error SSE event, got %q", rec.Body.String())
  2431. }
  2432. }
  2433. func TestOpenAINonStreamingContentTypePassThrough(t *testing.T) {
  2434. gin.SetMode(gin.TestMode)
  2435. cfg := &config.Config{
  2436. Security: config.SecurityConfig{
  2437. ResponseHeaders: config.ResponseHeaderConfig{Enabled: false},
  2438. },
  2439. }
  2440. svc := &OpenAIGatewayService{cfg: cfg}
  2441. rec := httptest.NewRecorder()
  2442. c, _ := gin.CreateTestContext(rec)
  2443. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2444. body := []byte(`{"usage":{"input_tokens":1,"output_tokens":2,"input_tokens_details":{"cached_tokens":0}}}`)
  2445. resp := &http.Response{
  2446. StatusCode: http.StatusOK,
  2447. Body: io.NopCloser(bytes.NewReader(body)),
  2448. Header: http.Header{"Content-Type": []string{"application/vnd.test+json"}},
  2449. }
  2450. _, err := svc.handleNonStreamingResponse(c.Request.Context(), resp, c, &Account{}, "model", "model")
  2451. if err != nil {
  2452. t.Fatalf("handleNonStreamingResponse error: %v", err)
  2453. }
  2454. if !strings.Contains(rec.Header().Get("Content-Type"), "application/vnd.test+json") {
  2455. t.Fatalf("expected Content-Type passthrough, got %q", rec.Header().Get("Content-Type"))
  2456. }
  2457. }
  2458. func TestOpenAINonStreamingContentTypeDefault(t *testing.T) {
  2459. gin.SetMode(gin.TestMode)
  2460. cfg := &config.Config{
  2461. Security: config.SecurityConfig{
  2462. ResponseHeaders: config.ResponseHeaderConfig{Enabled: false},
  2463. },
  2464. }
  2465. svc := &OpenAIGatewayService{cfg: cfg}
  2466. rec := httptest.NewRecorder()
  2467. c, _ := gin.CreateTestContext(rec)
  2468. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2469. body := []byte(`{"usage":{"input_tokens":1,"output_tokens":2,"input_tokens_details":{"cached_tokens":0}}}`)
  2470. resp := &http.Response{
  2471. StatusCode: http.StatusOK,
  2472. Body: io.NopCloser(bytes.NewReader(body)),
  2473. Header: http.Header{},
  2474. }
  2475. _, err := svc.handleNonStreamingResponse(c.Request.Context(), resp, c, &Account{}, "model", "model")
  2476. if err != nil {
  2477. t.Fatalf("handleNonStreamingResponse error: %v", err)
  2478. }
  2479. if !strings.Contains(rec.Header().Get("Content-Type"), "application/json") {
  2480. t.Fatalf("expected default Content-Type, got %q", rec.Header().Get("Content-Type"))
  2481. }
  2482. }
  2483. func TestOpenAIStreamingHeadersOverride(t *testing.T) {
  2484. gin.SetMode(gin.TestMode)
  2485. cfg := &config.Config{
  2486. Security: config.SecurityConfig{
  2487. ResponseHeaders: config.ResponseHeaderConfig{Enabled: false},
  2488. },
  2489. Gateway: config.GatewayConfig{
  2490. StreamDataIntervalTimeout: 0,
  2491. StreamKeepaliveInterval: 0,
  2492. MaxLineSize: defaultMaxLineSize,
  2493. },
  2494. }
  2495. svc := &OpenAIGatewayService{cfg: cfg}
  2496. rec := httptest.NewRecorder()
  2497. c, _ := gin.CreateTestContext(rec)
  2498. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2499. pr, pw := io.Pipe()
  2500. resp := &http.Response{
  2501. StatusCode: http.StatusOK,
  2502. Body: pr,
  2503. Header: http.Header{
  2504. "Cache-Control": []string{"upstream"},
  2505. "X-Request-Id": []string{"req-123"},
  2506. "Content-Type": []string{"application/custom"},
  2507. },
  2508. }
  2509. go func() {
  2510. defer func() { _ = pw.Close() }()
  2511. _, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{}}\n\n"))
  2512. }()
  2513. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
  2514. _ = pr.Close()
  2515. if err != nil {
  2516. t.Fatalf("handleStreamingResponse error: %v", err)
  2517. }
  2518. if rec.Header().Get("Cache-Control") != "no-cache" {
  2519. t.Fatalf("expected Cache-Control override, got %q", rec.Header().Get("Cache-Control"))
  2520. }
  2521. if rec.Header().Get("Content-Type") != "text/event-stream" {
  2522. t.Fatalf("expected Content-Type override, got %q", rec.Header().Get("Content-Type"))
  2523. }
  2524. if rec.Header().Get("X-Request-Id") != "req-123" {
  2525. t.Fatalf("expected X-Request-Id passthrough, got %q", rec.Header().Get("X-Request-Id"))
  2526. }
  2527. }
  2528. func TestOpenAIStreamingReuseScannerBufferAndStillWorks(t *testing.T) {
  2529. gin.SetMode(gin.TestMode)
  2530. cfg := &config.Config{
  2531. Gateway: config.GatewayConfig{
  2532. StreamDataIntervalTimeout: 0,
  2533. StreamKeepaliveInterval: 0,
  2534. MaxLineSize: defaultMaxLineSize,
  2535. },
  2536. }
  2537. svc := &OpenAIGatewayService{cfg: cfg}
  2538. rec := httptest.NewRecorder()
  2539. c, _ := gin.CreateTestContext(rec)
  2540. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2541. pr, pw := io.Pipe()
  2542. resp := &http.Response{
  2543. StatusCode: http.StatusOK,
  2544. Body: pr,
  2545. Header: http.Header{},
  2546. }
  2547. go func() {
  2548. defer func() { _ = pw.Close() }()
  2549. _, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"input_tokens_details\":{\"cached_tokens\":3}}}}\n\n"))
  2550. }()
  2551. result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
  2552. _ = pr.Close()
  2553. require.NoError(t, err)
  2554. require.NotNil(t, result)
  2555. require.NotNil(t, result.usage)
  2556. require.Equal(t, 1, result.usage.InputTokens)
  2557. require.Equal(t, 2, result.usage.OutputTokens)
  2558. require.Equal(t, 3, result.usage.CacheReadInputTokens)
  2559. }
  2560. func TestOpenAIInvalidBaseURLWhenAllowlistDisabled(t *testing.T) {
  2561. gin.SetMode(gin.TestMode)
  2562. cfg := &config.Config{
  2563. Security: config.SecurityConfig{
  2564. URLAllowlist: config.URLAllowlistConfig{Enabled: false},
  2565. },
  2566. }
  2567. svc := &OpenAIGatewayService{cfg: cfg}
  2568. rec := httptest.NewRecorder()
  2569. c, _ := gin.CreateTestContext(rec)
  2570. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  2571. account := &Account{
  2572. Platform: PlatformOpenAI,
  2573. Type: AccountTypeAPIKey,
  2574. Credentials: map[string]any{"base_url": "://invalid-url"},
  2575. }
  2576. _, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte("{}"), "token", false, "", false)
  2577. if err == nil {
  2578. t.Fatalf("expected error for invalid base_url when allowlist disabled")
  2579. }
  2580. }
  2581. func TestOpenAIValidateUpstreamBaseURLDisabledRequiresHTTPS(t *testing.T) {
  2582. cfg := &config.Config{
  2583. Security: config.SecurityConfig{
  2584. URLAllowlist: config.URLAllowlistConfig{Enabled: false},
  2585. },
  2586. }
  2587. svc := &OpenAIGatewayService{cfg: cfg}
  2588. if _, err := svc.validateUpstreamBaseURL("http://not-https.example.com"); err == nil {
  2589. t.Fatalf("expected http to be rejected when allow_insecure_http is false")
  2590. }
  2591. normalized, err := svc.validateUpstreamBaseURL("https://example.com")
  2592. if err != nil {
  2593. t.Fatalf("expected https to be allowed when allowlist disabled, got %v", err)
  2594. }
  2595. if normalized != "https://example.com" {
  2596. t.Fatalf("expected raw url passthrough, got %q", normalized)
  2597. }
  2598. }
  2599. func TestOpenAIValidateUpstreamBaseURLDisabledAllowsHTTP(t *testing.T) {
  2600. cfg := &config.Config{
  2601. Security: config.SecurityConfig{
  2602. URLAllowlist: config.URLAllowlistConfig{
  2603. Enabled: false,
  2604. AllowInsecureHTTP: true,
  2605. },
  2606. },
  2607. }
  2608. svc := &OpenAIGatewayService{cfg: cfg}
  2609. normalized, err := svc.validateUpstreamBaseURL("http://not-https.example.com")
  2610. if err != nil {
  2611. t.Fatalf("expected http allowed when allow_insecure_http is true, got %v", err)
  2612. }
  2613. if normalized != "http://not-https.example.com" {
  2614. t.Fatalf("expected raw url passthrough, got %q", normalized)
  2615. }
  2616. }
  2617. func TestOpenAIValidateUpstreamBaseURLEnabledEnforcesAllowlist(t *testing.T) {
  2618. cfg := &config.Config{
  2619. Security: config.SecurityConfig{
  2620. URLAllowlist: config.URLAllowlistConfig{
  2621. Enabled: true,
  2622. UpstreamHosts: []string{"example.com"},
  2623. },
  2624. },
  2625. }
  2626. svc := &OpenAIGatewayService{cfg: cfg}
  2627. if _, err := svc.validateUpstreamBaseURL("https://example.com"); err != nil {
  2628. t.Fatalf("expected allowlisted host to pass, got %v", err)
  2629. }
  2630. if _, err := svc.validateUpstreamBaseURL("https://evil.com"); err == nil {
  2631. t.Fatalf("expected non-allowlisted host to fail")
  2632. }
  2633. }
  2634. func TestOpenAIUpdateCodexUsageSnapshotFromHeaders(t *testing.T) {
  2635. repo := &snapshotUpdateAccountRepo{updateExtraCalls: make(chan map[string]any, 1)}
  2636. svc := &OpenAIGatewayService{accountRepo: repo}
  2637. headers := http.Header{}
  2638. headers.Set("x-codex-primary-used-percent", "12")
  2639. headers.Set("x-codex-secondary-used-percent", "34")
  2640. headers.Set("x-codex-primary-window-minutes", "300")
  2641. headers.Set("x-codex-secondary-window-minutes", "10080")
  2642. headers.Set("x-codex-primary-reset-after-seconds", "600")
  2643. headers.Set("x-codex-secondary-reset-after-seconds", "86400")
  2644. svc.UpdateCodexUsageSnapshotFromHeaders(context.Background(), 123, headers)
  2645. select {
  2646. case updates := <-repo.updateExtraCalls:
  2647. require.Equal(t, 12.0, updates["codex_5h_used_percent"])
  2648. require.Equal(t, 34.0, updates["codex_7d_used_percent"])
  2649. require.Equal(t, 600, updates["codex_5h_reset_after_seconds"])
  2650. require.Equal(t, 86400, updates["codex_7d_reset_after_seconds"])
  2651. case <-time.After(2 * time.Second):
  2652. t.Fatal("expected UpdateExtra to be called")
  2653. }
  2654. }
  2655. func TestOpenAIResponsesRequestPathSuffix(t *testing.T) {
  2656. gin.SetMode(gin.TestMode)
  2657. rec := httptest.NewRecorder()
  2658. c, _ := gin.CreateTestContext(rec)
  2659. tests := []struct {
  2660. name string
  2661. path string
  2662. want string
  2663. }{
  2664. {name: "exact v1 responses", path: "/v1/responses", want: ""},
  2665. {name: "compact v1 responses", path: "/v1/responses/compact", want: "/compact"},
  2666. {name: "compact alias responses", path: "/responses/compact/", want: "/compact"},
  2667. {name: "nested suffix", path: "/openai/v1/responses/compact/detail", want: "/compact/detail"},
  2668. {name: "unrelated path", path: "/v1/chat/completions", want: ""},
  2669. }
  2670. for _, tt := range tests {
  2671. t.Run(tt.name, func(t *testing.T) {
  2672. c.Request = httptest.NewRequest(http.MethodPost, tt.path, nil)
  2673. require.Equal(t, tt.want, openAIResponsesRequestPathSuffix(c))
  2674. })
  2675. }
  2676. }
  2677. func TestNormalizeOpenAICompactRequestBodyPreservesCurrentCodexPayloadFields(t *testing.T) {
  2678. 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"}`)
  2679. normalized, changed, err := normalizeOpenAICompactRequestBody(body)
  2680. require.NoError(t, err)
  2681. require.True(t, changed)
  2682. require.Equal(t, "gpt-5.5", gjson.GetBytes(normalized, "model").String())
  2683. require.True(t, gjson.GetBytes(normalized, "tools").Exists())
  2684. require.True(t, gjson.GetBytes(normalized, "parallel_tool_calls").Bool())
  2685. require.Equal(t, "high", gjson.GetBytes(normalized, "reasoning.effort").String())
  2686. require.Equal(t, "low", gjson.GetBytes(normalized, "text.verbosity").String())
  2687. require.Equal(t, "resp_123", gjson.GetBytes(normalized, "previous_response_id").String())
  2688. require.False(t, gjson.GetBytes(normalized, "store").Exists())
  2689. require.False(t, gjson.GetBytes(normalized, "stream").Exists())
  2690. require.False(t, gjson.GetBytes(normalized, "prompt_cache_key").Exists())
  2691. }
  2692. func TestNormalizeOpenAICompactRequestBodyDropsParallelToolCallsWithoutUsableTools(t *testing.T) {
  2693. for _, body := range []string{
  2694. `{"model":"gpt-5.5","input":[],"parallel_tool_calls":true}`,
  2695. `{"model":"gpt-5.5","input":[],"tools":[],"parallel_tool_calls":false}`,
  2696. `{"model":"gpt-5.5","input":[],"tools":null,"parallel_tool_calls":true}`,
  2697. `{"model":"gpt-5.5","input":[],"tools":{},"parallel_tool_calls":true}`,
  2698. } {
  2699. normalized, changed, err := normalizeOpenAICompactRequestBody([]byte(body))
  2700. require.NoError(t, err)
  2701. require.True(t, changed)
  2702. require.False(t, gjson.GetBytes(normalized, "parallel_tool_calls").Exists(), string(normalized))
  2703. }
  2704. withTools := []byte(`{"model":"gpt-5.5","tools":[{"type":"function","name":"lookup"}],"parallel_tool_calls":false}`)
  2705. normalized, _, err := normalizeOpenAICompactRequestBody(withTools)
  2706. require.NoError(t, err)
  2707. require.True(t, gjson.GetBytes(normalized, "parallel_tool_calls").Exists())
  2708. require.False(t, gjson.GetBytes(normalized, "parallel_tool_calls").Bool())
  2709. }
  2710. func TestOpenAIBuildUpstreamRequestOpenAIPassthroughPreservesCompactPath(t *testing.T) {
  2711. gin.SetMode(gin.TestMode)
  2712. rec := httptest.NewRecorder()
  2713. c, _ := gin.CreateTestContext(rec)
  2714. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
  2715. svc := &OpenAIGatewayService{}
  2716. account := &Account{Type: AccountTypeOAuth}
  2717. req, err := svc.buildUpstreamRequestOpenAIPassthrough(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token")
  2718. require.NoError(t, err)
  2719. require.Equal(t, chatgptCodexURL+"/compact", req.URL.String())
  2720. require.Equal(t, "application/json", req.Header.Get("Accept"))
  2721. require.Equal(t, codexCLIVersion, req.Header.Get("Version"))
  2722. require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header")
  2723. require.NotEmpty(t, req.Header.Get("Session_Id"))
  2724. require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context()))
  2725. }
  2726. func TestOpenAIBuildUpstreamRequestOpenAIPassthroughPreservesExplicitAPIKeyBetaHeader(t *testing.T) {
  2727. gin.SetMode(gin.TestMode)
  2728. rec := httptest.NewRecorder()
  2729. c, _ := gin.CreateTestContext(rec)
  2730. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
  2731. c.Request.Header.Set("OpenAI-Beta", "api-key-specific-beta")
  2732. svc := &OpenAIGatewayService{cfg: &config.Config{
  2733. Security: config.SecurityConfig{
  2734. URLAllowlist: config.URLAllowlistConfig{Enabled: false},
  2735. },
  2736. }}
  2737. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  2738. req, err := svc.buildUpstreamRequestOpenAIPassthrough(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token")
  2739. require.NoError(t, err)
  2740. require.Equal(t, "api-key-specific-beta", req.Header.Get("OpenAI-Beta"), "OAuth-only backport must not alter API-key passthrough headers")
  2741. }
  2742. func TestOpenAIBuildUpstreamRequestCompactForcesJSONAcceptForOAuth(t *testing.T) {
  2743. gin.SetMode(gin.TestMode)
  2744. rec := httptest.NewRecorder()
  2745. c, _ := gin.CreateTestContext(rec)
  2746. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
  2747. svc := &OpenAIGatewayService{}
  2748. account := &Account{
  2749. Type: AccountTypeOAuth,
  2750. Credentials: map[string]any{"chatgpt_account_id": "chatgpt-acc"},
  2751. }
  2752. req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", true)
  2753. require.NoError(t, err)
  2754. require.Equal(t, chatgptCodexURL+"/compact", req.URL.String())
  2755. require.Equal(t, "application/json", req.Header.Get("Accept"))
  2756. require.Equal(t, codexCLIVersion, req.Header.Get("Version"))
  2757. require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header")
  2758. require.NotEmpty(t, req.Header.Get("Session_Id"))
  2759. require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context()))
  2760. }
  2761. func TestOpenAIBuildUpstreamRequestOAuthMessagesBridgeUsesSessionOnly(t *testing.T) {
  2762. gin.SetMode(gin.TestMode)
  2763. rec := httptest.NewRecorder()
  2764. c, _ := gin.CreateTestContext(rec)
  2765. 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"}]}`)
  2766. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  2767. c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
  2768. c.Request.Header.Set("originator", "codex_cli_rs")
  2769. svc := &OpenAIGatewayService{}
  2770. account := &Account{
  2771. Type: AccountTypeOAuth,
  2772. Credentials: map[string]any{"chatgpt_account_id": "chatgpt-acc"},
  2773. }
  2774. req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, body, "token", true, "anthropic-metadata-session-1", false)
  2775. require.NoError(t, err)
  2776. require.NotEmpty(t, req.Header.Get("Session_Id"))
  2777. require.Empty(t, req.Header.Get("Conversation_Id"))
  2778. require.Empty(t, req.Header.Get("OpenAI-Beta"))
  2779. require.Empty(t, req.Header.Get("originator"))
  2780. }
  2781. func TestOpenAIBuildUpstreamRequestPreservesCompactPathForAPIKeyBaseURL(t *testing.T) {
  2782. gin.SetMode(gin.TestMode)
  2783. rec := httptest.NewRecorder()
  2784. c, _ := gin.CreateTestContext(rec)
  2785. c.Request = httptest.NewRequest(http.MethodPost, "/responses/compact", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
  2786. svc := &OpenAIGatewayService{cfg: &config.Config{
  2787. Security: config.SecurityConfig{
  2788. URLAllowlist: config.URLAllowlistConfig{Enabled: false},
  2789. },
  2790. }}
  2791. account := &Account{
  2792. Type: AccountTypeAPIKey,
  2793. Platform: PlatformOpenAI,
  2794. Credentials: map[string]any{"base_url": "https://example.com/v1"},
  2795. }
  2796. req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", false)
  2797. require.NoError(t, err)
  2798. require.Equal(t, "https://example.com/v1/responses/compact", req.URL.String())
  2799. }
  2800. func TestOpenAIBuildUpstreamRequestPreservesCodexIdentityHeaders(t *testing.T) {
  2801. gin.SetMode(gin.TestMode)
  2802. rec := httptest.NewRecorder()
  2803. c, _ := gin.CreateTestContext(rec)
  2804. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  2805. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
  2806. c.Request.Header.Set("X-Codex-Window-ID", "window-http")
  2807. c.Request.Header.Set("X-Codex-Installation-ID", "installation-http")
  2808. c.Request.Header.Set("X-Test", "blocked")
  2809. body := []byte(`{"model":"gpt-5","input":"hello"}`)
  2810. svc := &OpenAIGatewayService{cfg: &config.Config{
  2811. Security: config.SecurityConfig{
  2812. URLAllowlist: config.URLAllowlistConfig{Enabled: false},
  2813. },
  2814. }}
  2815. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  2816. req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, body, "token", false, "", true)
  2817. require.NoError(t, err)
  2818. require.Equal(t, "window-http", req.Header.Get("X-Codex-Window-ID"))
  2819. require.Equal(t, "installation-http", req.Header.Get("X-Codex-Installation-ID"))
  2820. require.Empty(t, req.Header.Get("X-Test"))
  2821. require.True(t, openai.EvaluateEngineFingerprint(req.Header, body, openai.DefaultEngineFingerprintSignals))
  2822. }
  2823. func TestOpenAIBuildUpstreamRequestOAuthOfficialClientOriginatorCompatibility(t *testing.T) {
  2824. gin.SetMode(gin.TestMode)
  2825. // 强制统一出口:客户端自报的 originator / User-Agent 都不参与上游身份构造,
  2826. // 一律改写为网关规范身份,天然满足 originator 与 UA 首段配套的上游校验(issue #3901)。
  2827. tests := []struct {
  2828. name string
  2829. userAgent string
  2830. originator string
  2831. }{
  2832. {name: "official desktop ua", userAgent: "Codex Desktop/1.2.3"},
  2833. {
  2834. name: "mismatched originator",
  2835. userAgent: "codex_vscode/0.140.2 (Mac OS X 14.0; arm64) vscode (codex_vscode; 0.140.2)",
  2836. originator: "codex_cli_rs",
  2837. },
  2838. {
  2839. name: "tui identity",
  2840. userAgent: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
  2841. originator: "codex-tui",
  2842. },
  2843. {name: "official originator without ua", originator: "codex_vscode"},
  2844. {name: "third-party ua", userAgent: "luna/1.2.0"},
  2845. }
  2846. for _, tt := range tests {
  2847. t.Run(tt.name, func(t *testing.T) {
  2848. rec := httptest.NewRecorder()
  2849. c, _ := gin.CreateTestContext(rec)
  2850. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
  2851. if tt.userAgent != "" {
  2852. c.Request.Header.Set("User-Agent", tt.userAgent)
  2853. }
  2854. if tt.originator != "" {
  2855. c.Request.Header.Set("originator", tt.originator)
  2856. }
  2857. svc := &OpenAIGatewayService{}
  2858. account := &Account{
  2859. Type: AccountTypeOAuth,
  2860. Credentials: map[string]any{"chatgpt_account_id": "chatgpt-acc"},
  2861. }
  2862. isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator"))
  2863. req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", isCodexCLI)
  2864. require.NoError(t, err)
  2865. require.Equal(t, openai.CodexDefaultOriginator, req.Header.Get("originator"))
  2866. require.Equal(t, codexCLIUserAgent, req.Header.Get("User-Agent"))
  2867. require.Equal(t, codexCLIVersion, req.Header.Get("version"))
  2868. })
  2869. }
  2870. }
  2871. // ==================== P1-08 修复:model 替换性能优化测试 ====================
  2872. // ==================== P1-08 修复:model 替换性能优化测试 =============
  2873. func TestReplaceModelInSSELine(t *testing.T) {
  2874. svc := &OpenAIGatewayService{}
  2875. tests := []struct {
  2876. name string
  2877. line string
  2878. from string
  2879. to string
  2880. expected string
  2881. }{
  2882. {
  2883. name: "顶层 model 字段替换",
  2884. line: `data: {"id":"chatcmpl-123","model":"gpt-4o","choices":[]}`,
  2885. from: "gpt-4o",
  2886. to: "my-custom-model",
  2887. expected: `data: {"id":"chatcmpl-123","model":"my-custom-model","choices":[]}`,
  2888. },
  2889. {
  2890. name: "嵌套 response.model 替换",
  2891. line: `data: {"type":"response","response":{"id":"resp-1","model":"gpt-4o","output":[]}}`,
  2892. from: "gpt-4o",
  2893. to: "my-model",
  2894. expected: `data: {"type":"response","response":{"id":"resp-1","model":"my-model","output":[]}}`,
  2895. },
  2896. {
  2897. name: "model 不匹配时不替换",
  2898. line: `data: {"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
  2899. from: "gpt-4o",
  2900. to: "my-model",
  2901. expected: `data: {"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
  2902. },
  2903. {
  2904. name: "无 model 字段时不替换",
  2905. line: `data: {"id":"chatcmpl-123","choices":[]}`,
  2906. from: "gpt-4o",
  2907. to: "my-model",
  2908. expected: `data: {"id":"chatcmpl-123","choices":[]}`,
  2909. },
  2910. {
  2911. name: "空 data 行",
  2912. line: `data: `,
  2913. from: "gpt-4o",
  2914. to: "my-model",
  2915. expected: `data: `,
  2916. },
  2917. {
  2918. name: "[DONE] 行",
  2919. line: `data: [DONE]`,
  2920. from: "gpt-4o",
  2921. to: "my-model",
  2922. expected: `data: [DONE]`,
  2923. },
  2924. {
  2925. name: "非 data: 前缀行",
  2926. line: `event: message`,
  2927. from: "gpt-4o",
  2928. to: "my-model",
  2929. expected: `event: message`,
  2930. },
  2931. {
  2932. name: "非法 JSON 不替换",
  2933. line: `data: {invalid json}`,
  2934. from: "gpt-4o",
  2935. to: "my-model",
  2936. expected: `data: {invalid json}`,
  2937. },
  2938. {
  2939. name: "无空格 data: 格式",
  2940. line: `data:{"id":"x","model":"gpt-4o"}`,
  2941. from: "gpt-4o",
  2942. to: "my-model",
  2943. expected: `data: {"id":"x","model":"my-model"}`,
  2944. },
  2945. {
  2946. name: "model 名含特殊字符",
  2947. line: `data: {"model":"org/model-v2.1-beta"}`,
  2948. from: "org/model-v2.1-beta",
  2949. to: "custom/alias",
  2950. expected: `data: {"model":"custom/alias"}`,
  2951. },
  2952. {
  2953. name: "空行",
  2954. line: "",
  2955. from: "gpt-4o",
  2956. to: "my-model",
  2957. expected: "",
  2958. },
  2959. {
  2960. name: "保持其他字段不变",
  2961. line: `data: {"id":"abc","object":"chat.completion.chunk","model":"gpt-4o","created":1234567890,"choices":[{"index":0,"delta":{"content":"hi"}}]}`,
  2962. from: "gpt-4o",
  2963. to: "alias",
  2964. expected: `data: {"id":"abc","object":"chat.completion.chunk","model":"alias","created":1234567890,"choices":[{"index":0,"delta":{"content":"hi"}}]}`,
  2965. },
  2966. {
  2967. name: "顶层优先于嵌套:同时存在两个 model",
  2968. line: `data: {"model":"gpt-4o","response":{"model":"gpt-4o"}}`,
  2969. from: "gpt-4o",
  2970. to: "replaced",
  2971. expected: `data: {"model":"replaced","response":{"model":"gpt-4o"}}`,
  2972. },
  2973. }
  2974. for _, tt := range tests {
  2975. t.Run(tt.name, func(t *testing.T) {
  2976. got := svc.replaceModelInSSELine(tt.line, tt.from, tt.to)
  2977. require.Equal(t, tt.expected, got)
  2978. })
  2979. }
  2980. }
  2981. func TestReplaceModelInSSEBody(t *testing.T) {
  2982. svc := &OpenAIGatewayService{}
  2983. tests := []struct {
  2984. name string
  2985. body string
  2986. from string
  2987. to string
  2988. expected string
  2989. }{
  2990. {
  2991. name: "多行 SSE body 替换",
  2992. body: "data: {\"model\":\"gpt-4o\",\"choices\":[]}\n\ndata: {\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n",
  2993. from: "gpt-4o",
  2994. to: "alias",
  2995. expected: "data: {\"model\":\"alias\",\"choices\":[]}\n\ndata: {\"model\":\"alias\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n",
  2996. },
  2997. {
  2998. name: "无需替换的 body",
  2999. body: "data: {\"model\":\"gpt-3.5-turbo\"}\n\ndata: [DONE]\n",
  3000. from: "gpt-4o",
  3001. to: "alias",
  3002. expected: "data: {\"model\":\"gpt-3.5-turbo\"}\n\ndata: [DONE]\n",
  3003. },
  3004. {
  3005. name: "混合 event 和 data 行",
  3006. body: "event: message\ndata: {\"model\":\"gpt-4o\"}\n\n",
  3007. from: "gpt-4o",
  3008. to: "alias",
  3009. expected: "event: message\ndata: {\"model\":\"alias\"}\n\n",
  3010. },
  3011. {
  3012. name: "空 body",
  3013. body: "",
  3014. from: "gpt-4o",
  3015. to: "alias",
  3016. expected: "",
  3017. },
  3018. }
  3019. for _, tt := range tests {
  3020. t.Run(tt.name, func(t *testing.T) {
  3021. got := svc.replaceModelInSSEBody(tt.body, tt.from, tt.to)
  3022. require.Equal(t, tt.expected, got)
  3023. })
  3024. }
  3025. }
  3026. func TestReplaceModelInResponseBody(t *testing.T) {
  3027. svc := &OpenAIGatewayService{}
  3028. tests := []struct {
  3029. name string
  3030. body string
  3031. from string
  3032. to string
  3033. expected string
  3034. }{
  3035. {
  3036. name: "替换顶层 model",
  3037. body: `{"id":"chatcmpl-123","model":"gpt-4o","choices":[]}`,
  3038. from: "gpt-4o",
  3039. to: "alias",
  3040. expected: `{"id":"chatcmpl-123","model":"alias","choices":[]}`,
  3041. },
  3042. {
  3043. name: "model 不匹配不替换",
  3044. body: `{"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
  3045. from: "gpt-4o",
  3046. to: "alias",
  3047. expected: `{"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
  3048. },
  3049. {
  3050. name: "无 model 字段不替换",
  3051. body: `{"id":"chatcmpl-123","choices":[]}`,
  3052. from: "gpt-4o",
  3053. to: "alias",
  3054. expected: `{"id":"chatcmpl-123","choices":[]}`,
  3055. },
  3056. {
  3057. name: "非法 JSON 返回原值",
  3058. body: `not json`,
  3059. from: "gpt-4o",
  3060. to: "alias",
  3061. expected: `not json`,
  3062. },
  3063. {
  3064. name: "空 body 返回原值",
  3065. body: ``,
  3066. from: "gpt-4o",
  3067. to: "alias",
  3068. expected: ``,
  3069. },
  3070. {
  3071. name: "保持嵌套结构不变",
  3072. body: `{"model":"gpt-4o","usage":{"prompt_tokens":10,"completion_tokens":20},"choices":[{"message":{"role":"assistant","content":"hello"}}]}`,
  3073. from: "gpt-4o",
  3074. to: "alias",
  3075. expected: `{"model":"alias","usage":{"prompt_tokens":10,"completion_tokens":20},"choices":[{"message":{"role":"assistant","content":"hello"}}]}`,
  3076. },
  3077. }
  3078. for _, tt := range tests {
  3079. t.Run(tt.name, func(t *testing.T) {
  3080. got := svc.replaceModelInResponseBody([]byte(tt.body), tt.from, tt.to)
  3081. require.Equal(t, tt.expected, string(got))
  3082. })
  3083. }
  3084. }
  3085. func TestExtractOpenAISSEDataLine(t *testing.T) {
  3086. tests := []struct {
  3087. name string
  3088. line string
  3089. wantData string
  3090. wantOK bool
  3091. }{
  3092. {name: "标准格式", line: `data: {"type":"x"}`, wantData: `{"type":"x"}`, wantOK: true},
  3093. {name: "无空格格式", line: `data:{"type":"x"}`, wantData: `{"type":"x"}`, wantOK: true},
  3094. {name: "纯空数据", line: `data: `, wantData: ``, wantOK: true},
  3095. {name: "非 data 行", line: `event: message`, wantData: ``, wantOK: false},
  3096. }
  3097. for _, tt := range tests {
  3098. t.Run(tt.name, func(t *testing.T) {
  3099. got, ok := extractOpenAISSEDataLine(tt.line)
  3100. require.Equal(t, tt.wantOK, ok)
  3101. require.Equal(t, tt.wantData, got)
  3102. })
  3103. }
  3104. }
  3105. func TestParseSSEUsage_SelectiveParsing(t *testing.T) {
  3106. svc := &OpenAIGatewayService{}
  3107. usage := &OpenAIUsage{InputTokens: 9, OutputTokens: 8, CacheReadInputTokens: 7}
  3108. // 非终态事件中的显式 usage 作为兼容 fallback,非零字段会被合并。
  3109. svc.parseSSEUsage(`{"type":"response.in_progress","response":{"usage":{"input_tokens":1,"output_tokens":2}}}`, usage)
  3110. require.Equal(t, 1, usage.InputTokens)
  3111. require.Equal(t, 2, usage.OutputTokens)
  3112. require.Equal(t, 7, usage.CacheReadInputTokens)
  3113. // completed 事件,应提取 usage
  3114. svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"input_tokens":3,"output_tokens":5,"input_tokens_details":{"cached_tokens":2}}}}`, usage)
  3115. require.Equal(t, 3, usage.InputTokens)
  3116. require.Equal(t, 5, usage.OutputTokens)
  3117. require.Equal(t, 2, usage.CacheReadInputTokens)
  3118. // done 事件同样可能携带最终 usage
  3119. svc.parseSSEUsage(`{"type":"response.done","response":{"usage":{"input_tokens":13,"output_tokens":15,"input_tokens_details":{"cached_tokens":4}}}}`, usage)
  3120. require.Equal(t, 13, usage.InputTokens)
  3121. require.Equal(t, 15, usage.OutputTokens)
  3122. require.Equal(t, 4, usage.CacheReadInputTokens)
  3123. // failed 事件在部分上游路径也会携带已消耗 usage,应与 WS/passthrough 保持一致
  3124. svc.parseSSEUsage(`{"type":"response.failed","response":{"usage":{"input_tokens":17,"output_tokens":19,"input_tokens_details":{"cached_tokens":6}}}}`, usage)
  3125. require.Equal(t, 17, usage.InputTokens)
  3126. require.Equal(t, 19, usage.OutputTokens)
  3127. require.Equal(t, 6, usage.CacheReadInputTokens)
  3128. svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"prompt_tokens":21,"completion_tokens":8,"prompt_tokens_details":{"cached_tokens":6}}}}`, usage)
  3129. require.Equal(t, 21, usage.InputTokens)
  3130. require.Equal(t, 8, usage.OutputTokens)
  3131. require.Equal(t, 6, usage.CacheReadInputTokens)
  3132. }
  3133. func TestParseSSEUsage_NonTerminalUsageMergesNonZeroFields(t *testing.T) {
  3134. svc := &OpenAIGatewayService{}
  3135. usage := &OpenAIUsage{}
  3136. svc.parseSSEUsage(`{"type":"response.in_progress","usage":{"input_tokens":17,"output_tokens":1,"input_tokens_details":{"cached_tokens":4}}}`, usage)
  3137. 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)
  3138. require.Equal(t, 17, usage.InputTokens)
  3139. require.Equal(t, 5, usage.OutputTokens)
  3140. require.Equal(t, 4, usage.CacheReadInputTokens)
  3141. require.Equal(t, 3, usage.CacheCreationInputTokens)
  3142. }
  3143. func TestParseSSEUsage_TerminalUsageReplacesFallback(t *testing.T) {
  3144. svc := &OpenAIGatewayService{}
  3145. usage := &OpenAIUsage{}
  3146. svc.parseSSEUsage(`{"type":"response.output_text.done","usage":{"input_tokens":17,"output_tokens":5,"input_tokens_details":{"cached_tokens":4}}}`, usage)
  3147. svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"input_tokens":19,"output_tokens":7}}}`, usage)
  3148. require.Equal(t, 19, usage.InputTokens)
  3149. require.Equal(t, 7, usage.OutputTokens)
  3150. require.Zero(t, usage.CacheReadInputTokens)
  3151. }
  3152. func TestParseSSEUsage_TerminalWithoutUsageKeepsFallback(t *testing.T) {
  3153. svc := &OpenAIGatewayService{}
  3154. usage := &OpenAIUsage{}
  3155. svc.parseSSEUsage(`{"type":"response.in_progress","usage":{"input_tokens":17,"output_tokens":5}}`, usage)
  3156. svc.parseSSEUsage(`{"type":"response.completed","response":{"id":"resp_1"}}`, usage)
  3157. svc.parseSSEUsage(" [DONE]\n", usage)
  3158. require.Equal(t, 17, usage.InputTokens)
  3159. require.Equal(t, 5, usage.OutputTokens)
  3160. }
  3161. func TestExtractOpenAIUsageFromJSONBytes_AcceptsResponseAndChatUsageShapes(t *testing.T) {
  3162. usage, ok := extractOpenAIUsageFromJSONBytes([]byte(`{"id":"resp_1","usage":{"input_tokens":9,"output_tokens":5,"input_tokens_details":{"cached_tokens":2,"cache_write_tokens":4}}}`))
  3163. require.True(t, ok)
  3164. require.Equal(t, 9, usage.InputTokens)
  3165. require.Equal(t, 5, usage.OutputTokens)
  3166. require.Equal(t, 2, usage.CacheReadInputTokens)
  3167. require.Equal(t, 4, usage.CacheCreationInputTokens)
  3168. 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}}}}`))
  3169. require.True(t, ok)
  3170. require.Equal(t, 13, usage.InputTokens)
  3171. require.Equal(t, 7, usage.OutputTokens)
  3172. require.Equal(t, 4, usage.CacheReadInputTokens)
  3173. require.Equal(t, 3, usage.CacheCreationInputTokens)
  3174. usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":11,"output_tokens":2,"cache_write_input_tokens":6}}`))
  3175. require.True(t, ok)
  3176. require.Equal(t, 6, usage.CacheCreationInputTokens)
  3177. usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_creation_input_tokens":19,"input_tokens_details":{"cache_write_tokens":7}}}`))
  3178. require.True(t, ok)
  3179. require.Equal(t, 7, usage.CacheCreationInputTokens, "官方嵌套字段应优先于兼容顶层别名")
  3180. usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_creation_input_tokens":19,"input_tokens_details":{"cache_write_tokens":0}}}`))
  3181. require.True(t, ok)
  3182. require.Zero(t, usage.CacheCreationInputTokens, "官方嵌套字段显式为零时仍应优先于兼容顶层别名")
  3183. usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_read_input_tokens":19,"input_tokens_details":{"cached_tokens":0}}}`))
  3184. require.True(t, ok)
  3185. require.Zero(t, usage.CacheReadInputTokens, "官方嵌套缓存读取字段显式为零时仍应优先于兼容顶层别名")
  3186. // xAI reports reasoning_tokens outside visible output_tokens. Only the
  3187. // arithmetic-consistent shape is independent; OpenAI's canonical shape
  3188. // already includes reasoning in completion/output_tokens.
  3189. usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":10000,"output_tokens":500,"total_tokens":10800,"output_tokens_details":{"reasoning_tokens":300}}}`))
  3190. require.True(t, ok)
  3191. require.Equal(t, 800, usage.OutputTokens)
  3192. usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":10000,"output_tokens":500,"total_tokens":10500,"output_tokens_details":{"reasoning_tokens":300}}}`))
  3193. require.True(t, ok)
  3194. require.Equal(t, 500, usage.OutputTokens)
  3195. }
  3196. func TestExtractOpenAIUsageFromJSONBytes_IncludesGrokReasoningTokens(t *testing.T) {
  3197. usage, ok := extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"prompt_tokens":32,"completion_tokens":9,"total_tokens":135,"completion_tokens_details":{"reasoning_tokens":94}}}`))
  3198. require.True(t, ok)
  3199. require.Equal(t, 103, usage.OutputTokens, "Grok Chat usage bills visible completion plus reasoning tokens")
  3200. usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":32,"output_tokens":103,"total_tokens":135,"output_tokens_details":{"reasoning_tokens":94}}}`))
  3201. require.True(t, ok)
  3202. require.Equal(t, 103, usage.OutputTokens, "Responses output_tokens already includes reasoning when total confirms it")
  3203. usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":32,"output_tokens":9,"total_tokens":135,"output_tokens_details":{"reasoning_tokens":94}}}`))
  3204. require.True(t, ok)
  3205. require.Equal(t, 103, usage.OutputTokens, "Responses detail-only shape is normalized when total exposes the full output")
  3206. }
  3207. func TestExtractCodexFinalResponse_SampleReplay(t *testing.T) {
  3208. body := strings.Join([]string{
  3209. `event: message`,
  3210. `data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
  3211. `data: {"type":"response.completed","response":{"id":"resp_1","model":"gpt-4o","usage":{"input_tokens":11,"output_tokens":22,"input_tokens_details":{"cached_tokens":3}}}}`,
  3212. `data: [DONE]`,
  3213. }, "\n")
  3214. finalResp, ok := extractCodexFinalResponse(body)
  3215. require.True(t, ok)
  3216. require.Contains(t, string(finalResp), `"id":"resp_1"`)
  3217. require.Contains(t, string(finalResp), `"input_tokens":11`)
  3218. }
  3219. func TestHandleSSEToJSON_CompletedEventReturnsJSON(t *testing.T) {
  3220. gin.SetMode(gin.TestMode)
  3221. rec := httptest.NewRecorder()
  3222. c, _ := gin.CreateTestContext(rec)
  3223. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  3224. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  3225. resp := &http.Response{
  3226. StatusCode: http.StatusOK,
  3227. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  3228. }
  3229. body := []byte(strings.Join([]string{
  3230. `data: {"type":"response.in_progress","response":{"id":"resp_2"}}`,
  3231. `data: {"type":"response.completed","response":{"id":"resp_2","model":"gpt-4o","usage":{"input_tokens":7,"output_tokens":9,"input_tokens_details":{"cached_tokens":1}}}}`,
  3232. `data: [DONE]`,
  3233. }, "\n"))
  3234. usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-4o", "gpt-4o")
  3235. require.NoError(t, err)
  3236. require.NotNil(t, usage)
  3237. require.Equal(t, 7, usage.InputTokens)
  3238. require.Equal(t, 9, usage.OutputTokens)
  3239. require.Equal(t, 1, usage.CacheReadInputTokens)
  3240. // Header 可能由上游 Content-Type 透传;关键是 body 已转换为最终 JSON 响应。
  3241. require.NotContains(t, rec.Body.String(), "event:")
  3242. require.Contains(t, rec.Body.String(), `"id":"resp_2"`)
  3243. require.NotContains(t, rec.Body.String(), "data:")
  3244. }
  3245. func TestHandleNonStreamingResponse_APIKeyFallsBackToSSEBodyWhenContentTypeIsWrong(t *testing.T) {
  3246. gin.SetMode(gin.TestMode)
  3247. rec := httptest.NewRecorder()
  3248. c, _ := gin.CreateTestContext(rec)
  3249. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  3250. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  3251. resp := &http.Response{
  3252. StatusCode: http.StatusOK,
  3253. Header: http.Header{"Content-Type": []string{"application/json"}},
  3254. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  3255. `data: {"type":"response.output_text.delta","delta":"hel"}`,
  3256. `data: {"type":"response.output_text.delta","delta":"lo"}`,
  3257. `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}}}`,
  3258. `data: [DONE]`,
  3259. }, "\n"))),
  3260. }
  3261. account := &Account{ID: 1, Type: AccountTypeAPIKey}
  3262. result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.4", "gpt-5.4")
  3263. require.NoError(t, err)
  3264. require.NotNil(t, result)
  3265. require.Equal(t, 3, result.InputTokens)
  3266. require.Equal(t, 2, result.OutputTokens)
  3267. require.NotContains(t, rec.Body.String(), "data:")
  3268. require.Equal(t, "resp_api_key_sse", gjson.Get(rec.Body.String(), "id").String())
  3269. require.Equal(t, "hello", gjson.Get(rec.Body.String(), "output.0.content.0.text").String())
  3270. }
  3271. func TestHandleNonStreamingResponse_OAuthJSONBodyWithDataEventTextKeepsJSONUsage(t *testing.T) {
  3272. gin.SetMode(gin.TestMode)
  3273. rec := httptest.NewRecorder()
  3274. c, _ := gin.CreateTestContext(rec)
  3275. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", nil)
  3276. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  3277. // Plain JSON compact response whose output text happens to contain the
  3278. // literal substrings "data:" and "event:" (e.g. echoing shell/log output).
  3279. // This must NOT be misdetected as SSE framing: it has a top-level usage
  3280. // object and no upstream text/event-stream Content-Type.
  3281. jsonBody := `{"id":"resp_oauth_compact","object":"response","model":"gpt-5.4","status":"completed",` +
  3282. `"output":[{"type":"message","content":[{"type":"output_text",` +
  3283. `"text":"processing data: 1,2,3 then event: click finished"}]}],` +
  3284. `"usage":{"input_tokens":11,"output_tokens":22,"total_tokens":33}}`
  3285. resp := &http.Response{
  3286. StatusCode: http.StatusOK,
  3287. Header: http.Header{"Content-Type": []string{"application/json"}},
  3288. Body: io.NopCloser(strings.NewReader(jsonBody)),
  3289. }
  3290. account := &Account{ID: 146, Type: AccountTypeOAuth}
  3291. result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.4", "gpt-5.4")
  3292. require.NoError(t, err)
  3293. require.NotNil(t, result)
  3294. require.Equal(t, 11, result.InputTokens)
  3295. require.Equal(t, 22, result.OutputTokens)
  3296. // Response must remain the original JSON body (not routed through the SSE
  3297. // path, which would rewrite/lose the body or usage).
  3298. require.Equal(t, "application/json", rec.Header().Get("Content-Type"))
  3299. require.Equal(t, "resp_oauth_compact", gjson.Get(rec.Body.String(), "id").String())
  3300. require.Equal(t, int64(33), gjson.Get(rec.Body.String(), "usage.total_tokens").Int())
  3301. require.Contains(t, rec.Body.String(), "processing data: 1,2,3 then event: click finished")
  3302. }
  3303. func TestHandleNonStreamingResponse_ObservesUpstreamModelBeforeClientRewrite(t *testing.T) {
  3304. gin.SetMode(gin.TestMode)
  3305. rec := httptest.NewRecorder()
  3306. c, _ := gin.CreateTestContext(rec)
  3307. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  3308. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  3309. resp := &http.Response{
  3310. StatusCode: http.StatusOK,
  3311. Header: http.Header{"Content-Type": []string{"application/json"}},
  3312. Body: io.NopCloser(strings.NewReader(
  3313. `{"id":"resp_model_audit","object":"response","model":"gpt-5.5","status":"completed","output":[],"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}`,
  3314. )),
  3315. }
  3316. account := &Account{ID: 1, Type: AccountTypeAPIKey}
  3317. result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.6-sol", "gpt-5.5")
  3318. require.NoError(t, err)
  3319. require.NotNil(t, result)
  3320. // 客户端仍看到自己请求的模型名,审计观察器则保留改写前的上游声明。
  3321. require.Equal(t, "gpt-5.6-sol", gjson.Get(rec.Body.String(), "model").String())
  3322. require.Equal(t, "gpt-5.5", observedUpstreamResponseModel(c))
  3323. require.False(t, observedUpstreamResponseModelConflict(c))
  3324. }
  3325. func TestHandleSSEToJSON_ReconstructsImageGenerationOutputItemDone(t *testing.T) {
  3326. gin.SetMode(gin.TestMode)
  3327. rec := httptest.NewRecorder()
  3328. c, _ := gin.CreateTestContext(rec)
  3329. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  3330. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  3331. resp := &http.Response{
  3332. StatusCode: http.StatusOK,
  3333. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  3334. }
  3335. body := []byte(strings.Join([]string{
  3336. `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"}}`,
  3337. `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}}}}`,
  3338. `data: [DONE]`,
  3339. }, "\n"))
  3340. usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-5.4", "gpt-5.4")
  3341. require.NoError(t, err)
  3342. require.NotNil(t, usage)
  3343. require.Equal(t, 4, usage.ImageOutputTokens)
  3344. require.NotContains(t, rec.Body.String(), "data:")
  3345. require.Equal(t, "image_generation_call", gjson.Get(rec.Body.String(), "output.0.type").String())
  3346. require.Equal(t, "completed", gjson.Get(rec.Body.String(), "output.0.status").String())
  3347. require.Equal(t, "aGVsbG8=", gjson.Get(rec.Body.String(), "output.0.result").String())
  3348. require.Equal(t, "draw a cat", gjson.Get(rec.Body.String(), "output.0.revised_prompt").String())
  3349. }
  3350. func TestHandleSSEToJSON_NoFinalResponseKeepsSSEBody(t *testing.T) {
  3351. gin.SetMode(gin.TestMode)
  3352. rec := httptest.NewRecorder()
  3353. c, _ := gin.CreateTestContext(rec)
  3354. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  3355. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  3356. resp := &http.Response{
  3357. StatusCode: http.StatusOK,
  3358. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  3359. }
  3360. body := []byte(strings.Join([]string{
  3361. `data: {"type":"response.in_progress","response":{"id":"resp_3"}}`,
  3362. `data: [DONE]`,
  3363. }, "\n"))
  3364. usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-4o", "gpt-4o")
  3365. require.NoError(t, err)
  3366. require.NotNil(t, usage)
  3367. require.Equal(t, 0, usage.InputTokens)
  3368. require.Contains(t, rec.Header().Get("Content-Type"), "text/event-stream")
  3369. require.Contains(t, rec.Body.String(), `data: {"type":"response.in_progress"`)
  3370. }
  3371. // 无账号时没有可换的对象:newOpenAIStreamFailoverError 要拿 account 记录 ops 归属与
  3372. // 账号健康,故这一支保持原有的协议错误行为。带真实账号的同一报文改为换号,
  3373. // 由 TestNonStreamingSSEToJSON_UnclassifiedFailedEventFailsOver 钉死(issue #5281)。
  3374. func TestHandleSSEToJSON_ResponseFailedWithoutAccountReturnsProtocolError(t *testing.T) {
  3375. gin.SetMode(gin.TestMode)
  3376. rec := httptest.NewRecorder()
  3377. c, _ := gin.CreateTestContext(rec)
  3378. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  3379. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  3380. resp := &http.Response{
  3381. StatusCode: http.StatusOK,
  3382. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  3383. }
  3384. body := []byte(strings.Join([]string{
  3385. `data: {"type":"response.failed","error":{"message":"upstream rejected request"}}`,
  3386. `data: [DONE]`,
  3387. }, "\n"))
  3388. usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-4o", "gpt-4o")
  3389. require.Nil(t, usage)
  3390. require.Error(t, err)
  3391. require.Equal(t, http.StatusBadGateway, rec.Code)
  3392. require.Contains(t, rec.Body.String(), "upstream rejected request")
  3393. require.Contains(t, rec.Header().Get("Content-Type"), "application/json")
  3394. }
  3395. func TestOpenAICompatSSEFrameParserResetsEventTypeAtFrameBoundary(t *testing.T) {
  3396. var parser openAICompatSSEFrameParser
  3397. frame, ok := parser.AddLine("event: response.created")
  3398. require.False(t, ok)
  3399. require.Empty(t, frame)
  3400. frame, ok = parser.AddLine(`data: {"response":{"id":"resp_1"}}`)
  3401. require.False(t, ok)
  3402. require.Empty(t, frame)
  3403. frame, ok = parser.AddLine("")
  3404. require.True(t, ok)
  3405. require.Equal(t, "response.created", frame.EventType)
  3406. require.JSONEq(t, `{"response":{"id":"resp_1"}}`, frame.Data)
  3407. frame, ok = parser.AddLine(`data: {"delta":"ok"}`)
  3408. require.False(t, ok)
  3409. require.Empty(t, frame.EventType)
  3410. frame, ok = parser.AddLine("")
  3411. require.True(t, ok)
  3412. require.Empty(t, frame.EventType)
  3413. require.JSONEq(t, `{"delta":"ok"}`, frame.Data)
  3414. }
  3415. func TestStreamingPassthroughCyberPolicyMarksAndPassesThrough(t *testing.T) {
  3416. gin.SetMode(gin.TestMode)
  3417. cfg := &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}
  3418. svc := &OpenAIGatewayService{cfg: cfg}
  3419. rec := httptest.NewRecorder()
  3420. c, _ := gin.CreateTestContext(rec)
  3421. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  3422. resp := &http.Response{
  3423. StatusCode: http.StatusOK,
  3424. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  3425. "event: response.created",
  3426. `data: {"type":"response.created","response":{"id":"r1"}}`,
  3427. "",
  3428. "event: response.failed",
  3429. `data: {"type":"response.failed","response":{"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}}`,
  3430. "",
  3431. }, "\n"))),
  3432. Header: http.Header{"X-Request-Id": []string{"rid-cyber"}},
  3433. }
  3434. _, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, time.Now(), "m", "m")
  3435. require.Error(t, err)
  3436. var failoverErr *UpstreamFailoverError
  3437. require.False(t, errors.As(err, &failoverErr), "cyber must NOT failover")
  3438. require.Contains(t, rec.Body.String(), "cyber_policy", "response.failed passed through to client")
  3439. mark := GetOpsCyberPolicy(c)
  3440. require.NotNil(t, mark)
  3441. require.Equal(t, "flagged for cyber policy", mark.Message)
  3442. }
  3443. func TestHandleStreamingResponseCyberPolicyMarks(t *testing.T) {
  3444. gin.SetMode(gin.TestMode)
  3445. cfg := &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}
  3446. svc := &OpenAIGatewayService{cfg: cfg}
  3447. rec := httptest.NewRecorder()
  3448. c, _ := gin.CreateTestContext(rec)
  3449. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  3450. resp := &http.Response{
  3451. StatusCode: http.StatusOK,
  3452. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  3453. "event: response.created",
  3454. `data: {"type":"response.created","response":{"id":"r1"}}`,
  3455. "",
  3456. "event: response.failed",
  3457. `data: {"type":"response.failed","error":{"code":"cyber_policy","message":"flagged"}}`,
  3458. "",
  3459. }, "\n"))),
  3460. Header: http.Header{"X-Request-Id": []string{"rid"}},
  3461. }
  3462. _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, time.Now(), "m", "m")
  3463. require.Error(t, err)
  3464. var fo *UpstreamFailoverError
  3465. require.False(t, errors.As(err, &fo))
  3466. require.NotNil(t, GetOpsCyberPolicy(c))
  3467. }
  3468. func TestHandleErrorResponseCyberPolicyPassthrough(t *testing.T) {
  3469. gin.SetMode(gin.TestMode)
  3470. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  3471. rec := httptest.NewRecorder()
  3472. c, _ := gin.CreateTestContext(rec)
  3473. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  3474. cyberBody := `{"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}`
  3475. resp := &http.Response{
  3476. StatusCode: http.StatusBadRequest,
  3477. Header: http.Header{"Content-Type": []string{"application/json"}, "X-Request-Id": []string{"rid"}},
  3478. Body: io.NopCloser(strings.NewReader(cyberBody)),
  3479. }
  3480. _, err := svc.handleErrorResponse(context.Background(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, nil)
  3481. require.Error(t, err)
  3482. require.Equal(t, http.StatusBadRequest, rec.Code, "passthrough upstream 400, not rewrapped 502")
  3483. require.Contains(t, rec.Body.String(), "cyber_policy", "client sees original cyber body")
  3484. require.NotContains(t, rec.Body.String(), "Upstream request failed", "must not 502-rewrap")
  3485. mark := GetOpsCyberPolicy(c)
  3486. require.NotNil(t, mark)
  3487. require.Equal(t, http.StatusBadRequest, mark.UpstreamStatus)
  3488. }
  3489. func TestHandleCompatErrorResponseCyberPolicyEarlyReturn(t *testing.T) {
  3490. gin.SetMode(gin.TestMode)
  3491. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  3492. rec := httptest.NewRecorder()
  3493. c, _ := gin.CreateTestContext(rec)
  3494. c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
  3495. cyberBody := `{"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}`
  3496. resp := &http.Response{
  3497. StatusCode: http.StatusBadRequest,
  3498. Header: http.Header{"Content-Type": []string{"application/json"}},
  3499. Body: io.NopCloser(strings.NewReader(cyberBody)),
  3500. }
  3501. var gotStatus int
  3502. var gotType, gotMsg string
  3503. writeError := func(_ *gin.Context, statusCode int, errType, message string) {
  3504. gotStatus, gotType, gotMsg = statusCode, errType, message
  3505. }
  3506. // cyber 命中应早返回(写兼容错误 + 不冷却账号),而非落到通用 "Upstream request failed"。
  3507. _, err := svc.handleCompatErrorResponse(resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, writeError)
  3508. require.Error(t, err)
  3509. require.Equal(t, http.StatusBadRequest, gotStatus)
  3510. require.Equal(t, "invalid_request_error", gotType)
  3511. require.Contains(t, gotMsg, "flagged for cyber policy")
  3512. require.NotContains(t, gotMsg, "Upstream request failed")
  3513. require.NotNil(t, GetOpsCyberPolicy(c))
  3514. }