openai_ws_forwarder_ingress_session_test.go 170 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546354735483549355035513552355335543555355635573558355935603561356235633564356535663567356835693570357135723573357435753576357735783579358035813582358335843585358635873588358935903591359235933594359535963597359835993600360136023603360436053606360736083609361036113612361336143615361636173618361936203621362236233624362536263627362836293630363136323633363436353636363736383639364036413642364336443645364636473648364936503651365236533654365536563657365836593660366136623663366436653666366736683669367036713672367336743675367636773678367936803681368236833684368536863687368836893690369136923693369436953696369736983699370037013702370337043705370637073708370937103711371237133714371537163717371837193720372137223723372437253726372737283729373037313732373337343735373637373738373937403741374237433744374537463747374837493750375137523753375437553756375737583759376037613762376337643765376637673768376937703771377237733774377537763777377837793780378137823783378437853786378737883789379037913792379337943795379637973798379938003801380238033804380538063807380838093810381138123813381438153816381738183819382038213822382338243825382638273828382938303831383238333834383538363837383838393840384138423843384438453846384738483849385038513852385338543855385638573858385938603861386238633864386538663867386838693870387138723873387438753876387738783879388038813882388338843885388638873888388938903891389238933894389538963897389838993900390139023903390439053906390739083909391039113912391339143915391639173918391939203921392239233924392539263927392839293930393139323933393439353936393739383939394039413942394339443945394639473948394939503951395239533954395539563957395839593960396139623963396439653966396739683969397039713972397339743975397639773978397939803981398239833984398539863987398839893990399139923993399439953996399739983999400040014002400340044005400640074008400940104011401240134014401540164017401840194020402140224023402440254026402740284029403040314032403340344035403640374038403940404041404240434044404540464047404840494050405140524053405440554056405740584059406040614062406340644065406640674068406940704071407240734074407540764077407840794080408140824083408440854086408740884089409040914092409340944095409640974098409941004101410241034104410541064107410841094110411141124113411441154116411741184119412041214122412341244125412641274128412941304131413241334134413541364137413841394140414141424143414441454146414741484149415041514152415341544155415641574158415941604161416241634164416541664167416841694170417141724173417441754176417741784179418041814182418341844185418641874188418941904191419241934194419541964197419841994200420142024203420442054206420742084209421042114212421342144215421642174218421942204221422242234224422542264227422842294230423142324233423442354236423742384239424042414242424342444245424642474248424942504251425242534254425542564257425842594260426142624263426442654266426742684269427042714272427342744275427642774278427942804281428242834284428542864287428842894290429142924293429442954296429742984299430043014302430343044305430643074308430943104311431243134314431543164317431843194320432143224323432443254326432743284329433043314332433343344335433643374338433943404341434243434344434543464347434843494350435143524353435443554356435743584359436043614362436343644365436643674368436943704371437243734374437543764377437843794380438143824383438443854386438743884389439043914392439343944395439643974398439944004401440244034404440544064407440844094410441144124413441444154416441744184419442044214422442344244425442644274428442944304431443244334434443544364437443844394440444144424443444444454446444744484449445044514452445344544455445644574458445944604461446244634464446544664467446844694470447144724473447444754476447744784479448044814482448344844485448644874488448944904491449244934494449544964497449844994500450145024503450445054506450745084509451045114512451345144515451645174518451945204521452245234524452545264527452845294530453145324533453445354536453745384539454045414542454345444545454645474548454945504551455245534554455545564557455845594560456145624563456445654566456745684569457045714572457345744575457645774578457945804581458245834584458545864587458845894590459145924593459445954596459745984599460046014602460346044605460646074608460946104611461246134614461546164617461846194620462146224623462446254626462746284629463046314632463346344635463646374638463946404641464246434644464546464647464846494650465146524653465446554656465746584659466046614662466346644665466646674668466946704671467246734674467546764677467846794680468146824683468446854686468746884689469046914692469346944695469646974698469947004701470247034704470547064707470847094710471147124713471447154716471747184719472047214722472347244725472647274728472947304731473247334734473547364737
  1. package service
  2. import (
  3. "context"
  4. "errors"
  5. "io"
  6. "net/http"
  7. "net/http/httptest"
  8. "strings"
  9. "sync"
  10. "testing"
  11. "time"
  12. "github.com/Wei-Shaw/sub2api/internal/config"
  13. coderws "github.com/coder/websocket"
  14. "github.com/gin-gonic/gin"
  15. "github.com/stretchr/testify/require"
  16. "github.com/tidwall/gjson"
  17. )
  18. type openAIWSLeaseLossAfterReadConn struct {
  19. *openAIWSCaptureConn
  20. cancel context.CancelCauseFunc
  21. once sync.Once
  22. }
  23. func (c *openAIWSLeaseLossAfterReadConn) ReadMessage(ctx context.Context) ([]byte, error) {
  24. message, err := c.openAIWSCaptureConn.ReadMessage(ctx)
  25. if err == nil {
  26. c.once.Do(func() {
  27. c.cancel(ErrOpenAIWSIngressLeaseLost)
  28. })
  29. }
  30. return message, err
  31. }
  32. type openAIWSSingleConnDialer struct {
  33. conn openAIWSClientConn
  34. }
  35. func (d *openAIWSSingleConnDialer) Dial(
  36. ctx context.Context,
  37. wsURL string,
  38. headers http.Header,
  39. proxyURL string,
  40. ) (openAIWSClientConn, int, http.Header, error) {
  41. return d.conn, 0, nil, nil
  42. }
  43. func TestOpenAIWSDownstreamWriteContext_CancellationOwnership(t *testing.T) {
  44. t.Run("pre-canceled ordinary context is canceled before return", func(t *testing.T) {
  45. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  46. cancelControl(context.Canceled)
  47. writeCtx, cancelWrite := newOpenAIWSDownstreamWriteContext(controlCtx, nil, time.Second)
  48. defer cancelWrite()
  49. require.ErrorIs(t, writeCtx.Err(), context.Canceled)
  50. })
  51. t.Run("lease loss keeps current write alive", func(t *testing.T) {
  52. lifecycleCtx, cancelLifecycle := context.WithCancelCause(context.Background())
  53. controlCtx, cancelControl := context.WithCancelCause(lifecycleCtx)
  54. hooks := &OpenAIWSIngressHooks{ClientLifecycleContext: lifecycleCtx}
  55. writeCtx, cancelWrite := newOpenAIWSDownstreamWriteContext(controlCtx, hooks, time.Second)
  56. defer cancelWrite()
  57. cancelControl(ErrOpenAIWSIngressLeaseLost)
  58. select {
  59. case <-writeCtx.Done():
  60. t.Fatalf("lease loss unexpectedly canceled downstream write: %v", writeCtx.Err())
  61. case <-time.After(20 * time.Millisecond):
  62. }
  63. clientDisconnected := errors.New("client disconnected")
  64. cancelLifecycle(clientDisconnected)
  65. <-writeCtx.Done()
  66. require.ErrorIs(t, context.Cause(writeCtx), clientDisconnected)
  67. })
  68. t.Run("ordinary cancellation is direct and preserves cause", func(t *testing.T) {
  69. lifecycleCtx, cancelLifecycle := context.WithCancelCause(context.Background())
  70. controlCtx, cancelControl := context.WithCancelCause(lifecycleCtx)
  71. defer cancelControl(context.Canceled)
  72. hooks := &OpenAIWSIngressHooks{ClientLifecycleContext: lifecycleCtx}
  73. writeCtx, cancelWrite := newOpenAIWSDownstreamWriteContext(controlCtx, hooks, time.Second)
  74. defer cancelWrite()
  75. serverShutdown := errors.New("server shutdown")
  76. cancelLifecycle(serverShutdown)
  77. require.ErrorIs(t, writeCtx.Err(), context.Canceled)
  78. require.ErrorIs(t, context.Cause(writeCtx), serverShutdown)
  79. })
  80. }
  81. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossTurns(t *testing.T) {
  82. gin.SetMode(gin.TestMode)
  83. cfg := &config.Config{}
  84. cfg.Security.URLAllowlist.Enabled = false
  85. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  86. cfg.Gateway.OpenAIWS.Enabled = true
  87. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  88. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  89. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  90. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  91. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  92. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  93. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  94. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  95. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  96. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  97. captureConn := &openAIWSCaptureConn{
  98. events: [][]byte{
  99. []byte(`{"type":"response.output_item.done","item":{"id":"ig_ingress_1","type":"image_generation_call","status":"generating","result":"iVBORw0KGgoAAAANSUhEUg/+=="}}`),
  100. []byte(`{"type":"response.completed","response":{"id":"resp_ingress_turn_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  101. []byte(`{"type":"response.completed","response":{"id":"resp_ingress_turn_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  102. },
  103. }
  104. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  105. pool := newOpenAIWSConnPool(cfg)
  106. pool.setClientDialerForTest(captureDialer)
  107. svc := &OpenAIGatewayService{
  108. cfg: cfg,
  109. httpUpstream: &httpUpstreamRecorder{},
  110. cache: &stubGatewayCache{},
  111. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  112. toolCorrector: NewCodexToolCorrector(),
  113. openaiWSPool: pool,
  114. }
  115. account := &Account{
  116. ID: 114,
  117. Name: "openai-ingress-session-lease",
  118. Platform: PlatformOpenAI,
  119. Type: AccountTypeAPIKey,
  120. Status: StatusActive,
  121. Schedulable: true,
  122. Concurrency: 1,
  123. Credentials: map[string]any{
  124. "api_key": "sk-test",
  125. },
  126. Extra: map[string]any{
  127. "responses_websockets_v2_enabled": true,
  128. },
  129. }
  130. serverErrCh := make(chan error, 1)
  131. turnTerminalCh := make(chan string, 2)
  132. hooks := &OpenAIWSIngressHooks{
  133. AfterTurn: func(_ int, result *OpenAIForwardResult, turnErr error) {
  134. if turnErr == nil && result != nil {
  135. turnTerminalCh <- result.UpstreamTerminalEvent
  136. }
  137. },
  138. }
  139. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  140. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  141. CompressionMode: coderws.CompressionContextTakeover,
  142. })
  143. if err != nil {
  144. serverErrCh <- err
  145. return
  146. }
  147. defer func() {
  148. _ = conn.CloseNow()
  149. }()
  150. rec := httptest.NewRecorder()
  151. ginCtx, _ := gin.CreateTestContext(rec)
  152. req := r.Clone(r.Context())
  153. req.Header = req.Header.Clone()
  154. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  155. ginCtx.Request = req
  156. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  157. msgType, firstMessage, readErr := conn.Read(readCtx)
  158. cancel()
  159. if readErr != nil {
  160. serverErrCh <- readErr
  161. return
  162. }
  163. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  164. serverErrCh <- errors.New("unsupported websocket client message type")
  165. return
  166. }
  167. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, hooks)
  168. }))
  169. defer wsServer.Close()
  170. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  171. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  172. cancelDial()
  173. require.NoError(t, err)
  174. defer func() {
  175. _ = clientConn.CloseNow()
  176. }()
  177. writeMessage := func(payload string) {
  178. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  179. defer cancel()
  180. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  181. }
  182. readMessage := func() []byte {
  183. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  184. defer cancel()
  185. msgType, message, readErr := clientConn.Read(readCtx)
  186. require.NoError(t, readErr)
  187. require.Equal(t, coderws.MessageText, msgType)
  188. return message
  189. }
  190. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false}`)
  191. firstTurnImageEvent := readMessage()
  192. require.Equal(t, "response.output_item.done", gjson.GetBytes(firstTurnImageEvent, "type").String())
  193. require.Equal(t, "completed", gjson.GetBytes(firstTurnImageEvent, "item.status").String())
  194. require.Equal(t, "iVBORw0KGgoAAAANSUhEUg/+==", gjson.GetBytes(firstTurnImageEvent, "item.result").String())
  195. firstTurnEvent := readMessage()
  196. require.Equal(t, "response.completed", gjson.GetBytes(firstTurnEvent, "type").String())
  197. require.Equal(t, "resp_ingress_turn_1", gjson.GetBytes(firstTurnEvent, "response.id").String())
  198. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_ingress_turn_1"}`)
  199. secondTurnEvent := readMessage()
  200. require.Equal(t, "response.completed", gjson.GetBytes(secondTurnEvent, "type").String())
  201. require.Equal(t, "resp_ingress_turn_2", gjson.GetBytes(secondTurnEvent, "response.id").String())
  202. require.Equal(t, "response.completed", <-turnTerminalCh, "首轮 turn 应保留成功终态")
  203. require.Equal(t, "response.completed", <-turnTerminalCh, "第二轮 turn 应保留成功终态")
  204. _ = clientConn.Close(coderws.StatusNormalClosure, "done")
  205. select {
  206. case serverErr := <-serverErrCh:
  207. require.NoError(t, serverErr)
  208. case <-time.After(5 * time.Second):
  209. t.Fatal("等待 ingress websocket 结束超时")
  210. }
  211. metrics := svc.SnapshotOpenAIWSPoolMetrics()
  212. require.Equal(t, int64(1), metrics.AcquireTotal, "同一 ingress 会话多 turn 应只获取一次上游 lease")
  213. require.Equal(t, 1, captureDialer.DialCount(), "同一 ingress 会话应保持同一上游连接")
  214. require.Len(t, captureConn.writes, 2, "应向同一上游连接发送两轮 response.create")
  215. }
  216. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_LeaseLossSendsRetryClose(t *testing.T) {
  217. gin.SetMode(gin.TestMode)
  218. cfg := &config.Config{}
  219. cfg.Security.URLAllowlist.Enabled = false
  220. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  221. cfg.Gateway.OpenAIWS.Enabled = true
  222. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  223. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  224. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  225. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  226. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  227. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  228. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  229. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  230. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  231. lifecycleCtx, cancelLifecycle := context.WithCancelCause(context.Background())
  232. defer cancelLifecycle(context.Canceled)
  233. controlCtx, cancelControl := context.WithCancelCause(lifecycleCtx)
  234. upstreamConn := &openAIWSLeaseLossAfterReadConn{
  235. openAIWSCaptureConn: &openAIWSCaptureConn{events: [][]byte{
  236. []byte(`{"type":"response.completed","response":{"id":"resp_lease_loss","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  237. }},
  238. cancel: cancelControl,
  239. }
  240. pool := newOpenAIWSConnPool(cfg)
  241. pool.setClientDialerForTest(&openAIWSSingleConnDialer{conn: upstreamConn})
  242. defer pool.Close()
  243. svc := &OpenAIGatewayService{
  244. cfg: cfg,
  245. httpUpstream: &httpUpstreamRecorder{},
  246. cache: &stubGatewayCache{},
  247. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  248. toolCorrector: NewCodexToolCorrector(),
  249. openaiWSPool: pool,
  250. }
  251. account := &Account{
  252. ID: 118,
  253. Name: "openai-ingress-lease-loss",
  254. Platform: PlatformOpenAI,
  255. Type: AccountTypeAPIKey,
  256. Status: StatusActive,
  257. Schedulable: true,
  258. Concurrency: 1,
  259. Credentials: map[string]any{"api_key": "sk-test"},
  260. Extra: map[string]any{"responses_websockets_v2_enabled": true},
  261. }
  262. serverErrCh := make(chan error, 1)
  263. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  264. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  265. if err != nil {
  266. serverErrCh <- err
  267. return
  268. }
  269. defer func() {
  270. _ = conn.CloseNow()
  271. }()
  272. rec := httptest.NewRecorder()
  273. ginCtx, _ := gin.CreateTestContext(rec)
  274. req := r.Clone(controlCtx)
  275. req.Header = req.Header.Clone()
  276. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  277. ginCtx.Request = req
  278. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  279. msgType, firstMessage, readErr := conn.Read(readCtx)
  280. cancelRead()
  281. if readErr != nil {
  282. serverErrCh <- readErr
  283. return
  284. }
  285. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  286. serverErrCh <- errors.New("unsupported websocket client message type")
  287. return
  288. }
  289. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(
  290. controlCtx,
  291. ginCtx,
  292. conn,
  293. account,
  294. "sk-test",
  295. firstMessage,
  296. &OpenAIWSIngressHooks{ClientLifecycleContext: lifecycleCtx},
  297. )
  298. }))
  299. defer wsServer.Close()
  300. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  301. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  302. cancelDial()
  303. require.NoError(t, err)
  304. defer func() {
  305. _ = clientConn.CloseNow()
  306. }()
  307. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  308. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`))
  309. cancelWrite()
  310. require.NoError(t, err)
  311. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  312. msgType, event, err := clientConn.Read(readCtx)
  313. cancelRead()
  314. require.NoError(t, err)
  315. require.Equal(t, coderws.MessageText, msgType)
  316. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  317. closeReadCtx, cancelCloseRead := context.WithTimeout(context.Background(), 3*time.Second)
  318. _, _, err = clientConn.Read(closeReadCtx)
  319. cancelCloseRead()
  320. var closeErr coderws.CloseError
  321. require.ErrorAs(t, err, &closeErr)
  322. require.Equal(t, coderws.StatusTryAgainLater, closeErr.Code)
  323. require.Equal(t, "websocket ingress capacity lease lost; please reconnect", closeErr.Reason)
  324. select {
  325. case serverErr := <-serverErrCh:
  326. var clientCloseErr *OpenAIWSClientCloseError
  327. require.ErrorAs(t, serverErr, &clientCloseErr)
  328. require.Equal(t, coderws.StatusTryAgainLater, clientCloseErr.StatusCode())
  329. require.ErrorIs(t, serverErr, ErrOpenAIWSIngressLeaseLost)
  330. case <-time.After(3 * time.Second):
  331. t.Fatal("ingress lease-loss reader did not exit")
  332. }
  333. }
  334. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_IdleTimeoutReleasesStoreDisabledSession(t *testing.T) {
  335. gin.SetMode(gin.TestMode)
  336. cfg := &config.Config{}
  337. cfg.Security.URLAllowlist.Enabled = false
  338. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  339. cfg.Gateway.OpenAIWS.Enabled = true
  340. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  341. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  342. cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 1
  343. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  344. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  345. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  346. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  347. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  348. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  349. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  350. captureConn := &openAIWSCaptureConn{events: [][]byte{
  351. []byte(`{"type":"response.completed","response":{"id":"resp_idle_timeout","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  352. }}
  353. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  354. pool := newOpenAIWSConnPool(cfg)
  355. pool.setClientDialerForTest(captureDialer)
  356. defer pool.Close()
  357. svc := &OpenAIGatewayService{
  358. cfg: cfg,
  359. httpUpstream: &httpUpstreamRecorder{},
  360. cache: &stubGatewayCache{},
  361. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  362. toolCorrector: NewCodexToolCorrector(),
  363. openaiWSPool: pool,
  364. }
  365. account := &Account{
  366. ID: 116,
  367. Name: "openai-ingress-idle-timeout",
  368. Platform: PlatformOpenAI,
  369. Type: AccountTypeAPIKey,
  370. Status: StatusActive,
  371. Schedulable: true,
  372. Concurrency: 1,
  373. Credentials: map[string]any{"api_key": "sk-test"},
  374. Extra: map[string]any{"responses_websockets_v2_enabled": true},
  375. }
  376. serverErrCh := make(chan error, 1)
  377. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  378. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  379. if err != nil {
  380. serverErrCh <- err
  381. return
  382. }
  383. defer func() { _ = conn.CloseNow() }()
  384. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  385. _, firstMessage, err := conn.Read(readCtx)
  386. cancelRead()
  387. if err != nil {
  388. serverErrCh <- err
  389. return
  390. }
  391. rec := httptest.NewRecorder()
  392. ginCtx, _ := gin.CreateTestContext(rec)
  393. ginCtx.Request = r.Clone(r.Context())
  394. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  395. }))
  396. defer wsServer.Close()
  397. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  398. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  399. cancelDial()
  400. require.NoError(t, err)
  401. defer func() { _ = clientConn.CloseNow() }()
  402. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  403. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false}`))
  404. cancelWrite()
  405. require.NoError(t, err)
  406. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  407. _, event, err := clientConn.Read(readCtx)
  408. cancelRead()
  409. require.NoError(t, err)
  410. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  411. closeReadCtx, cancelCloseRead := context.WithTimeout(context.Background(), 3*time.Second)
  412. _, _, err = clientConn.Read(closeReadCtx)
  413. cancelCloseRead()
  414. var clientClose coderws.CloseError
  415. require.ErrorAs(t, err, &clientClose)
  416. require.Equal(t, coderws.StatusNormalClosure, clientClose.Code)
  417. require.Equal(t, "websocket idle timeout", clientClose.Reason)
  418. select {
  419. case proxyErr := <-serverErrCh:
  420. var closeErr *OpenAIWSClientCloseError
  421. require.ErrorAs(t, proxyErr, &closeErr)
  422. require.Equal(t, coderws.StatusNormalClosure, closeErr.StatusCode())
  423. require.Equal(t, "websocket idle timeout", closeErr.Reason())
  424. case <-time.After(4 * time.Second):
  425. t.Fatal("timed out waiting for idle ingress session to close")
  426. }
  427. ap, ok := pool.getAccountPool(account.ID)
  428. require.True(t, ok)
  429. ap.mu.Lock()
  430. require.Empty(t, ap.pinnedConns, "idle close must unpin a store=false session")
  431. for _, conn := range ap.conns {
  432. require.False(t, conn.isLeased(), "idle close must release the upstream lease")
  433. }
  434. ap.mu.Unlock()
  435. }
  436. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_FollowupCreateCanOmitModel(t *testing.T) {
  437. gin.SetMode(gin.TestMode)
  438. cfg := &config.Config{}
  439. cfg.Security.URLAllowlist.Enabled = false
  440. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  441. cfg.Gateway.OpenAIWS.Enabled = true
  442. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  443. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  444. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  445. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  446. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  447. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  448. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  449. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  450. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  451. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  452. captureConn := &openAIWSCaptureConn{
  453. events: [][]byte{
  454. []byte(`{"type":"response.completed","response":{"id":"resp_omit_model_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  455. []byte(`{"type":"response.completed","response":{"id":"resp_omit_model_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  456. },
  457. }
  458. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  459. pool := newOpenAIWSConnPool(cfg)
  460. pool.setClientDialerForTest(captureDialer)
  461. svc := &OpenAIGatewayService{
  462. cfg: cfg,
  463. httpUpstream: &httpUpstreamRecorder{},
  464. cache: &stubGatewayCache{},
  465. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  466. toolCorrector: NewCodexToolCorrector(),
  467. openaiWSPool: pool,
  468. }
  469. account := &Account{
  470. ID: 115,
  471. Name: "openai-ingress-omit-model",
  472. Platform: PlatformOpenAI,
  473. Type: AccountTypeAPIKey,
  474. Status: StatusActive,
  475. Schedulable: true,
  476. Concurrency: 1,
  477. Credentials: map[string]any{
  478. "api_key": "sk-test",
  479. "model_mapping": map[string]any{
  480. "client-model": "gpt-5.1",
  481. },
  482. },
  483. Extra: map[string]any{
  484. "responses_websockets_v2_enabled": true,
  485. },
  486. }
  487. serverErrCh := make(chan error, 1)
  488. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  489. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  490. CompressionMode: coderws.CompressionContextTakeover,
  491. })
  492. if err != nil {
  493. serverErrCh <- err
  494. return
  495. }
  496. defer func() {
  497. _ = conn.CloseNow()
  498. }()
  499. rec := httptest.NewRecorder()
  500. ginCtx, _ := gin.CreateTestContext(rec)
  501. req := r.Clone(r.Context())
  502. req.Header = req.Header.Clone()
  503. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  504. ginCtx.Request = req
  505. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  506. msgType, firstMessage, readErr := conn.Read(readCtx)
  507. cancel()
  508. if readErr != nil {
  509. serverErrCh <- readErr
  510. return
  511. }
  512. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  513. serverErrCh <- errors.New("unsupported websocket client message type")
  514. return
  515. }
  516. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  517. }))
  518. defer wsServer.Close()
  519. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  520. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  521. cancelDial()
  522. require.NoError(t, err)
  523. defer func() {
  524. _ = clientConn.CloseNow()
  525. }()
  526. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  527. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"client-model","stream":false}`))
  528. cancelWrite()
  529. require.NoError(t, err)
  530. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  531. _, firstEvent, readErr := clientConn.Read(readCtx)
  532. cancelRead()
  533. require.NoError(t, readErr)
  534. require.Equal(t, "resp_omit_model_1", gjson.GetBytes(firstEvent, "response.id").String())
  535. writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
  536. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","stream":false,"previous_response_id":"resp_omit_model_1"}`))
  537. cancelWrite()
  538. require.NoError(t, err)
  539. readCtx, cancelRead = context.WithTimeout(context.Background(), 3*time.Second)
  540. _, secondEvent, readErr := clientConn.Read(readCtx)
  541. cancelRead()
  542. require.NoError(t, readErr)
  543. require.Equal(t, "resp_omit_model_2", gjson.GetBytes(secondEvent, "response.id").String())
  544. _ = clientConn.Close(coderws.StatusNormalClosure, "done")
  545. select {
  546. case serverErr := <-serverErrCh:
  547. require.NoError(t, serverErr)
  548. case <-time.After(5 * time.Second):
  549. t.Fatal("等待 ingress websocket 结束超时")
  550. }
  551. require.Len(t, captureConn.writes, 2)
  552. require.Equal(t, "gpt-5.1", gjson.Get(requestToJSONString(captureConn.writes[0]), "model").String())
  553. require.Equal(t, "gpt-5.1", gjson.Get(requestToJSONString(captureConn.writes[1]), "model").String())
  554. require.Equal(t, "resp_omit_model_1", gjson.Get(requestToJSONString(captureConn.writes[1]), "previous_response_id").String())
  555. }
  556. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_CodexImageBridgeRespectsResponsesLite(t *testing.T) {
  557. gin.SetMode(gin.TestMode)
  558. cfg := &config.Config{}
  559. cfg.Security.URLAllowlist.Enabled = false
  560. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  561. cfg.Gateway.OpenAIWS.Enabled = true
  562. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  563. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  564. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  565. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  566. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  567. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  568. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  569. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  570. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  571. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  572. captureConn := &openAIWSCaptureConn{
  573. events: [][]byte{
  574. []byte(`{"type":"response.completed","response":{"id":"resp_codex_image_bridge","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}}`),
  575. []byte(`{"type":"response.completed","response":{"id":"resp_codex_image_lite","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}}`),
  576. []byte(`{"type":"response.completed","response":{"id":"resp_codex_image_function","model":"gpt-5.5","usage":{"input_tokens":1,"output_tokens":1}}}`),
  577. },
  578. }
  579. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  580. pool := newOpenAIWSConnPool(cfg)
  581. pool.setClientDialerForTest(captureDialer)
  582. svc := &OpenAIGatewayService{
  583. cfg: cfg,
  584. httpUpstream: &httpUpstreamRecorder{},
  585. cache: &stubGatewayCache{},
  586. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  587. toolCorrector: NewCodexToolCorrector(),
  588. openaiWSPool: pool,
  589. }
  590. groupID := int64(3)
  591. apiKey := &APIKey{
  592. ID: 1,
  593. UserID: 1,
  594. GroupID: &groupID,
  595. Group: &Group{
  596. ID: groupID,
  597. AllowImageGeneration: true,
  598. },
  599. }
  600. account := &Account{
  601. ID: 31,
  602. Name: "openai-codex-image-ws",
  603. Platform: PlatformOpenAI,
  604. Type: AccountTypeOAuth,
  605. Status: StatusActive,
  606. Schedulable: true,
  607. Concurrency: 1,
  608. Credentials: map[string]any{
  609. "access_token": "test-token",
  610. },
  611. Extra: map[string]any{
  612. "openai_oauth_responses_websockets_v2_enabled": true,
  613. "codex_image_generation_bridge": true,
  614. },
  615. }
  616. serverErrCh := make(chan error, 1)
  617. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  618. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  619. CompressionMode: coderws.CompressionContextTakeover,
  620. })
  621. if err != nil {
  622. serverErrCh <- err
  623. return
  624. }
  625. defer func() {
  626. _ = conn.CloseNow()
  627. }()
  628. rec := httptest.NewRecorder()
  629. ginCtx, _ := gin.CreateTestContext(rec)
  630. req := r.Clone(r.Context())
  631. req.Header = req.Header.Clone()
  632. req.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
  633. ginCtx.Request = req
  634. ginCtx.Set("api_key", apiKey)
  635. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  636. msgType, firstMessage, readErr := conn.Read(readCtx)
  637. cancel()
  638. if readErr != nil {
  639. serverErrCh <- readErr
  640. return
  641. }
  642. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  643. serverErrCh <- errors.New("unsupported websocket client message type")
  644. return
  645. }
  646. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "test-token", firstMessage, nil)
  647. }))
  648. defer wsServer.Close()
  649. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  650. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  651. cancelDial()
  652. require.NoError(t, err)
  653. defer func() {
  654. _ = clientConn.CloseNow()
  655. }()
  656. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  657. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.5","stream":false,"input":"draw a cat","parallel_tool_calls":true,"sequence":900719925474099312345}`))
  658. cancelWrite()
  659. require.NoError(t, err)
  660. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  661. msgType, message, err := clientConn.Read(readCtx)
  662. cancelRead()
  663. require.NoError(t, err)
  664. require.Equal(t, coderws.MessageText, msgType)
  665. require.Equal(t, "resp_codex_image_bridge", gjson.GetBytes(message, "response.id").String())
  666. writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
  667. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{
  668. "type":"response.create",
  669. "model":"gpt-5.5",
  670. "stream":false,
  671. "previous_response_id":"resp_codex_image_bridge",
  672. "reasoning":{"effort":"high"},
  673. "parallel_tool_calls":true,
  674. "client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"},
  675. "tools":[{"type":"namespace","name":"collaboration","tools":[{"type":"function","name":"spawn_agent"}]}],
  676. "input":[
  677. {"type":"additional_tools","role":"developer","tools":[{"type":"custom","name":"exec","description":"Execute code-mode tools, including image_gen.imagegen."}]},
  678. {"type":"message","role":"user","content":[{"type":"input_text","text":"draw a cat"}]}
  679. ],
  680. "tool_choice":{"type":"namespace","name":"collaboration"}
  681. }`))
  682. cancelWrite()
  683. require.NoError(t, err)
  684. readCtx, cancelRead = context.WithTimeout(context.Background(), 3*time.Second)
  685. msgType, message, err = clientConn.Read(readCtx)
  686. cancelRead()
  687. require.NoError(t, err)
  688. require.Equal(t, coderws.MessageText, msgType)
  689. require.Equal(t, "resp_codex_image_lite", gjson.GetBytes(message, "response.id").String())
  690. writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
  691. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{
  692. "type":"response.create",
  693. "model":"gpt-5.5",
  694. "stream":false,
  695. "previous_response_id":"resp_codex_image_lite",
  696. "input":"draw a cat",
  697. "tools":[{"type":"function","name":"image_gen.imagegen","parameters":{"type":"object"}}]
  698. }`))
  699. cancelWrite()
  700. require.NoError(t, err)
  701. readCtx, cancelRead = context.WithTimeout(context.Background(), 3*time.Second)
  702. msgType, message, err = clientConn.Read(readCtx)
  703. cancelRead()
  704. require.NoError(t, err)
  705. require.Equal(t, coderws.MessageText, msgType)
  706. require.Equal(t, "resp_codex_image_function", gjson.GetBytes(message, "response.id").String())
  707. writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
  708. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{
  709. "type":"response.create",
  710. "model":"gpt-5.5",
  711. "parallel_tool_calls":"false",
  712. "client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"},
  713. "tools":[{"type":"function","name":"shell"}]
  714. }`))
  715. cancelWrite()
  716. require.NoError(t, err)
  717. select {
  718. case serverErr := <-serverErrCh:
  719. var closeErr *OpenAIWSClientCloseError
  720. require.ErrorAs(t, serverErr, &closeErr)
  721. require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode())
  722. require.Contains(t, closeErr.Reason(), "parallel_tool_calls to be a boolean")
  723. case <-time.After(5 * time.Second):
  724. t.Fatal("等待 ingress websocket 结束超时")
  725. }
  726. require.Len(t, captureConn.writes, 3)
  727. nonLitePayload := requestToJSONString(captureConn.writes[0])
  728. require.True(t, gjson.Get(nonLitePayload, `tools.#(type=="image_generation")`).Exists())
  729. require.Equal(t, "png", gjson.Get(nonLitePayload, `tools.#(type=="image_generation").output_format`).String())
  730. require.Equal(t, "auto", gjson.Get(nonLitePayload, "tool_choice").String())
  731. require.Contains(t, gjson.Get(nonLitePayload, "instructions").String(), "image_generation")
  732. require.False(t, gjson.Get(nonLitePayload, "reasoning.context").Exists())
  733. require.True(t, gjson.Get(nonLitePayload, "parallel_tool_calls").Bool())
  734. require.Equal(t, "900719925474099312345", gjson.Get(nonLitePayload, "sequence").Raw)
  735. litePayload := requestToJSONString(captureConn.writes[1])
  736. require.False(t, gjson.Get(litePayload, `tools.#(type=="image_generation")`).Exists())
  737. require.NotContains(t, gjson.Get(litePayload, "instructions").String(), "image_generation")
  738. require.Equal(t, "exec", gjson.Get(litePayload, `input.#(type=="additional_tools").tools.0.name`).String())
  739. require.Contains(t, gjson.Get(litePayload, `input.#(type=="additional_tools").tools.0.description`).String(), "image_gen.imagegen")
  740. require.False(t, gjson.Get(litePayload, `tools.#(type=="namespace")`).Exists())
  741. require.Equal(t, "collaboration", gjson.Get(litePayload, `input.#(type=="additional_tools").tools.1.name`).String())
  742. require.Equal(t, "namespace", gjson.Get(litePayload, "tool_choice.type").String())
  743. require.Equal(t, "collaboration", gjson.Get(litePayload, "tool_choice.name").String())
  744. require.Equal(t, "high", gjson.Get(litePayload, "reasoning.effort").String())
  745. require.Equal(t, "all_turns", gjson.Get(litePayload, "reasoning.context").String())
  746. require.True(t, gjson.Get(litePayload, "parallel_tool_calls").Exists())
  747. require.False(t, gjson.Get(litePayload, "parallel_tool_calls").Bool())
  748. functionPayload := requestToJSONString(captureConn.writes[2])
  749. require.True(t, gjson.Get(functionPayload, `tools.#(name=="image_gen.imagegen")`).Exists())
  750. require.False(t, gjson.Get(functionPayload, `tools.#(type=="image_generation")`).Exists())
  751. require.False(t, gjson.Get(functionPayload, "tool_choice").Exists())
  752. require.NotContains(t, gjson.Get(functionPayload, "instructions").String(), codexImageGenerationBridgeMarker)
  753. }
  754. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_DedicatedModeDoesNotReuseConnAcrossSessions(t *testing.T) {
  755. gin.SetMode(gin.TestMode)
  756. cfg := &config.Config{}
  757. cfg.Security.URLAllowlist.Enabled = false
  758. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  759. cfg.Gateway.OpenAIWS.Enabled = true
  760. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  761. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  762. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  763. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  764. cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeShared
  765. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  766. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  767. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  768. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  769. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  770. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  771. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  772. upstreamConn1 := &openAIWSCaptureConn{
  773. events: [][]byte{
  774. []byte(`{"type":"response.completed","response":{"id":"resp_dedicated_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  775. },
  776. }
  777. upstreamConn2 := &openAIWSCaptureConn{
  778. events: [][]byte{
  779. []byte(`{"type":"response.completed","response":{"id":"resp_dedicated_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  780. },
  781. }
  782. dialer := &openAIWSQueueDialer{
  783. conns: []openAIWSClientConn{upstreamConn1, upstreamConn2},
  784. }
  785. pool := newOpenAIWSConnPool(cfg)
  786. pool.setClientDialerForTest(dialer)
  787. svc := &OpenAIGatewayService{
  788. cfg: cfg,
  789. httpUpstream: &httpUpstreamRecorder{},
  790. cache: &stubGatewayCache{},
  791. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  792. toolCorrector: NewCodexToolCorrector(),
  793. openaiWSPool: pool,
  794. }
  795. account := &Account{
  796. ID: 441,
  797. Name: "openai-ingress-dedicated",
  798. Platform: PlatformOpenAI,
  799. Type: AccountTypeAPIKey,
  800. Status: StatusActive,
  801. Schedulable: true,
  802. Concurrency: 1,
  803. Credentials: map[string]any{
  804. "api_key": "sk-test",
  805. },
  806. Extra: map[string]any{
  807. "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModeDedicated,
  808. },
  809. }
  810. serverErrCh := make(chan error, 2)
  811. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  812. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  813. CompressionMode: coderws.CompressionContextTakeover,
  814. })
  815. if err != nil {
  816. serverErrCh <- err
  817. return
  818. }
  819. defer func() {
  820. _ = conn.CloseNow()
  821. }()
  822. rec := httptest.NewRecorder()
  823. ginCtx, _ := gin.CreateTestContext(rec)
  824. req := r.Clone(r.Context())
  825. req.Header = req.Header.Clone()
  826. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  827. ginCtx.Request = req
  828. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  829. msgType, firstMessage, readErr := conn.Read(readCtx)
  830. cancel()
  831. if readErr != nil {
  832. serverErrCh <- readErr
  833. return
  834. }
  835. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  836. serverErrCh <- errors.New("unsupported websocket client message type")
  837. return
  838. }
  839. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  840. }))
  841. defer wsServer.Close()
  842. runSingleTurnSession := func(expectedResponseID string) {
  843. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  844. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  845. cancelDial()
  846. require.NoError(t, err)
  847. defer func() {
  848. _ = clientConn.CloseNow()
  849. }()
  850. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  851. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`))
  852. cancelWrite()
  853. require.NoError(t, err)
  854. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  855. msgType, event, readErr := clientConn.Read(readCtx)
  856. cancelRead()
  857. require.NoError(t, readErr)
  858. require.Equal(t, coderws.MessageText, msgType)
  859. require.Equal(t, expectedResponseID, gjson.GetBytes(event, "response.id").String())
  860. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  861. select {
  862. case serverErr := <-serverErrCh:
  863. require.NoError(t, serverErr)
  864. case <-time.After(5 * time.Second):
  865. t.Fatal("等待 ingress websocket 结束超时")
  866. }
  867. }
  868. runSingleTurnSession("resp_dedicated_1")
  869. runSingleTurnSession("resp_dedicated_2")
  870. require.Equal(t, 2, dialer.DialCount(), "dedicated 模式下跨客户端会话不应复用上游连接")
  871. }
  872. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughModeRelaysByCaddyAdapter(t *testing.T) {
  873. gin.SetMode(gin.TestMode)
  874. cfg := &config.Config{}
  875. cfg.Security.URLAllowlist.Enabled = false
  876. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  877. cfg.Gateway.OpenAIWS.Enabled = true
  878. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  879. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  880. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  881. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  882. cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool
  883. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  884. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  885. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  886. upstreamConn := &openAIWSCaptureConn{
  887. events: [][]byte{
  888. []byte(`{"type":"response.completed","response":{"id":"resp_passthrough_turn_1","model":"gpt-5.1","output":[{"id":"ig_passthrough_1","type":"image_generation_call","status":"generating","result":"final-image"}],"usage":{"input_tokens":2,"output_tokens":3}}}`),
  889. },
  890. }
  891. captureDialer := &openAIWSCaptureDialer{conn: upstreamConn}
  892. svc := &OpenAIGatewayService{
  893. cfg: cfg,
  894. httpUpstream: &httpUpstreamRecorder{},
  895. cache: &stubGatewayCache{},
  896. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  897. toolCorrector: NewCodexToolCorrector(),
  898. openaiWSPassthroughDialer: captureDialer,
  899. }
  900. account := &Account{
  901. ID: 452,
  902. Name: "openai-ingress-passthrough",
  903. Platform: PlatformOpenAI,
  904. Type: AccountTypeAPIKey,
  905. Status: StatusActive,
  906. Schedulable: true,
  907. Concurrency: 1,
  908. Credentials: map[string]any{
  909. "api_key": "sk-test",
  910. },
  911. Extra: map[string]any{
  912. "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough,
  913. },
  914. }
  915. serverErrCh := make(chan error, 1)
  916. resultCh := make(chan *OpenAIForwardResult, 1)
  917. hooks := &OpenAIWSIngressHooks{
  918. AfterTurn: func(_ int, result *OpenAIForwardResult, turnErr error) {
  919. if turnErr == nil && result != nil {
  920. resultCh <- result
  921. }
  922. },
  923. }
  924. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  925. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  926. CompressionMode: coderws.CompressionContextTakeover,
  927. })
  928. if err != nil {
  929. serverErrCh <- err
  930. return
  931. }
  932. defer func() {
  933. _ = conn.CloseNow()
  934. }()
  935. rec := httptest.NewRecorder()
  936. ginCtx, _ := gin.CreateTestContext(rec)
  937. req := r.Clone(r.Context())
  938. req.Header = req.Header.Clone()
  939. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  940. ginCtx.Request = req
  941. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  942. msgType, firstMessage, readErr := conn.Read(readCtx)
  943. cancel()
  944. if readErr != nil {
  945. serverErrCh <- readErr
  946. return
  947. }
  948. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  949. serverErrCh <- errors.New("unsupported websocket client message type")
  950. return
  951. }
  952. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, hooks)
  953. }))
  954. defer wsServer.Close()
  955. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  956. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  957. cancelDial()
  958. require.NoError(t, err)
  959. defer func() {
  960. _ = clientConn.CloseNow()
  961. }()
  962. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  963. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false,"service_tier":"fast","reasoning":{"effort":"HIGH"}}`))
  964. cancelWrite()
  965. require.NoError(t, err)
  966. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  967. _, event, readErr := clientConn.Read(readCtx)
  968. cancelRead()
  969. require.NoError(t, readErr)
  970. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  971. require.Equal(t, "resp_passthrough_turn_1", gjson.GetBytes(event, "response.id").String())
  972. require.Equal(t, "completed", gjson.GetBytes(event, "response.output.0.status").String())
  973. _ = clientConn.Close(coderws.StatusNormalClosure, "done")
  974. select {
  975. case serverErr := <-serverErrCh:
  976. // After normal client close, the server goroutine may receive the close frame
  977. // as an error — this is expected behavior, not a test failure.
  978. if serverErr != nil {
  979. require.Contains(t, serverErr.Error(), "StatusNormalClosure",
  980. "server error should only be a normal close frame, got: %v", serverErr)
  981. }
  982. case <-time.After(5 * time.Second):
  983. t.Fatal("等待 passthrough websocket 结束超时")
  984. }
  985. select {
  986. case result := <-resultCh:
  987. require.Equal(t, "resp_passthrough_turn_1", result.RequestID)
  988. require.True(t, result.OpenAIWSMode)
  989. require.Equal(t, 2, result.Usage.InputTokens)
  990. require.Equal(t, 3, result.Usage.OutputTokens)
  991. require.NotNil(t, result.ServiceTier)
  992. require.Equal(t, "priority", *result.ServiceTier)
  993. require.NotNil(t, result.ReasoningEffort)
  994. require.Equal(t, "high", *result.ReasoningEffort)
  995. case <-time.After(2 * time.Second):
  996. t.Fatal("未收到 passthrough turn 结果回调")
  997. }
  998. require.Equal(t, 1, captureDialer.DialCount(), "passthrough 模式应直接建立上游 websocket")
  999. require.Len(t, upstreamConn.writes, 1, "passthrough 模式应透传首条 response.create")
  1000. }
  1001. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughBridgeBoundaries(t *testing.T) {
  1002. tests := []struct {
  1003. name string
  1004. payload string
  1005. threshold int64
  1006. wantRelayReject bool
  1007. }{
  1008. {
  1009. name: "small response create",
  1010. payload: `{"type":"response.create","model":"gpt-5.1"}`,
  1011. threshold: 1024,
  1012. },
  1013. {
  1014. name: "continuation response create",
  1015. payload: `{"type":"response.create","previous_response_id":"resp_previous","model":"gpt-5.1"}`,
  1016. threshold: 1,
  1017. },
  1018. {
  1019. name: "other event type",
  1020. payload: `{"type":"session.update","padding":"xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"}`,
  1021. threshold: 1,
  1022. },
  1023. {
  1024. name: "malformed data",
  1025. payload: `{"type":"response.create","padding":"xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"`,
  1026. threshold: 1,
  1027. wantRelayReject: true,
  1028. },
  1029. {
  1030. name: "duplicate type",
  1031. payload: `{"type":"response.create","type":"response.create","model":"gpt-5.1"}`,
  1032. threshold: 1,
  1033. },
  1034. {
  1035. name: "duplicate previous response id",
  1036. payload: `{"type":"response.create","previous_response_id":null,"previous_response_id":null,"model":"gpt-5.1"}`,
  1037. threshold: 1,
  1038. },
  1039. }
  1040. for _, tt := range tests {
  1041. t.Run(tt.name, func(t *testing.T) {
  1042. gin.SetMode(gin.TestMode)
  1043. cfg := &config.Config{}
  1044. cfg.Security.URLAllowlist.Enabled = false
  1045. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1046. cfg.Gateway.OpenAIWS.Enabled = true
  1047. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1048. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1049. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  1050. cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool
  1051. cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true
  1052. cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = tt.threshold
  1053. upstreamConn := &openAIWSCaptureConn{events: [][]byte{
  1054. []byte(`{"type":"response.completed","response":{"id":"resp_duplicate_keys","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1055. }}
  1056. dialer := &openAIWSCaptureDialer{conn: upstreamConn}
  1057. httpUpstream := &httpUpstreamRecorder{}
  1058. svc := &OpenAIGatewayService{
  1059. cfg: cfg,
  1060. httpUpstream: httpUpstream,
  1061. cache: &stubGatewayCache{},
  1062. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1063. toolCorrector: NewCodexToolCorrector(),
  1064. openaiWSPassthroughDialer: dialer,
  1065. }
  1066. account := &Account{
  1067. ID: 453,
  1068. Platform: PlatformOpenAI,
  1069. Type: AccountTypeAPIKey,
  1070. Status: StatusActive,
  1071. Schedulable: true,
  1072. Concurrency: 1,
  1073. Credentials: map[string]any{"api_key": "sk-test"},
  1074. Extra: map[string]any{
  1075. "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough,
  1076. },
  1077. }
  1078. errCh := make(chan error, 1)
  1079. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1080. conn, err := coderws.Accept(w, r, nil)
  1081. if err != nil {
  1082. errCh <- err
  1083. return
  1084. }
  1085. defer func() { _ = conn.CloseNow() }()
  1086. _, firstMessage, err := conn.Read(r.Context())
  1087. if err != nil {
  1088. errCh <- err
  1089. return
  1090. }
  1091. rec := httptest.NewRecorder()
  1092. ginCtx, _ := gin.CreateTestContext(rec)
  1093. ginCtx.Request = r
  1094. errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  1095. }))
  1096. defer wsServer.Close()
  1097. clientConn, _, err := coderws.Dial(context.Background(), "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1098. require.NoError(t, err)
  1099. defer func() { _ = clientConn.CloseNow() }()
  1100. require.NoError(t, clientConn.Write(context.Background(), coderws.MessageText, []byte(tt.payload)))
  1101. _, event, err := clientConn.Read(context.Background())
  1102. if tt.wantRelayReject {
  1103. require.Error(t, err)
  1104. } else {
  1105. require.NoError(t, err)
  1106. require.Equal(t, "resp_duplicate_keys", gjson.GetBytes(event, "response.id").String())
  1107. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  1108. }
  1109. select {
  1110. case proxyErr := <-errCh:
  1111. if tt.wantRelayReject {
  1112. require.Error(t, proxyErr)
  1113. } else if proxyErr != nil {
  1114. require.Contains(t, proxyErr.Error(), "StatusNormalClosure")
  1115. }
  1116. case <-time.After(5 * time.Second):
  1117. t.Fatal("等待 boundary passthrough websocket 结束超时")
  1118. }
  1119. require.Empty(t, httpUpstream.requests, "boundary frame must not use the HTTP bridge")
  1120. if tt.wantRelayReject {
  1121. require.Zero(t, dialer.DialCount())
  1122. require.Empty(t, upstreamConn.writes)
  1123. } else {
  1124. require.Equal(t, 1, dialer.DialCount())
  1125. require.Len(t, upstreamConn.writes, 1)
  1126. }
  1127. })
  1128. }
  1129. }
  1130. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughHeadersUsePromptCacheAndTurnState(t *testing.T) {
  1131. gin.SetMode(gin.TestMode)
  1132. cfg := &config.Config{}
  1133. cfg.Security.URLAllowlist.Enabled = false
  1134. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1135. cfg.Gateway.OpenAIWS.Enabled = true
  1136. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1137. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1138. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1139. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  1140. cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool
  1141. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  1142. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  1143. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1144. upstreamConn := &openAIWSCaptureConn{
  1145. events: [][]byte{
  1146. []byte(`{"type":"response.completed","response":{"id":"resp_passthrough_headers","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1147. },
  1148. }
  1149. captureDialer := &openAIWSCaptureDialer{conn: upstreamConn}
  1150. svc := &OpenAIGatewayService{
  1151. cfg: cfg,
  1152. httpUpstream: &httpUpstreamRecorder{},
  1153. cache: &stubGatewayCache{},
  1154. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1155. toolCorrector: NewCodexToolCorrector(),
  1156. openaiWSPassthroughDialer: captureDialer,
  1157. }
  1158. account := &Account{
  1159. ID: 453,
  1160. Name: "openai-ingress-passthrough-headers",
  1161. Platform: PlatformOpenAI,
  1162. Type: AccountTypeOAuth,
  1163. Status: StatusActive,
  1164. Schedulable: true,
  1165. Concurrency: 1,
  1166. Credentials: map[string]any{
  1167. "access_token": "oauth-token",
  1168. },
  1169. Extra: map[string]any{
  1170. "openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough,
  1171. },
  1172. }
  1173. serverErrCh := make(chan error, 1)
  1174. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1175. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  1176. CompressionMode: coderws.CompressionContextTakeover,
  1177. })
  1178. if err != nil {
  1179. serverErrCh <- err
  1180. return
  1181. }
  1182. defer func() {
  1183. _ = conn.CloseNow()
  1184. }()
  1185. rec := httptest.NewRecorder()
  1186. ginCtx, _ := gin.CreateTestContext(rec)
  1187. req := r.Clone(r.Context())
  1188. req.Header = req.Header.Clone()
  1189. req.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
  1190. req.Header.Set(openAIWSTurnStateHeader, "turn-state-1")
  1191. req.Header.Set(openAIWSTurnMetadataHeader, "turn-meta-1")
  1192. ginCtx.Request = req
  1193. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  1194. msgType, firstMessage, readErr := conn.Read(readCtx)
  1195. cancel()
  1196. if readErr != nil {
  1197. serverErrCh <- readErr
  1198. return
  1199. }
  1200. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  1201. serverErrCh <- errors.New("unsupported websocket client message type")
  1202. return
  1203. }
  1204. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "oauth-token", firstMessage, nil)
  1205. }))
  1206. defer wsServer.Close()
  1207. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1208. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1209. cancelDial()
  1210. require.NoError(t, err)
  1211. defer func() {
  1212. _ = clientConn.CloseNow()
  1213. }()
  1214. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  1215. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{
  1216. "type":"response.create",
  1217. "model":"gpt-5.1",
  1218. "stream":false,
  1219. "prompt_cache_key":"pcache_passthrough",
  1220. "reasoning":{"effort":"medium","context":"current_turn"},
  1221. "parallel_tool_calls":true,
  1222. "client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"},
  1223. "tools":[{"type":"namespace","name":"collaboration","tools":[{"type":"function","name":"spawn_agent"}]}],
  1224. "input":[{"type":"message","role":"user","content":"hello"}],
  1225. "tool_choice":{"type":"namespace","name":"collaboration"}
  1226. }`))
  1227. cancelWrite()
  1228. require.NoError(t, err)
  1229. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  1230. _, event, readErr := clientConn.Read(readCtx)
  1231. cancelRead()
  1232. require.NoError(t, readErr)
  1233. require.Equal(t, "resp_passthrough_headers", gjson.GetBytes(event, "response.id").String())
  1234. _ = clientConn.Close(coderws.StatusNormalClosure, "done")
  1235. select {
  1236. case serverErr := <-serverErrCh:
  1237. if serverErr != nil {
  1238. require.Contains(t, serverErr.Error(), "StatusNormalClosure")
  1239. }
  1240. case <-time.After(5 * time.Second):
  1241. t.Fatal("等待 passthrough websocket 结束超时")
  1242. }
  1243. require.Equal(t, isolateOpenAIUpstreamSessionID(0, account, "pcache_passthrough"), captureDialer.lastHeaders.Get("session_id"))
  1244. require.Equal(t, "turn-state-1", captureDialer.lastHeaders.Get(openAIWSTurnStateHeader))
  1245. require.Equal(t, "turn-meta-1", captureDialer.lastHeaders.Get(openAIWSTurnMetadataHeader))
  1246. require.Len(t, upstreamConn.writes, 1)
  1247. forwarded := requestToJSONString(upstreamConn.writes[0])
  1248. require.False(t, gjson.Get(forwarded, `tools.#(type=="namespace")`).Exists())
  1249. require.Equal(t, "collaboration", gjson.Get(forwarded, `input.#(type=="additional_tools").tools.0.name`).String())
  1250. require.Equal(t, "namespace", gjson.Get(forwarded, "tool_choice.type").String())
  1251. require.Equal(t, "collaboration", gjson.Get(forwarded, "tool_choice.name").String())
  1252. require.Equal(t, "medium", gjson.Get(forwarded, "reasoning.effort").String())
  1253. require.Equal(t, "all_turns", gjson.Get(forwarded, "reasoning.context").String())
  1254. require.True(t, gjson.Get(forwarded, "parallel_tool_calls").Exists())
  1255. require.False(t, gjson.Get(forwarded, "parallel_tool_calls").Bool())
  1256. }
  1257. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_HTTPBridgeModeRelaysHTTPStream(t *testing.T) {
  1258. gin.SetMode(gin.TestMode)
  1259. cfg := &config.Config{}
  1260. cfg.Security.URLAllowlist.Enabled = false
  1261. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1262. cfg.Gateway.OpenAIWS.Enabled = true
  1263. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1264. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1265. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1266. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  1267. cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool
  1268. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  1269. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1270. upstream := &httpUpstreamRecorder{
  1271. resp: &http.Response{
  1272. StatusCode: http.StatusOK,
  1273. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_bridge_1"}},
  1274. Body: io.NopCloser(strings.NewReader(
  1275. "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n" +
  1276. "data: {\"type\":\"response.done\",\"response\":{\"id\":\"resp_http_bridge_1\",\"output\":[{\"id\":\"ig_bridge_1\",\"type\":\"image_generation_call\",\"status\":\"in_progress\",\"result\":\"final-image\"}],\"usage\":{\"input_tokens\":2,\"output_tokens\":1,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n" +
  1277. "data: [DONE]\n\n",
  1278. )),
  1279. },
  1280. }
  1281. svc := &OpenAIGatewayService{
  1282. cfg: cfg,
  1283. httpUpstream: upstream,
  1284. cache: &stubGatewayCache{},
  1285. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1286. toolCorrector: NewCodexToolCorrector(),
  1287. }
  1288. account := &Account{
  1289. ID: 552,
  1290. Name: "openai-ingress-http-bridge",
  1291. Platform: PlatformOpenAI,
  1292. Type: AccountTypeAPIKey,
  1293. Status: StatusActive,
  1294. Schedulable: true,
  1295. Concurrency: 1,
  1296. Credentials: map[string]any{
  1297. "api_key": "sk-test",
  1298. },
  1299. Extra: map[string]any{
  1300. "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModeHTTPBridge,
  1301. },
  1302. }
  1303. serverErrCh := make(chan error, 1)
  1304. resultCh := make(chan *OpenAIForwardResult, 1)
  1305. hooks := &OpenAIWSIngressHooks{
  1306. AfterTurn: func(_ int, result *OpenAIForwardResult, turnErr error) {
  1307. if turnErr == nil && result != nil {
  1308. resultCh <- result
  1309. }
  1310. },
  1311. }
  1312. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1313. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  1314. CompressionMode: coderws.CompressionContextTakeover,
  1315. })
  1316. if err != nil {
  1317. serverErrCh <- err
  1318. return
  1319. }
  1320. defer func() {
  1321. _ = conn.CloseNow()
  1322. }()
  1323. rec := httptest.NewRecorder()
  1324. ginCtx, _ := gin.CreateTestContext(rec)
  1325. req := r.Clone(r.Context())
  1326. req.Header = req.Header.Clone()
  1327. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  1328. ginCtx.Request = req
  1329. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  1330. msgType, firstMessage, readErr := conn.Read(readCtx)
  1331. cancel()
  1332. if readErr != nil {
  1333. serverErrCh <- readErr
  1334. return
  1335. }
  1336. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  1337. serverErrCh <- errors.New("unsupported websocket client message type")
  1338. return
  1339. }
  1340. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, hooks)
  1341. }))
  1342. defer wsServer.Close()
  1343. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1344. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1345. cancelDial()
  1346. require.NoError(t, err)
  1347. defer func() {
  1348. _ = clientConn.CloseNow()
  1349. }()
  1350. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  1351. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{
  1352. "type":"response.create","model":"gpt-5.1","stream":false,
  1353. "parallel_tool_calls":true,
  1354. "client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"}
  1355. }`))
  1356. cancelWrite()
  1357. require.NoError(t, err)
  1358. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  1359. _, event1, readErr1 := clientConn.Read(readCtx)
  1360. cancelRead()
  1361. require.NoError(t, readErr1)
  1362. require.Equal(t, "response.output_text.delta", gjson.GetBytes(event1, "type").String())
  1363. require.Equal(t, "hello", gjson.GetBytes(event1, "delta").String())
  1364. readCtx2, cancelRead2 := context.WithTimeout(context.Background(), 3*time.Second)
  1365. _, event2, readErr2 := clientConn.Read(readCtx2)
  1366. cancelRead2()
  1367. require.NoError(t, readErr2)
  1368. require.Equal(t, "response.done", gjson.GetBytes(event2, "type").String())
  1369. require.Equal(t, "resp_http_bridge_1", gjson.GetBytes(event2, "response.id").String())
  1370. require.Equal(t, "completed", gjson.GetBytes(event2, "response.output.0.status").String())
  1371. _ = clientConn.Close(coderws.StatusNormalClosure, "done")
  1372. select {
  1373. case serverErr := <-serverErrCh:
  1374. require.NoError(t, serverErr)
  1375. case <-time.After(5 * time.Second):
  1376. t.Fatal("等待 http_bridge websocket 结束超时")
  1377. }
  1378. select {
  1379. case result := <-resultCh:
  1380. require.Equal(t, "resp_http_bridge_1", result.RequestID)
  1381. require.True(t, result.OpenAIWSMode)
  1382. require.Equal(t, 2, result.Usage.InputTokens)
  1383. require.Equal(t, 1, result.Usage.OutputTokens)
  1384. require.Equal(t, 1, result.Usage.CacheReadInputTokens)
  1385. require.NotNil(t, result.FirstTokenMs)
  1386. case <-time.After(2 * time.Second):
  1387. t.Fatal("未收到 http_bridge turn 结果回调")
  1388. }
  1389. require.NotNil(t, upstream.lastReq, "http_bridge 模式应调用 HTTP 上游")
  1390. require.Equal(t, "true", upstream.lastReq.Header.Get(responsesLiteHeader))
  1391. require.True(t, gjson.GetBytes(upstream.lastBody, "parallel_tool_calls").Exists())
  1392. require.False(t, gjson.GetBytes(upstream.lastBody, "parallel_tool_calls").Bool())
  1393. }
  1394. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_ModeOffReturnsPolicyViolation(t *testing.T) {
  1395. gin.SetMode(gin.TestMode)
  1396. cfg := &config.Config{}
  1397. cfg.Security.URLAllowlist.Enabled = false
  1398. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1399. cfg.Gateway.OpenAIWS.Enabled = true
  1400. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1401. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1402. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1403. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  1404. cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeShared
  1405. svc := &OpenAIGatewayService{
  1406. cfg: cfg,
  1407. httpUpstream: &httpUpstreamRecorder{},
  1408. cache: &stubGatewayCache{},
  1409. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1410. toolCorrector: NewCodexToolCorrector(),
  1411. openaiWSPool: newOpenAIWSConnPool(cfg),
  1412. }
  1413. account := &Account{
  1414. ID: 442,
  1415. Name: "openai-ingress-off",
  1416. Platform: PlatformOpenAI,
  1417. Type: AccountTypeAPIKey,
  1418. Status: StatusActive,
  1419. Schedulable: true,
  1420. Concurrency: 1,
  1421. Credentials: map[string]any{
  1422. "api_key": "sk-test",
  1423. },
  1424. Extra: map[string]any{
  1425. "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModeOff,
  1426. },
  1427. }
  1428. serverErrCh := make(chan error, 1)
  1429. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1430. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  1431. CompressionMode: coderws.CompressionContextTakeover,
  1432. })
  1433. if err != nil {
  1434. serverErrCh <- err
  1435. return
  1436. }
  1437. defer func() {
  1438. _ = conn.CloseNow()
  1439. }()
  1440. rec := httptest.NewRecorder()
  1441. ginCtx, _ := gin.CreateTestContext(rec)
  1442. req := r.Clone(r.Context())
  1443. req.Header = req.Header.Clone()
  1444. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  1445. ginCtx.Request = req
  1446. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  1447. msgType, firstMessage, readErr := conn.Read(readCtx)
  1448. cancel()
  1449. if readErr != nil {
  1450. serverErrCh <- readErr
  1451. return
  1452. }
  1453. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  1454. serverErrCh <- errors.New("unsupported websocket client message type")
  1455. return
  1456. }
  1457. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  1458. }))
  1459. defer wsServer.Close()
  1460. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1461. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1462. cancelDial()
  1463. require.NoError(t, err)
  1464. defer func() {
  1465. _ = clientConn.CloseNow()
  1466. }()
  1467. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  1468. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`))
  1469. cancelWrite()
  1470. require.NoError(t, err)
  1471. select {
  1472. case serverErr := <-serverErrCh:
  1473. var closeErr *OpenAIWSClientCloseError
  1474. require.ErrorAs(t, serverErr, &closeErr)
  1475. require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode())
  1476. require.Equal(t, "websocket mode is disabled for this account", closeErr.Reason())
  1477. case <-time.After(5 * time.Second):
  1478. t.Fatal("等待 ingress websocket 结束超时")
  1479. }
  1480. }
  1481. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPrevResponseStrictDropToFullCreate(t *testing.T) {
  1482. gin.SetMode(gin.TestMode)
  1483. cfg := &config.Config{}
  1484. cfg.Security.URLAllowlist.Enabled = false
  1485. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1486. cfg.Gateway.OpenAIWS.Enabled = true
  1487. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1488. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1489. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1490. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1491. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1492. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1493. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  1494. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  1495. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  1496. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1497. captureConn := &openAIWSCaptureConn{
  1498. events: [][]byte{
  1499. []byte(`{"type":"response.completed","response":{"id":"resp_preflight_rewrite_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1500. []byte(`{"type":"response.completed","response":{"id":"resp_preflight_rewrite_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1501. },
  1502. }
  1503. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1504. pool := newOpenAIWSConnPool(cfg)
  1505. pool.setClientDialerForTest(captureDialer)
  1506. svc := &OpenAIGatewayService{
  1507. cfg: cfg,
  1508. httpUpstream: &httpUpstreamRecorder{},
  1509. cache: &stubGatewayCache{},
  1510. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1511. toolCorrector: NewCodexToolCorrector(),
  1512. openaiWSPool: pool,
  1513. }
  1514. account := &Account{
  1515. ID: 140,
  1516. Name: "openai-ingress-prev-preflight-rewrite",
  1517. Platform: PlatformOpenAI,
  1518. Type: AccountTypeAPIKey,
  1519. Status: StatusActive,
  1520. Schedulable: true,
  1521. Concurrency: 1,
  1522. Credentials: map[string]any{
  1523. "api_key": "sk-test",
  1524. },
  1525. Extra: map[string]any{
  1526. "responses_websockets_v2_enabled": true,
  1527. },
  1528. }
  1529. serverErrCh := make(chan error, 1)
  1530. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1531. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  1532. CompressionMode: coderws.CompressionContextTakeover,
  1533. })
  1534. if err != nil {
  1535. serverErrCh <- err
  1536. return
  1537. }
  1538. defer func() {
  1539. _ = conn.CloseNow()
  1540. }()
  1541. rec := httptest.NewRecorder()
  1542. ginCtx, _ := gin.CreateTestContext(rec)
  1543. req := r.Clone(r.Context())
  1544. req.Header = req.Header.Clone()
  1545. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  1546. ginCtx.Request = req
  1547. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  1548. msgType, firstMessage, readErr := conn.Read(readCtx)
  1549. cancel()
  1550. if readErr != nil {
  1551. serverErrCh <- readErr
  1552. return
  1553. }
  1554. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  1555. serverErrCh <- errors.New("unsupported websocket client message type")
  1556. return
  1557. }
  1558. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  1559. }))
  1560. defer wsServer.Close()
  1561. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1562. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1563. cancelDial()
  1564. require.NoError(t, err)
  1565. defer func() {
  1566. _ = clientConn.CloseNow()
  1567. }()
  1568. writeMessage := func(payload string) {
  1569. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  1570. defer cancel()
  1571. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  1572. }
  1573. readMessage := func() []byte {
  1574. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  1575. defer cancel()
  1576. msgType, message, readErr := clientConn.Read(readCtx)
  1577. require.NoError(t, readErr)
  1578. require.Equal(t, coderws.MessageText, msgType)
  1579. return message
  1580. }
  1581. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  1582. firstTurn := readMessage()
  1583. require.Equal(t, "resp_preflight_rewrite_1", gjson.GetBytes(firstTurn, "response.id").String())
  1584. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"previous_response_id":"resp_stale_external","input":[{"type":"input_text","text":"world"}]}`)
  1585. secondTurn := readMessage()
  1586. require.Equal(t, "resp_preflight_rewrite_2", gjson.GetBytes(secondTurn, "response.id").String())
  1587. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  1588. select {
  1589. case serverErr := <-serverErrCh:
  1590. require.NoError(t, serverErr)
  1591. case <-time.After(5 * time.Second):
  1592. t.Fatal("等待 ingress websocket 结束超时")
  1593. }
  1594. require.Equal(t, 1, captureDialer.DialCount(), "严格增量不成立时应在同一连接内降级为 full create")
  1595. require.Len(t, captureConn.writes, 2)
  1596. secondWrite := requestToJSONString(captureConn.writes[1])
  1597. require.False(t, gjson.Get(secondWrite, "previous_response_id").Exists(), "严格增量不成立时应移除 previous_response_id,改为 full create")
  1598. require.Equal(t, 2, len(gjson.Get(secondWrite, "input").Array()), "严格降级为 full create 时应重放完整 input 上下文")
  1599. require.Equal(t, "hello", gjson.Get(secondWrite, "input.0.text").String())
  1600. require.Equal(t, "world", gjson.Get(secondWrite, "input.1.text").String())
  1601. }
  1602. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPrevResponseStrictDropBeforePreflightPingFailReconnects(t *testing.T) {
  1603. gin.SetMode(gin.TestMode)
  1604. prevPreflightPingIdle := openAIWSIngressPreflightPingIdle
  1605. openAIWSIngressPreflightPingIdle = 0
  1606. defer func() {
  1607. openAIWSIngressPreflightPingIdle = prevPreflightPingIdle
  1608. }()
  1609. cfg := &config.Config{}
  1610. cfg.Security.URLAllowlist.Enabled = false
  1611. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1612. cfg.Gateway.OpenAIWS.Enabled = true
  1613. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1614. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1615. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1616. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  1617. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1618. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  1619. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  1620. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  1621. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  1622. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1623. firstConn := &openAIWSPreflightFailConn{
  1624. events: [][]byte{
  1625. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_drop_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1626. },
  1627. }
  1628. secondConn := &openAIWSCaptureConn{
  1629. events: [][]byte{
  1630. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_drop_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1631. },
  1632. }
  1633. dialer := &openAIWSQueueDialer{
  1634. conns: []openAIWSClientConn{firstConn, secondConn},
  1635. }
  1636. pool := newOpenAIWSConnPool(cfg)
  1637. pool.setClientDialerForTest(dialer)
  1638. svc := &OpenAIGatewayService{
  1639. cfg: cfg,
  1640. httpUpstream: &httpUpstreamRecorder{},
  1641. cache: &stubGatewayCache{},
  1642. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1643. toolCorrector: NewCodexToolCorrector(),
  1644. openaiWSPool: pool,
  1645. }
  1646. account := &Account{
  1647. ID: 142,
  1648. Name: "openai-ingress-prev-strict-drop-before-ping",
  1649. Platform: PlatformOpenAI,
  1650. Type: AccountTypeAPIKey,
  1651. Status: StatusActive,
  1652. Schedulable: true,
  1653. Concurrency: 1,
  1654. Credentials: map[string]any{
  1655. "api_key": "sk-test",
  1656. },
  1657. Extra: map[string]any{
  1658. "responses_websockets_v2_enabled": true,
  1659. },
  1660. }
  1661. serverErrCh := make(chan error, 1)
  1662. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1663. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  1664. CompressionMode: coderws.CompressionContextTakeover,
  1665. })
  1666. if err != nil {
  1667. serverErrCh <- err
  1668. return
  1669. }
  1670. defer func() {
  1671. _ = conn.CloseNow()
  1672. }()
  1673. rec := httptest.NewRecorder()
  1674. ginCtx, _ := gin.CreateTestContext(rec)
  1675. req := r.Clone(r.Context())
  1676. req.Header = req.Header.Clone()
  1677. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  1678. ginCtx.Request = req
  1679. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  1680. msgType, firstMessage, readErr := conn.Read(readCtx)
  1681. cancel()
  1682. if readErr != nil {
  1683. serverErrCh <- readErr
  1684. return
  1685. }
  1686. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  1687. serverErrCh <- errors.New("unsupported websocket client message type")
  1688. return
  1689. }
  1690. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  1691. }))
  1692. defer wsServer.Close()
  1693. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1694. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1695. cancelDial()
  1696. require.NoError(t, err)
  1697. defer func() {
  1698. _ = clientConn.CloseNow()
  1699. }()
  1700. writeMessage := func(payload string) {
  1701. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  1702. defer cancel()
  1703. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  1704. }
  1705. readMessage := func() []byte {
  1706. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  1707. defer cancel()
  1708. msgType, message, readErr := clientConn.Read(readCtx)
  1709. require.NoError(t, readErr)
  1710. require.Equal(t, coderws.MessageText, msgType)
  1711. return message
  1712. }
  1713. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  1714. firstTurn := readMessage()
  1715. require.Equal(t, "resp_turn_ping_drop_1", gjson.GetBytes(firstTurn, "response.id").String())
  1716. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"previous_response_id":"resp_stale_external","input":[{"type":"input_text","text":"world"}]}`)
  1717. secondTurn := readMessage()
  1718. require.Equal(t, "resp_turn_ping_drop_2", gjson.GetBytes(secondTurn, "response.id").String())
  1719. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  1720. select {
  1721. case serverErr := <-serverErrCh:
  1722. require.NoError(t, serverErr)
  1723. case <-time.After(5 * time.Second):
  1724. t.Fatal("等待 ingress websocket 严格降级后预检换连超时")
  1725. }
  1726. require.Equal(t, 2, dialer.DialCount(), "严格降级为 full create 后,预检 ping 失败应允许换连")
  1727. require.Equal(t, 1, firstConn.WriteCount(), "首连接在预检失败后不应继续发送第二轮")
  1728. require.GreaterOrEqual(t, firstConn.PingCount(), 1, "第二轮前应执行 preflight ping")
  1729. secondConn.mu.Lock()
  1730. secondWrites := append([]map[string]any(nil), secondConn.writes...)
  1731. secondConn.mu.Unlock()
  1732. require.Len(t, secondWrites, 1)
  1733. secondWrite := requestToJSONString(secondWrites[0])
  1734. require.False(t, gjson.Get(secondWrite, "previous_response_id").Exists(), "严格降级后重试应移除 previous_response_id")
  1735. require.Equal(t, 2, len(gjson.Get(secondWrite, "input").Array()))
  1736. require.Equal(t, "hello", gjson.Get(secondWrite, "input.0.text").String())
  1737. require.Equal(t, "world", gjson.Get(secondWrite, "input.1.text").String())
  1738. }
  1739. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreEnabledSkipsStrictPrevResponseEval(t *testing.T) {
  1740. gin.SetMode(gin.TestMode)
  1741. cfg := &config.Config{}
  1742. cfg.Security.URLAllowlist.Enabled = false
  1743. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1744. cfg.Gateway.OpenAIWS.Enabled = true
  1745. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1746. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1747. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1748. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1749. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1750. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1751. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  1752. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  1753. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  1754. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1755. captureConn := &openAIWSCaptureConn{
  1756. events: [][]byte{
  1757. []byte(`{"type":"response.completed","response":{"id":"resp_store_enabled_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1758. []byte(`{"type":"response.completed","response":{"id":"resp_store_enabled_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1759. },
  1760. }
  1761. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1762. pool := newOpenAIWSConnPool(cfg)
  1763. pool.setClientDialerForTest(captureDialer)
  1764. svc := &OpenAIGatewayService{
  1765. cfg: cfg,
  1766. httpUpstream: &httpUpstreamRecorder{},
  1767. cache: &stubGatewayCache{},
  1768. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1769. toolCorrector: NewCodexToolCorrector(),
  1770. openaiWSPool: pool,
  1771. }
  1772. account := &Account{
  1773. ID: 143,
  1774. Name: "openai-ingress-store-enabled-skip-strict",
  1775. Platform: PlatformOpenAI,
  1776. Type: AccountTypeAPIKey,
  1777. Status: StatusActive,
  1778. Schedulable: true,
  1779. Concurrency: 1,
  1780. Credentials: map[string]any{
  1781. "api_key": "sk-test",
  1782. },
  1783. Extra: map[string]any{
  1784. "responses_websockets_v2_enabled": true,
  1785. },
  1786. }
  1787. serverErrCh := make(chan error, 1)
  1788. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1789. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  1790. CompressionMode: coderws.CompressionContextTakeover,
  1791. })
  1792. if err != nil {
  1793. serverErrCh <- err
  1794. return
  1795. }
  1796. defer func() {
  1797. _ = conn.CloseNow()
  1798. }()
  1799. rec := httptest.NewRecorder()
  1800. ginCtx, _ := gin.CreateTestContext(rec)
  1801. req := r.Clone(r.Context())
  1802. req.Header = req.Header.Clone()
  1803. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  1804. ginCtx.Request = req
  1805. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  1806. msgType, firstMessage, readErr := conn.Read(readCtx)
  1807. cancel()
  1808. if readErr != nil {
  1809. serverErrCh <- readErr
  1810. return
  1811. }
  1812. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  1813. serverErrCh <- errors.New("unsupported websocket client message type")
  1814. return
  1815. }
  1816. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  1817. }))
  1818. defer wsServer.Close()
  1819. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1820. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1821. cancelDial()
  1822. require.NoError(t, err)
  1823. defer func() {
  1824. _ = clientConn.CloseNow()
  1825. }()
  1826. writeMessage := func(payload string) {
  1827. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  1828. defer cancel()
  1829. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  1830. }
  1831. readMessage := func() []byte {
  1832. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  1833. defer cancel()
  1834. msgType, message, readErr := clientConn.Read(readCtx)
  1835. require.NoError(t, readErr)
  1836. require.Equal(t, coderws.MessageText, msgType)
  1837. return message
  1838. }
  1839. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":true}`)
  1840. firstTurn := readMessage()
  1841. require.Equal(t, "resp_store_enabled_1", gjson.GetBytes(firstTurn, "response.id").String())
  1842. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":true,"previous_response_id":"resp_stale_external"}`)
  1843. secondTurn := readMessage()
  1844. require.Equal(t, "resp_store_enabled_2", gjson.GetBytes(secondTurn, "response.id").String())
  1845. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  1846. select {
  1847. case serverErr := <-serverErrCh:
  1848. require.NoError(t, serverErr)
  1849. case <-time.After(5 * time.Second):
  1850. t.Fatal("等待 store=true 场景 websocket 结束超时")
  1851. }
  1852. require.Equal(t, 1, captureDialer.DialCount())
  1853. require.Len(t, captureConn.writes, 2)
  1854. require.Equal(t, "resp_stale_external", gjson.Get(requestToJSONString(captureConn.writes[1]), "previous_response_id").String(), "store=true 场景不应触发 store-disabled strict 规则")
  1855. }
  1856. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPrevResponsePreflightSkipForFunctionCallOutput(t *testing.T) {
  1857. gin.SetMode(gin.TestMode)
  1858. cfg := &config.Config{}
  1859. cfg.Security.URLAllowlist.Enabled = false
  1860. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1861. cfg.Gateway.OpenAIWS.Enabled = true
  1862. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1863. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1864. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1865. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1866. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1867. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1868. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  1869. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  1870. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  1871. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1872. captureConn := &openAIWSCaptureConn{
  1873. events: [][]byte{
  1874. []byte(`{"type":"response.completed","response":{"id":"resp_preflight_skip_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1875. []byte(`{"type":"response.completed","response":{"id":"resp_preflight_skip_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1876. },
  1877. }
  1878. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1879. pool := newOpenAIWSConnPool(cfg)
  1880. pool.setClientDialerForTest(captureDialer)
  1881. svc := &OpenAIGatewayService{
  1882. cfg: cfg,
  1883. httpUpstream: &httpUpstreamRecorder{},
  1884. cache: &stubGatewayCache{},
  1885. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1886. toolCorrector: NewCodexToolCorrector(),
  1887. openaiWSPool: pool,
  1888. }
  1889. account := &Account{
  1890. ID: 141,
  1891. Name: "openai-ingress-prev-preflight-skip-fco",
  1892. Platform: PlatformOpenAI,
  1893. Type: AccountTypeAPIKey,
  1894. Status: StatusActive,
  1895. Schedulable: true,
  1896. Concurrency: 1,
  1897. Credentials: map[string]any{
  1898. "api_key": "sk-test",
  1899. },
  1900. Extra: map[string]any{
  1901. "responses_websockets_v2_enabled": true,
  1902. },
  1903. }
  1904. serverErrCh := make(chan error, 1)
  1905. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1906. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  1907. CompressionMode: coderws.CompressionContextTakeover,
  1908. })
  1909. if err != nil {
  1910. serverErrCh <- err
  1911. return
  1912. }
  1913. defer func() {
  1914. _ = conn.CloseNow()
  1915. }()
  1916. rec := httptest.NewRecorder()
  1917. ginCtx, _ := gin.CreateTestContext(rec)
  1918. req := r.Clone(r.Context())
  1919. req.Header = req.Header.Clone()
  1920. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  1921. ginCtx.Request = req
  1922. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  1923. msgType, firstMessage, readErr := conn.Read(readCtx)
  1924. cancel()
  1925. if readErr != nil {
  1926. serverErrCh <- readErr
  1927. return
  1928. }
  1929. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  1930. serverErrCh <- errors.New("unsupported websocket client message type")
  1931. return
  1932. }
  1933. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  1934. }))
  1935. defer wsServer.Close()
  1936. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1937. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1938. cancelDial()
  1939. require.NoError(t, err)
  1940. defer func() {
  1941. _ = clientConn.CloseNow()
  1942. }()
  1943. writeMessage := func(payload string) {
  1944. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  1945. defer cancel()
  1946. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  1947. }
  1948. readMessage := func() []byte {
  1949. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  1950. defer cancel()
  1951. msgType, message, readErr := clientConn.Read(readCtx)
  1952. require.NoError(t, readErr)
  1953. require.Equal(t, coderws.MessageText, msgType)
  1954. return message
  1955. }
  1956. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false}`)
  1957. firstTurn := readMessage()
  1958. require.Equal(t, "resp_preflight_skip_1", gjson.GetBytes(firstTurn, "response.id").String())
  1959. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"previous_response_id":"resp_stale_external","input":[{"type":"function_call_output","call_id":"call_1","output":"ok"}]}`)
  1960. secondTurn := readMessage()
  1961. require.Equal(t, "resp_preflight_skip_2", gjson.GetBytes(secondTurn, "response.id").String())
  1962. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  1963. select {
  1964. case serverErr := <-serverErrCh:
  1965. require.NoError(t, serverErr)
  1966. case <-time.After(5 * time.Second):
  1967. t.Fatal("等待 ingress websocket 结束超时")
  1968. }
  1969. require.Equal(t, 1, captureDialer.DialCount())
  1970. require.Len(t, captureConn.writes, 2)
  1971. require.Equal(t, "resp_stale_external", gjson.Get(requestToJSONString(captureConn.writes[1]), "previous_response_id").String(), "function_call_output 场景不应预改写 previous_response_id")
  1972. }
  1973. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledFunctionCallOutputAutoAttachPreviousResponseID(t *testing.T) {
  1974. gin.SetMode(gin.TestMode)
  1975. cfg := &config.Config{}
  1976. cfg.Security.URLAllowlist.Enabled = false
  1977. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1978. cfg.Gateway.OpenAIWS.Enabled = true
  1979. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1980. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1981. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1982. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1983. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1984. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1985. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  1986. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  1987. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  1988. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1989. captureConn := &openAIWSCaptureConn{
  1990. events: [][]byte{
  1991. []byte(`{"type":"response.completed","response":{"id":"resp_auto_prev_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1992. []byte(`{"type":"response.completed","response":{"id":"resp_auto_prev_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1993. },
  1994. }
  1995. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1996. pool := newOpenAIWSConnPool(cfg)
  1997. pool.setClientDialerForTest(captureDialer)
  1998. svc := &OpenAIGatewayService{
  1999. cfg: cfg,
  2000. httpUpstream: &httpUpstreamRecorder{},
  2001. cache: &stubGatewayCache{},
  2002. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  2003. toolCorrector: NewCodexToolCorrector(),
  2004. openaiWSPool: pool,
  2005. }
  2006. account := &Account{
  2007. ID: 143,
  2008. Name: "openai-ingress-fco-auto-prev",
  2009. Platform: PlatformOpenAI,
  2010. Type: AccountTypeAPIKey,
  2011. Status: StatusActive,
  2012. Schedulable: true,
  2013. Concurrency: 1,
  2014. Credentials: map[string]any{
  2015. "api_key": "sk-test",
  2016. },
  2017. Extra: map[string]any{
  2018. "responses_websockets_v2_enabled": true,
  2019. },
  2020. }
  2021. serverErrCh := make(chan error, 1)
  2022. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2023. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  2024. CompressionMode: coderws.CompressionContextTakeover,
  2025. })
  2026. if err != nil {
  2027. serverErrCh <- err
  2028. return
  2029. }
  2030. defer func() {
  2031. _ = conn.CloseNow()
  2032. }()
  2033. rec := httptest.NewRecorder()
  2034. ginCtx, _ := gin.CreateTestContext(rec)
  2035. req := r.Clone(r.Context())
  2036. req.Header = req.Header.Clone()
  2037. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  2038. ginCtx.Request = req
  2039. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  2040. msgType, firstMessage, readErr := conn.Read(readCtx)
  2041. cancel()
  2042. if readErr != nil {
  2043. serverErrCh <- readErr
  2044. return
  2045. }
  2046. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  2047. serverErrCh <- errors.New("unsupported websocket client message type")
  2048. return
  2049. }
  2050. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  2051. }))
  2052. defer wsServer.Close()
  2053. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2054. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  2055. cancelDial()
  2056. require.NoError(t, err)
  2057. defer func() {
  2058. _ = clientConn.CloseNow()
  2059. }()
  2060. writeMessage := func(payload string) {
  2061. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2062. defer cancel()
  2063. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  2064. }
  2065. readMessage := func() []byte {
  2066. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2067. defer cancel()
  2068. msgType, message, readErr := clientConn.Read(readCtx)
  2069. require.NoError(t, readErr)
  2070. require.Equal(t, coderws.MessageText, msgType)
  2071. return message
  2072. }
  2073. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  2074. firstTurn := readMessage()
  2075. require.Equal(t, "resp_auto_prev_1", gjson.GetBytes(firstTurn, "response.id").String())
  2076. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"function_call_output","call_id":"call_auto_1","output":"ok"}]}`)
  2077. secondTurn := readMessage()
  2078. require.Equal(t, "resp_auto_prev_2", gjson.GetBytes(secondTurn, "response.id").String())
  2079. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  2080. select {
  2081. case serverErr := <-serverErrCh:
  2082. require.NoError(t, serverErr)
  2083. case <-time.After(5 * time.Second):
  2084. t.Fatal("等待 ingress websocket 结束超时")
  2085. }
  2086. require.Equal(t, 1, captureDialer.DialCount())
  2087. require.Len(t, captureConn.writes, 2)
  2088. require.Equal(t, "resp_auto_prev_1", gjson.Get(requestToJSONString(captureConn.writes[1]), "previous_response_id").String(), "function_call_output 缺失 previous_response_id 时应回填上一轮响应 ID")
  2089. }
  2090. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledToolSearchOutputAutoAttachesPreviousResponseID(t *testing.T) {
  2091. gin.SetMode(gin.TestMode)
  2092. cfg := &config.Config{}
  2093. cfg.Security.URLAllowlist.Enabled = false
  2094. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2095. cfg.Gateway.OpenAIWS.Enabled = true
  2096. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2097. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2098. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2099. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  2100. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  2101. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  2102. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  2103. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2104. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2105. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2106. captureConn := &openAIWSCaptureConn{
  2107. events: [][]byte{
  2108. []byte(`{"type":"response.completed","response":{"id":"resp_tool_search_prev_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2109. []byte(`{"type":"response.completed","response":{"id":"resp_tool_search_prev_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2110. },
  2111. }
  2112. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  2113. pool := newOpenAIWSConnPool(cfg)
  2114. pool.setClientDialerForTest(captureDialer)
  2115. svc := &OpenAIGatewayService{
  2116. cfg: cfg,
  2117. httpUpstream: &httpUpstreamRecorder{},
  2118. cache: &stubGatewayCache{},
  2119. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  2120. toolCorrector: NewCodexToolCorrector(),
  2121. openaiWSPool: pool,
  2122. }
  2123. account := &Account{
  2124. ID: 145,
  2125. Name: "openai-ingress-tool-search-output-auto-prev",
  2126. Platform: PlatformOpenAI,
  2127. Type: AccountTypeAPIKey,
  2128. Status: StatusActive,
  2129. Schedulable: true,
  2130. Concurrency: 1,
  2131. Credentials: map[string]any{
  2132. "api_key": "sk-test",
  2133. },
  2134. Extra: map[string]any{
  2135. "responses_websockets_v2_enabled": true,
  2136. },
  2137. }
  2138. serverErrCh := make(chan error, 1)
  2139. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2140. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  2141. CompressionMode: coderws.CompressionContextTakeover,
  2142. })
  2143. if err != nil {
  2144. serverErrCh <- err
  2145. return
  2146. }
  2147. defer func() {
  2148. _ = conn.CloseNow()
  2149. }()
  2150. rec := httptest.NewRecorder()
  2151. ginCtx, _ := gin.CreateTestContext(rec)
  2152. req := r.Clone(r.Context())
  2153. req.Header = req.Header.Clone()
  2154. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  2155. ginCtx.Request = req
  2156. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  2157. msgType, firstMessage, readErr := conn.Read(readCtx)
  2158. cancel()
  2159. if readErr != nil {
  2160. serverErrCh <- readErr
  2161. return
  2162. }
  2163. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  2164. serverErrCh <- errors.New("unsupported websocket client message type")
  2165. return
  2166. }
  2167. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  2168. }))
  2169. defer wsServer.Close()
  2170. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2171. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  2172. cancelDial()
  2173. require.NoError(t, err)
  2174. defer func() {
  2175. _ = clientConn.CloseNow()
  2176. }()
  2177. writeMessage := func(payload string) {
  2178. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2179. defer cancel()
  2180. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  2181. }
  2182. readMessage := func() []byte {
  2183. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2184. defer cancel()
  2185. msgType, message, readErr := clientConn.Read(readCtx)
  2186. require.NoError(t, readErr)
  2187. require.Equal(t, coderws.MessageText, msgType)
  2188. return message
  2189. }
  2190. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  2191. firstTurn := readMessage()
  2192. require.Equal(t, "resp_tool_search_prev_1", gjson.GetBytes(firstTurn, "response.id").String())
  2193. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"tool_search_output","call_id":"call_search_1","output":"ok"}]}`)
  2194. secondTurn := readMessage()
  2195. require.Equal(t, "resp_tool_search_prev_2", gjson.GetBytes(secondTurn, "response.id").String())
  2196. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  2197. select {
  2198. case serverErr := <-serverErrCh:
  2199. require.NoError(t, serverErr)
  2200. case <-time.After(5 * time.Second):
  2201. t.Fatal("等待 ingress websocket 结束超时")
  2202. }
  2203. require.Equal(t, 1, captureDialer.DialCount())
  2204. require.Len(t, captureConn.writes, 2)
  2205. secondWrite := requestToJSONString(captureConn.writes[1])
  2206. require.Equal(t, "resp_tool_search_prev_1", gjson.Get(secondWrite, "previous_response_id").String(), "tool_search_output 缺失 previous_response_id 时应回填上一轮响应 ID")
  2207. require.Equal(t, "tool_search_output", gjson.Get(secondWrite, "input.0.type").String())
  2208. require.Equal(t, "call_search_1", gjson.Get(secondWrite, "input.0.call_id").String())
  2209. }
  2210. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledFunctionCallOutputSkipsAutoAttachWhenLastResponseIDMissing(t *testing.T) {
  2211. gin.SetMode(gin.TestMode)
  2212. cfg := &config.Config{}
  2213. cfg.Security.URLAllowlist.Enabled = false
  2214. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2215. cfg.Gateway.OpenAIWS.Enabled = true
  2216. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2217. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2218. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2219. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  2220. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  2221. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  2222. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  2223. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2224. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2225. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2226. captureConn := &openAIWSCaptureConn{
  2227. events: [][]byte{
  2228. []byte(`{"type":"response.completed","response":{"model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2229. []byte(`{"type":"response.completed","response":{"id":"resp_auto_prev_skip_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2230. },
  2231. }
  2232. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  2233. pool := newOpenAIWSConnPool(cfg)
  2234. pool.setClientDialerForTest(captureDialer)
  2235. svc := &OpenAIGatewayService{
  2236. cfg: cfg,
  2237. httpUpstream: &httpUpstreamRecorder{},
  2238. cache: &stubGatewayCache{},
  2239. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  2240. toolCorrector: NewCodexToolCorrector(),
  2241. openaiWSPool: pool,
  2242. }
  2243. account := &Account{
  2244. ID: 144,
  2245. Name: "openai-ingress-fco-auto-prev-skip",
  2246. Platform: PlatformOpenAI,
  2247. Type: AccountTypeAPIKey,
  2248. Status: StatusActive,
  2249. Schedulable: true,
  2250. Concurrency: 1,
  2251. Credentials: map[string]any{
  2252. "api_key": "sk-test",
  2253. },
  2254. Extra: map[string]any{
  2255. "responses_websockets_v2_enabled": true,
  2256. },
  2257. }
  2258. serverErrCh := make(chan error, 1)
  2259. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2260. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  2261. CompressionMode: coderws.CompressionContextTakeover,
  2262. })
  2263. if err != nil {
  2264. serverErrCh <- err
  2265. return
  2266. }
  2267. defer func() {
  2268. _ = conn.CloseNow()
  2269. }()
  2270. rec := httptest.NewRecorder()
  2271. ginCtx, _ := gin.CreateTestContext(rec)
  2272. req := r.Clone(r.Context())
  2273. req.Header = req.Header.Clone()
  2274. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  2275. ginCtx.Request = req
  2276. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  2277. msgType, firstMessage, readErr := conn.Read(readCtx)
  2278. cancel()
  2279. if readErr != nil {
  2280. serverErrCh <- readErr
  2281. return
  2282. }
  2283. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  2284. serverErrCh <- errors.New("unsupported websocket client message type")
  2285. return
  2286. }
  2287. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  2288. }))
  2289. defer wsServer.Close()
  2290. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2291. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  2292. cancelDial()
  2293. require.NoError(t, err)
  2294. defer func() {
  2295. _ = clientConn.CloseNow()
  2296. }()
  2297. writeMessage := func(payload string) {
  2298. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2299. defer cancel()
  2300. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  2301. }
  2302. readMessage := func() []byte {
  2303. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2304. defer cancel()
  2305. msgType, message, readErr := clientConn.Read(readCtx)
  2306. require.NoError(t, readErr)
  2307. require.Equal(t, coderws.MessageText, msgType)
  2308. return message
  2309. }
  2310. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  2311. firstTurn := readMessage()
  2312. require.Equal(t, "response.completed", gjson.GetBytes(firstTurn, "type").String())
  2313. require.Empty(t, gjson.GetBytes(firstTurn, "response.id").String(), "首轮响应不返回 response.id,模拟无法推导续链锚点")
  2314. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"function_call_output","call_id":"call_auto_skip_1","output":"ok"}]}`)
  2315. secondTurn := readMessage()
  2316. require.Equal(t, "resp_auto_prev_skip_2", gjson.GetBytes(secondTurn, "response.id").String())
  2317. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  2318. select {
  2319. case serverErr := <-serverErrCh:
  2320. require.NoError(t, serverErr)
  2321. case <-time.After(5 * time.Second):
  2322. t.Fatal("等待 ingress websocket 结束超时")
  2323. }
  2324. require.Equal(t, 1, captureDialer.DialCount())
  2325. require.Len(t, captureConn.writes, 2)
  2326. require.False(t, gjson.Get(requestToJSONString(captureConn.writes[1]), "previous_response_id").Exists(), "上一轮缺失 response.id 时不应自动补齐 previous_response_id")
  2327. }
  2328. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledFunctionCallOutputSkipsAutoAttachWhenToolCallContextPresent(t *testing.T) {
  2329. gin.SetMode(gin.TestMode)
  2330. cfg := &config.Config{}
  2331. cfg.Security.URLAllowlist.Enabled = false
  2332. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2333. cfg.Gateway.OpenAIWS.Enabled = true
  2334. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2335. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2336. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2337. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  2338. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  2339. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  2340. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  2341. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2342. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2343. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2344. captureConn := &openAIWSCaptureConn{
  2345. events: [][]byte{
  2346. []byte(`{"type":"response.completed","response":{"id":"resp_auto_prev_ctx_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2347. []byte(`{"type":"response.completed","response":{"id":"resp_auto_prev_ctx_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2348. },
  2349. }
  2350. captureDialer := &openAIWSQueueDialer{
  2351. conns: []openAIWSClientConn{captureConn},
  2352. }
  2353. pool := newOpenAIWSConnPool(cfg)
  2354. pool.setClientDialerForTest(captureDialer)
  2355. svc := &OpenAIGatewayService{
  2356. cfg: cfg,
  2357. httpUpstream: &httpUpstreamRecorder{},
  2358. cache: &stubGatewayCache{},
  2359. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  2360. toolCorrector: NewCodexToolCorrector(),
  2361. openaiWSPool: pool,
  2362. }
  2363. account := &Account{
  2364. ID: 114,
  2365. Name: "openai-ingress-tool-context",
  2366. Platform: PlatformOpenAI,
  2367. Type: AccountTypeAPIKey,
  2368. Status: StatusActive,
  2369. Schedulable: true,
  2370. Concurrency: 1,
  2371. Credentials: map[string]any{
  2372. "api_key": "sk-test",
  2373. },
  2374. Extra: map[string]any{
  2375. "responses_websockets_v2_enabled": true,
  2376. },
  2377. }
  2378. serverErrCh := make(chan error, 1)
  2379. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2380. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  2381. CompressionMode: coderws.CompressionContextTakeover,
  2382. })
  2383. if err != nil {
  2384. serverErrCh <- err
  2385. return
  2386. }
  2387. defer func() {
  2388. _ = conn.CloseNow()
  2389. }()
  2390. rec := httptest.NewRecorder()
  2391. ginCtx, _ := gin.CreateTestContext(rec)
  2392. req := r.Clone(r.Context())
  2393. req.Header = req.Header.Clone()
  2394. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  2395. ginCtx.Request = req
  2396. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  2397. msgType, firstMessage, readErr := conn.Read(readCtx)
  2398. cancel()
  2399. if readErr != nil {
  2400. serverErrCh <- readErr
  2401. return
  2402. }
  2403. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  2404. serverErrCh <- errors.New("unsupported websocket client message type")
  2405. return
  2406. }
  2407. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  2408. }))
  2409. defer wsServer.Close()
  2410. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2411. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  2412. cancelDial()
  2413. require.NoError(t, err)
  2414. defer func() {
  2415. _ = clientConn.CloseNow()
  2416. }()
  2417. writeMessage := func(payload string) {
  2418. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2419. defer cancel()
  2420. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  2421. }
  2422. readMessage := func() []byte {
  2423. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2424. defer cancel()
  2425. msgType, message, readErr := clientConn.Read(readCtx)
  2426. require.NoError(t, readErr)
  2427. require.Equal(t, coderws.MessageText, msgType)
  2428. return message
  2429. }
  2430. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  2431. firstTurn := readMessage()
  2432. require.Equal(t, "resp_auto_prev_ctx_1", gjson.GetBytes(firstTurn, "response.id").String())
  2433. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"function_call","call_id":"call_ctx_1","name":"shell","arguments":"{}"},{"type":"function_call_output","call_id":"call_ctx_1","output":"ok"},{"type":"message","role":"user","content":[{"type":"input_text","text":"retry"}]}]}`)
  2434. secondTurn := readMessage()
  2435. require.Equal(t, "resp_auto_prev_ctx_2", gjson.GetBytes(secondTurn, "response.id").String())
  2436. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  2437. select {
  2438. case serverErr := <-serverErrCh:
  2439. require.NoError(t, serverErr)
  2440. case <-time.After(5 * time.Second):
  2441. t.Fatal("等待 ingress websocket 结束超时")
  2442. }
  2443. require.Equal(t, 1, captureDialer.DialCount())
  2444. require.Len(t, captureConn.writes, 2)
  2445. require.False(t, gjson.Get(requestToJSONString(captureConn.writes[1]), "previous_response_id").Exists(), "请求已包含 function_call 上下文时不应自动补齐 previous_response_id")
  2446. }
  2447. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledFunctionCallOutputAutoAttachWhenOnlyItemReferencesPresent(t *testing.T) {
  2448. gin.SetMode(gin.TestMode)
  2449. cfg := &config.Config{}
  2450. cfg.Security.URLAllowlist.Enabled = false
  2451. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2452. cfg.Gateway.OpenAIWS.Enabled = true
  2453. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2454. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2455. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2456. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  2457. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  2458. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  2459. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  2460. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2461. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2462. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2463. captureConn := &openAIWSCaptureConn{
  2464. events: [][]byte{
  2465. []byte(`{"type":"response.completed","response":{"id":"resp_auto_prev_ref_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2466. []byte(`{"type":"response.completed","response":{"id":"resp_auto_prev_ref_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2467. },
  2468. }
  2469. captureDialer := &openAIWSQueueDialer{
  2470. conns: []openAIWSClientConn{captureConn},
  2471. }
  2472. pool := newOpenAIWSConnPool(cfg)
  2473. pool.setClientDialerForTest(captureDialer)
  2474. svc := &OpenAIGatewayService{
  2475. cfg: cfg,
  2476. httpUpstream: &httpUpstreamRecorder{},
  2477. cache: &stubGatewayCache{},
  2478. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  2479. toolCorrector: NewCodexToolCorrector(),
  2480. openaiWSPool: pool,
  2481. }
  2482. account := &Account{
  2483. ID: 115,
  2484. Name: "openai-ingress-item-reference",
  2485. Platform: PlatformOpenAI,
  2486. Type: AccountTypeAPIKey,
  2487. Status: StatusActive,
  2488. Schedulable: true,
  2489. Concurrency: 1,
  2490. Credentials: map[string]any{
  2491. "api_key": "sk-test",
  2492. },
  2493. Extra: map[string]any{
  2494. "responses_websockets_v2_enabled": true,
  2495. },
  2496. }
  2497. serverErrCh := make(chan error, 1)
  2498. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2499. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  2500. CompressionMode: coderws.CompressionContextTakeover,
  2501. })
  2502. if err != nil {
  2503. serverErrCh <- err
  2504. return
  2505. }
  2506. defer func() {
  2507. _ = conn.CloseNow()
  2508. }()
  2509. rec := httptest.NewRecorder()
  2510. ginCtx, _ := gin.CreateTestContext(rec)
  2511. req := r.Clone(r.Context())
  2512. req.Header = req.Header.Clone()
  2513. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  2514. ginCtx.Request = req
  2515. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  2516. msgType, firstMessage, readErr := conn.Read(readCtx)
  2517. cancel()
  2518. if readErr != nil {
  2519. serverErrCh <- readErr
  2520. return
  2521. }
  2522. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  2523. serverErrCh <- errors.New("unsupported websocket client message type")
  2524. return
  2525. }
  2526. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  2527. }))
  2528. defer wsServer.Close()
  2529. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2530. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  2531. cancelDial()
  2532. require.NoError(t, err)
  2533. defer func() {
  2534. _ = clientConn.CloseNow()
  2535. }()
  2536. writeMessage := func(payload string) {
  2537. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2538. defer cancel()
  2539. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  2540. }
  2541. readMessage := func() []byte {
  2542. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2543. defer cancel()
  2544. msgType, message, readErr := clientConn.Read(readCtx)
  2545. require.NoError(t, readErr)
  2546. require.Equal(t, coderws.MessageText, msgType)
  2547. return message
  2548. }
  2549. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  2550. firstTurn := readMessage()
  2551. require.Equal(t, "resp_auto_prev_ref_1", gjson.GetBytes(firstTurn, "response.id").String())
  2552. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"item_reference","id":"call_ref_1"},{"type":"function_call_output","call_id":"call_ref_1","output":"ok"},{"type":"message","role":"user","content":[{"type":"input_text","text":"retry"}]}]}`)
  2553. secondTurn := readMessage()
  2554. require.Equal(t, "resp_auto_prev_ref_2", gjson.GetBytes(secondTurn, "response.id").String())
  2555. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  2556. select {
  2557. case serverErr := <-serverErrCh:
  2558. require.NoError(t, serverErr)
  2559. case <-time.After(5 * time.Second):
  2560. t.Fatal("等待 ingress websocket 结束超时")
  2561. }
  2562. require.Equal(t, 1, captureDialer.DialCount())
  2563. require.Len(t, captureConn.writes, 2)
  2564. require.Equal(t, "resp_auto_prev_ref_1", gjson.Get(requestToJSONString(captureConn.writes[1]), "previous_response_id").String(), "仅有 item_reference 不足以自包含 function_call_output,应回填上一轮响应 ID")
  2565. }
  2566. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PreflightPingFailReconnectsBeforeTurn(t *testing.T) {
  2567. gin.SetMode(gin.TestMode)
  2568. prevPreflightPingIdle := openAIWSIngressPreflightPingIdle
  2569. openAIWSIngressPreflightPingIdle = 0
  2570. defer func() {
  2571. openAIWSIngressPreflightPingIdle = prevPreflightPingIdle
  2572. }()
  2573. cfg := &config.Config{}
  2574. cfg.Security.URLAllowlist.Enabled = false
  2575. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2576. cfg.Gateway.OpenAIWS.Enabled = true
  2577. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2578. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2579. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2580. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  2581. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  2582. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  2583. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  2584. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2585. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2586. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2587. firstConn := &openAIWSPreflightFailConn{
  2588. events: [][]byte{
  2589. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2590. },
  2591. }
  2592. secondConn := &openAIWSCaptureConn{
  2593. events: [][]byte{
  2594. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2595. },
  2596. }
  2597. dialer := &openAIWSQueueDialer{
  2598. conns: []openAIWSClientConn{firstConn, secondConn},
  2599. }
  2600. pool := newOpenAIWSConnPool(cfg)
  2601. pool.setClientDialerForTest(dialer)
  2602. svc := &OpenAIGatewayService{
  2603. cfg: cfg,
  2604. httpUpstream: &httpUpstreamRecorder{},
  2605. cache: &stubGatewayCache{},
  2606. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  2607. toolCorrector: NewCodexToolCorrector(),
  2608. openaiWSPool: pool,
  2609. }
  2610. account := &Account{
  2611. ID: 116,
  2612. Name: "openai-ingress-preflight-ping",
  2613. Platform: PlatformOpenAI,
  2614. Type: AccountTypeAPIKey,
  2615. Status: StatusActive,
  2616. Schedulable: true,
  2617. Concurrency: 1,
  2618. Credentials: map[string]any{
  2619. "api_key": "sk-test",
  2620. },
  2621. Extra: map[string]any{
  2622. "responses_websockets_v2_enabled": true,
  2623. },
  2624. }
  2625. serverErrCh := make(chan error, 1)
  2626. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2627. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  2628. CompressionMode: coderws.CompressionContextTakeover,
  2629. })
  2630. if err != nil {
  2631. serverErrCh <- err
  2632. return
  2633. }
  2634. defer func() {
  2635. _ = conn.CloseNow()
  2636. }()
  2637. rec := httptest.NewRecorder()
  2638. ginCtx, _ := gin.CreateTestContext(rec)
  2639. req := r.Clone(r.Context())
  2640. req.Header = req.Header.Clone()
  2641. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  2642. ginCtx.Request = req
  2643. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  2644. msgType, firstMessage, readErr := conn.Read(readCtx)
  2645. cancel()
  2646. if readErr != nil {
  2647. serverErrCh <- readErr
  2648. return
  2649. }
  2650. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  2651. serverErrCh <- errors.New("unsupported websocket client message type")
  2652. return
  2653. }
  2654. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  2655. }))
  2656. defer wsServer.Close()
  2657. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2658. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  2659. cancelDial()
  2660. require.NoError(t, err)
  2661. defer func() {
  2662. _ = clientConn.CloseNow()
  2663. }()
  2664. writeMessage := func(payload string) {
  2665. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2666. defer cancel()
  2667. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  2668. }
  2669. readMessage := func() []byte {
  2670. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2671. defer cancel()
  2672. msgType, message, readErr := clientConn.Read(readCtx)
  2673. require.NoError(t, readErr)
  2674. require.Equal(t, coderws.MessageText, msgType)
  2675. return message
  2676. }
  2677. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false}`)
  2678. firstTurn := readMessage()
  2679. require.Equal(t, "resp_turn_ping_1", gjson.GetBytes(firstTurn, "response.id").String())
  2680. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_turn_ping_1"}`)
  2681. secondTurn := readMessage()
  2682. require.Equal(t, "resp_turn_ping_2", gjson.GetBytes(secondTurn, "response.id").String())
  2683. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  2684. select {
  2685. case serverErr := <-serverErrCh:
  2686. require.NoError(t, serverErr)
  2687. case <-time.After(5 * time.Second):
  2688. t.Fatal("等待 ingress websocket 结束超时")
  2689. }
  2690. require.Equal(t, 2, dialer.DialCount(), "第二轮 turn 前 ping 失败应触发换连")
  2691. require.Equal(t, 1, firstConn.WriteCount(), "preflight ping 失败后不应继续向旧连接发送第二轮 turn")
  2692. require.GreaterOrEqual(t, firstConn.PingCount(), 1, "第二轮前应对旧连接执行 preflight ping")
  2693. }
  2694. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledStrictAffinityPreflightPingFailAutoRecoveryReconnects(t *testing.T) {
  2695. gin.SetMode(gin.TestMode)
  2696. prevPreflightPingIdle := openAIWSIngressPreflightPingIdle
  2697. openAIWSIngressPreflightPingIdle = 0
  2698. defer func() {
  2699. openAIWSIngressPreflightPingIdle = prevPreflightPingIdle
  2700. }()
  2701. cfg := &config.Config{}
  2702. cfg.Security.URLAllowlist.Enabled = false
  2703. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2704. cfg.Gateway.OpenAIWS.Enabled = true
  2705. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2706. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2707. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2708. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  2709. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  2710. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  2711. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  2712. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2713. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2714. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2715. firstConn := &openAIWSPreflightFailConn{
  2716. events: [][]byte{
  2717. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_strict_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2718. },
  2719. }
  2720. secondConn := &openAIWSCaptureConn{
  2721. events: [][]byte{
  2722. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_strict_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2723. },
  2724. }
  2725. dialer := &openAIWSQueueDialer{
  2726. conns: []openAIWSClientConn{firstConn, secondConn},
  2727. }
  2728. pool := newOpenAIWSConnPool(cfg)
  2729. pool.setClientDialerForTest(dialer)
  2730. svc := &OpenAIGatewayService{
  2731. cfg: cfg,
  2732. httpUpstream: &httpUpstreamRecorder{},
  2733. cache: &stubGatewayCache{},
  2734. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  2735. toolCorrector: NewCodexToolCorrector(),
  2736. openaiWSPool: pool,
  2737. }
  2738. account := &Account{
  2739. ID: 121,
  2740. Name: "openai-ingress-preflight-ping-strict-affinity",
  2741. Platform: PlatformOpenAI,
  2742. Type: AccountTypeAPIKey,
  2743. Status: StatusActive,
  2744. Schedulable: true,
  2745. Concurrency: 1,
  2746. Credentials: map[string]any{
  2747. "api_key": "sk-test",
  2748. },
  2749. Extra: map[string]any{
  2750. "responses_websockets_v2_enabled": true,
  2751. },
  2752. }
  2753. serverErrCh := make(chan error, 1)
  2754. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2755. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  2756. CompressionMode: coderws.CompressionContextTakeover,
  2757. })
  2758. if err != nil {
  2759. serverErrCh <- err
  2760. return
  2761. }
  2762. defer func() {
  2763. _ = conn.CloseNow()
  2764. }()
  2765. rec := httptest.NewRecorder()
  2766. ginCtx, _ := gin.CreateTestContext(rec)
  2767. req := r.Clone(r.Context())
  2768. req.Header = req.Header.Clone()
  2769. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  2770. ginCtx.Request = req
  2771. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  2772. msgType, firstMessage, readErr := conn.Read(readCtx)
  2773. cancel()
  2774. if readErr != nil {
  2775. serverErrCh <- readErr
  2776. return
  2777. }
  2778. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  2779. serverErrCh <- errors.New("unsupported websocket client message type")
  2780. return
  2781. }
  2782. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  2783. }))
  2784. defer wsServer.Close()
  2785. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2786. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  2787. cancelDial()
  2788. require.NoError(t, err)
  2789. defer func() {
  2790. _ = clientConn.CloseNow()
  2791. }()
  2792. writeMessage := func(payload string) {
  2793. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2794. defer cancel()
  2795. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  2796. }
  2797. readMessage := func() []byte {
  2798. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2799. defer cancel()
  2800. msgType, message, readErr := clientConn.Read(readCtx)
  2801. require.NoError(t, readErr)
  2802. require.Equal(t, coderws.MessageText, msgType)
  2803. return message
  2804. }
  2805. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  2806. firstTurn := readMessage()
  2807. require.Equal(t, "resp_turn_ping_strict_1", gjson.GetBytes(firstTurn, "response.id").String())
  2808. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"previous_response_id":"resp_turn_ping_strict_1","input":[{"type":"input_text","text":"world"}]}`)
  2809. secondTurn := readMessage()
  2810. require.Equal(t, "resp_turn_ping_strict_2", gjson.GetBytes(secondTurn, "response.id").String())
  2811. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  2812. select {
  2813. case serverErr := <-serverErrCh:
  2814. require.NoError(t, serverErr)
  2815. case <-time.After(5 * time.Second):
  2816. t.Fatal("等待 ingress websocket 严格亲和自动恢复后结束超时")
  2817. }
  2818. require.Equal(t, 2, dialer.DialCount(), "严格亲和 preflight ping 失败后应自动降级并换连重放")
  2819. require.Equal(t, 1, firstConn.WriteCount(), "preflight ping 失败后不应继续在旧连接写第二轮")
  2820. require.GreaterOrEqual(t, firstConn.PingCount(), 1, "第二轮前应执行 preflight ping")
  2821. secondConn.mu.Lock()
  2822. secondWrites := append([]map[string]any(nil), secondConn.writes...)
  2823. secondConn.mu.Unlock()
  2824. require.Len(t, secondWrites, 1)
  2825. secondWrite := requestToJSONString(secondWrites[0])
  2826. require.False(t, gjson.Get(secondWrite, "previous_response_id").Exists(), "自动恢复重放应移除 previous_response_id")
  2827. require.Equal(t, 2, len(gjson.Get(secondWrite, "input").Array()), "自动恢复重放应使用完整 input 上下文")
  2828. require.Equal(t, "hello", gjson.Get(secondWrite, "input.0.text").String())
  2829. require.Equal(t, "world", gjson.Get(secondWrite, "input.1.text").String())
  2830. }
  2831. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPreflightPingFailReplaysFunctionCallOutputWithContext(t *testing.T) {
  2832. gin.SetMode(gin.TestMode)
  2833. prevPreflightPingIdle := openAIWSIngressPreflightPingIdle
  2834. openAIWSIngressPreflightPingIdle = 0
  2835. defer func() {
  2836. openAIWSIngressPreflightPingIdle = prevPreflightPingIdle
  2837. }()
  2838. cfg := &config.Config{}
  2839. cfg.Security.URLAllowlist.Enabled = false
  2840. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2841. cfg.Gateway.OpenAIWS.Enabled = true
  2842. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2843. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2844. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2845. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  2846. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  2847. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  2848. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  2849. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2850. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2851. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2852. firstConn := &openAIWSPreflightFailConn{
  2853. events: [][]byte{
  2854. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_replay_ctx_1","model":"gpt-5.1","output":[{"type":"function_call","id":"fc_replay_1","call_id":"call_replay_1","name":"shell","arguments":"{}"}],"usage":{"input_tokens":1,"output_tokens":1}}}`),
  2855. },
  2856. }
  2857. secondConn := &openAIWSCaptureConn{
  2858. events: [][]byte{
  2859. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_replay_ctx_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2860. },
  2861. }
  2862. dialer := &openAIWSQueueDialer{
  2863. conns: []openAIWSClientConn{firstConn, secondConn},
  2864. }
  2865. pool := newOpenAIWSConnPool(cfg)
  2866. pool.setClientDialerForTest(dialer)
  2867. svc := &OpenAIGatewayService{
  2868. cfg: cfg,
  2869. httpUpstream: &httpUpstreamRecorder{},
  2870. cache: &stubGatewayCache{},
  2871. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  2872. toolCorrector: NewCodexToolCorrector(),
  2873. openaiWSPool: pool,
  2874. }
  2875. account := &Account{
  2876. ID: 128,
  2877. Name: "openai-ingress-preflight-replay-function-output-with-context",
  2878. Platform: PlatformOpenAI,
  2879. Type: AccountTypeAPIKey,
  2880. Status: StatusActive,
  2881. Schedulable: true,
  2882. Concurrency: 1,
  2883. Credentials: map[string]any{
  2884. "api_key": "sk-test",
  2885. },
  2886. Extra: map[string]any{
  2887. "responses_websockets_v2_enabled": true,
  2888. },
  2889. }
  2890. serverErrCh := make(chan error, 1)
  2891. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2892. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  2893. CompressionMode: coderws.CompressionContextTakeover,
  2894. })
  2895. if err != nil {
  2896. serverErrCh <- err
  2897. return
  2898. }
  2899. defer func() {
  2900. _ = conn.CloseNow()
  2901. }()
  2902. rec := httptest.NewRecorder()
  2903. ginCtx, _ := gin.CreateTestContext(rec)
  2904. req := r.Clone(r.Context())
  2905. req.Header = req.Header.Clone()
  2906. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  2907. ginCtx.Request = req
  2908. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  2909. msgType, firstMessage, readErr := conn.Read(readCtx)
  2910. cancel()
  2911. if readErr != nil {
  2912. serverErrCh <- readErr
  2913. return
  2914. }
  2915. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  2916. serverErrCh <- errors.New("unsupported websocket client message type")
  2917. return
  2918. }
  2919. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  2920. }))
  2921. defer wsServer.Close()
  2922. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2923. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  2924. cancelDial()
  2925. require.NoError(t, err)
  2926. defer func() {
  2927. _ = clientConn.CloseNow()
  2928. }()
  2929. writeMessage := func(payload string) {
  2930. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2931. defer cancel()
  2932. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  2933. }
  2934. readMessage := func() []byte {
  2935. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  2936. defer cancel()
  2937. msgType, message, readErr := clientConn.Read(readCtx)
  2938. require.NoError(t, readErr)
  2939. require.Equal(t, coderws.MessageText, msgType)
  2940. return message
  2941. }
  2942. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"call tool"}]}]}`)
  2943. firstTurn := readMessage()
  2944. require.Equal(t, "resp_turn_ping_replay_ctx_1", gjson.GetBytes(firstTurn, "response.id").String())
  2945. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"previous_response_id":"resp_turn_ping_replay_ctx_1","input":[{"type":"function_call_output","call_id":"call_replay_1","output":"ok"}]}`)
  2946. secondTurn := readMessage()
  2947. require.Equal(t, "resp_turn_ping_replay_ctx_2", gjson.GetBytes(secondTurn, "response.id").String())
  2948. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  2949. select {
  2950. case serverErr := <-serverErrCh:
  2951. require.NoError(t, serverErr)
  2952. case <-time.After(5 * time.Second):
  2953. t.Fatal("等待 ingress websocket function_call_output 自包含重放后结束超时")
  2954. }
  2955. require.Equal(t, 2, dialer.DialCount(), "带完整 tool 上下文的 function_call_output 应在 ping 失败后换新连接重放")
  2956. require.Equal(t, 1, firstConn.WriteCount())
  2957. require.GreaterOrEqual(t, firstConn.PingCount(), 1)
  2958. secondConn.mu.Lock()
  2959. secondWrites := append([]map[string]any(nil), secondConn.writes...)
  2960. secondConn.mu.Unlock()
  2961. require.Len(t, secondWrites, 1)
  2962. secondWrite := requestToJSONString(secondWrites[0])
  2963. require.False(t, gjson.Get(secondWrite, "previous_response_id").Exists())
  2964. require.Equal(t, 3, len(gjson.Get(secondWrite, "input").Array()))
  2965. require.Equal(t, "message", gjson.Get(secondWrite, "input.0.type").String())
  2966. require.Equal(t, "function_call", gjson.Get(secondWrite, "input.1.type").String())
  2967. require.Equal(t, "call_replay_1", gjson.Get(secondWrite, "input.1.call_id").String())
  2968. require.Equal(t, "function_call_output", gjson.Get(secondWrite, "input.2.type").String())
  2969. require.Equal(t, "call_replay_1", gjson.Get(secondWrite, "input.2.call_id").String())
  2970. }
  2971. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPreflightPingFailClosesWhenFunctionCallOutputNeedsPreviousResponseID(t *testing.T) {
  2972. gin.SetMode(gin.TestMode)
  2973. prevPreflightPingIdle := openAIWSIngressPreflightPingIdle
  2974. openAIWSIngressPreflightPingIdle = 0
  2975. defer func() {
  2976. openAIWSIngressPreflightPingIdle = prevPreflightPingIdle
  2977. }()
  2978. cfg := &config.Config{}
  2979. cfg.Security.URLAllowlist.Enabled = false
  2980. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2981. cfg.Gateway.OpenAIWS.Enabled = true
  2982. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2983. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2984. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2985. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  2986. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  2987. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  2988. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  2989. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2990. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2991. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2992. firstConn := &openAIWSPreflightFailConn{
  2993. events: [][]byte{
  2994. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_replay_fc_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  2995. },
  2996. }
  2997. secondConn := &openAIWSCaptureConn{
  2998. events: [][]byte{
  2999. []byte(`{"type":"error","error":{"type":"invalid_request_error","code":"previous_response_not_found","message":"Previous response not found."}}`),
  3000. },
  3001. }
  3002. dialer := &openAIWSQueueDialer{
  3003. conns: []openAIWSClientConn{firstConn, secondConn},
  3004. }
  3005. pool := newOpenAIWSConnPool(cfg)
  3006. pool.setClientDialerForTest(dialer)
  3007. svc := &OpenAIGatewayService{
  3008. cfg: cfg,
  3009. httpUpstream: &httpUpstreamRecorder{},
  3010. cache: &stubGatewayCache{},
  3011. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  3012. toolCorrector: NewCodexToolCorrector(),
  3013. openaiWSPool: pool,
  3014. }
  3015. account := &Account{
  3016. ID: 129,
  3017. Name: "openai-ingress-preflight-replay-function-output",
  3018. Platform: PlatformOpenAI,
  3019. Type: AccountTypeAPIKey,
  3020. Status: StatusActive,
  3021. Schedulable: true,
  3022. Concurrency: 1,
  3023. Credentials: map[string]any{
  3024. "api_key": "sk-test",
  3025. },
  3026. Extra: map[string]any{
  3027. "responses_websockets_v2_enabled": true,
  3028. },
  3029. }
  3030. serverErrCh := make(chan error, 1)
  3031. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  3032. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  3033. CompressionMode: coderws.CompressionContextTakeover,
  3034. })
  3035. if err != nil {
  3036. serverErrCh <- err
  3037. return
  3038. }
  3039. defer func() {
  3040. _ = conn.CloseNow()
  3041. }()
  3042. rec := httptest.NewRecorder()
  3043. ginCtx, _ := gin.CreateTestContext(rec)
  3044. req := r.Clone(r.Context())
  3045. req.Header = req.Header.Clone()
  3046. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  3047. ginCtx.Request = req
  3048. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  3049. msgType, firstMessage, readErr := conn.Read(readCtx)
  3050. cancel()
  3051. if readErr != nil {
  3052. serverErrCh <- readErr
  3053. return
  3054. }
  3055. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  3056. serverErrCh <- errors.New("unsupported websocket client message type")
  3057. return
  3058. }
  3059. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  3060. }))
  3061. defer wsServer.Close()
  3062. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  3063. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  3064. cancelDial()
  3065. require.NoError(t, err)
  3066. defer func() {
  3067. _ = clientConn.CloseNow()
  3068. }()
  3069. writeMessage := func(payload string) {
  3070. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3071. defer cancel()
  3072. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  3073. }
  3074. readMessage := func() []byte {
  3075. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3076. defer cancel()
  3077. msgType, message, readErr := clientConn.Read(readCtx)
  3078. require.NoError(t, readErr)
  3079. require.Equal(t, coderws.MessageText, msgType)
  3080. return message
  3081. }
  3082. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"function_call","call_id":"call_other","name":"shell","arguments":"{}"},{"type":"function_call_output","call_id":"call_replay_1","output":"ok"}]}`)
  3083. firstTurn := readMessage()
  3084. require.Equal(t, "resp_turn_ping_replay_fc_1", gjson.GetBytes(firstTurn, "response.id").String())
  3085. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"previous_response_id":"resp_turn_ping_replay_fc_1","input":[{"type":"function_call_output","call_id":"call_replay_1","output":"ok"}]}`)
  3086. select {
  3087. case serverErr := <-serverErrCh:
  3088. require.Error(t, serverErr)
  3089. var closeErr *OpenAIWSClientCloseError
  3090. require.ErrorAs(t, serverErr, &closeErr)
  3091. require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode())
  3092. require.Contains(t, closeErr.Reason(), "upstream continuation connection is unavailable")
  3093. case <-time.After(5 * time.Second):
  3094. t.Fatal("等待 ingress websocket 结束超时")
  3095. }
  3096. require.Equal(t, 1, dialer.DialCount(), "需要 previous_response_id 的 function_call_output 在原连接不可用时不应换新连接重试")
  3097. secondConn.mu.Lock()
  3098. secondWrites := append([]map[string]any(nil), secondConn.writes...)
  3099. secondConn.mu.Unlock()
  3100. require.Empty(t, secondWrites, "不能把旧连接的 previous_response_id 发送到新上游,否则会触发 previous_response_not_found")
  3101. }
  3102. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledPreflightPingFailClosesWhenReplayHasFunctionCallOutput(t *testing.T) {
  3103. gin.SetMode(gin.TestMode)
  3104. prevPreflightPingIdle := openAIWSIngressPreflightPingIdle
  3105. openAIWSIngressPreflightPingIdle = 0
  3106. defer func() {
  3107. openAIWSIngressPreflightPingIdle = prevPreflightPingIdle
  3108. }()
  3109. cfg := &config.Config{}
  3110. cfg.Security.URLAllowlist.Enabled = false
  3111. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  3112. cfg.Gateway.OpenAIWS.Enabled = true
  3113. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  3114. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  3115. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  3116. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  3117. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  3118. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  3119. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  3120. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  3121. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  3122. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  3123. firstConn := &openAIWSPreflightFailConn{
  3124. events: [][]byte{
  3125. []byte(`{"type":"response.completed","response":{"id":"resp_turn_ping_replay_only_fc_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  3126. },
  3127. }
  3128. secondConn := &openAIWSCaptureConn{
  3129. events: [][]byte{
  3130. []byte(`{"type":"error","error":{"type":"invalid_request_error","message":"No tool call found for function call output with call_id call_replay_1.","param":"input"}}`),
  3131. },
  3132. }
  3133. dialer := &openAIWSQueueDialer{
  3134. conns: []openAIWSClientConn{firstConn, secondConn},
  3135. }
  3136. pool := newOpenAIWSConnPool(cfg)
  3137. pool.setClientDialerForTest(dialer)
  3138. svc := &OpenAIGatewayService{
  3139. cfg: cfg,
  3140. httpUpstream: &httpUpstreamRecorder{},
  3141. cache: &stubGatewayCache{},
  3142. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  3143. toolCorrector: NewCodexToolCorrector(),
  3144. openaiWSPool: pool,
  3145. }
  3146. account := &Account{
  3147. ID: 130,
  3148. Name: "openai-ingress-preflight-replay-only-function-output",
  3149. Platform: PlatformOpenAI,
  3150. Type: AccountTypeAPIKey,
  3151. Status: StatusActive,
  3152. Schedulable: true,
  3153. Concurrency: 1,
  3154. Credentials: map[string]any{
  3155. "api_key": "sk-test",
  3156. },
  3157. Extra: map[string]any{
  3158. "responses_websockets_v2_enabled": true,
  3159. },
  3160. }
  3161. serverErrCh := make(chan error, 1)
  3162. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  3163. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  3164. CompressionMode: coderws.CompressionContextTakeover,
  3165. })
  3166. if err != nil {
  3167. serverErrCh <- err
  3168. return
  3169. }
  3170. defer func() {
  3171. _ = conn.CloseNow()
  3172. }()
  3173. rec := httptest.NewRecorder()
  3174. ginCtx, _ := gin.CreateTestContext(rec)
  3175. req := r.Clone(r.Context())
  3176. req.Header = req.Header.Clone()
  3177. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  3178. ginCtx.Request = req
  3179. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  3180. msgType, firstMessage, readErr := conn.Read(readCtx)
  3181. cancel()
  3182. if readErr != nil {
  3183. serverErrCh <- readErr
  3184. return
  3185. }
  3186. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  3187. serverErrCh <- errors.New("unsupported websocket client message type")
  3188. return
  3189. }
  3190. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  3191. }))
  3192. defer wsServer.Close()
  3193. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  3194. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  3195. cancelDial()
  3196. require.NoError(t, err)
  3197. defer func() {
  3198. _ = clientConn.CloseNow()
  3199. }()
  3200. writeMessage := func(payload string) {
  3201. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3202. defer cancel()
  3203. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  3204. }
  3205. readMessage := func() []byte {
  3206. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3207. defer cancel()
  3208. msgType, message, readErr := clientConn.Read(readCtx)
  3209. require.NoError(t, readErr)
  3210. require.Equal(t, coderws.MessageText, msgType)
  3211. return message
  3212. }
  3213. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"function_call","call_id":"call_other","name":"shell","arguments":"{}"},{"type":"function_call_output","call_id":"call_replay_1","output":"ok"}]}`)
  3214. firstTurn := readMessage()
  3215. require.Equal(t, "resp_turn_ping_replay_only_fc_1", gjson.GetBytes(firstTurn, "response.id").String())
  3216. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"previous_response_id":"resp_turn_ping_replay_only_fc_1","input":[{"type":"input_text","text":"after tool output"}]}`)
  3217. select {
  3218. case serverErr := <-serverErrCh:
  3219. require.Error(t, serverErr)
  3220. var closeErr *OpenAIWSClientCloseError
  3221. require.ErrorAs(t, serverErr, &closeErr)
  3222. require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode())
  3223. require.Contains(t, closeErr.Reason(), "upstream continuation connection is unavailable")
  3224. case <-time.After(5 * time.Second):
  3225. t.Fatal("等待 ingress websocket 结束超时")
  3226. }
  3227. require.Equal(t, 1, dialer.DialCount(), "replay input 带 function_call_output 时不应换新连接重试")
  3228. secondConn.mu.Lock()
  3229. secondWrites := append([]map[string]any(nil), secondConn.writes...)
  3230. secondConn.mu.Unlock()
  3231. require.Empty(t, secondWrites, "不能把会触发 No tool call found 的重放请求发到新上游")
  3232. }
  3233. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_WriteFailBeforeDownstreamRetriesOnce(t *testing.T) {
  3234. gin.SetMode(gin.TestMode)
  3235. cfg := &config.Config{}
  3236. cfg.Security.URLAllowlist.Enabled = false
  3237. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  3238. cfg.Gateway.OpenAIWS.Enabled = true
  3239. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  3240. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  3241. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  3242. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  3243. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  3244. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  3245. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  3246. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  3247. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  3248. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  3249. firstConn := &openAIWSWriteFailAfterFirstTurnConn{
  3250. events: [][]byte{
  3251. []byte(`{"type":"response.completed","response":{"id":"resp_turn_write_retry_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  3252. },
  3253. }
  3254. secondConn := &openAIWSCaptureConn{
  3255. events: [][]byte{
  3256. []byte(`{"type":"response.completed","response":{"id":"resp_turn_write_retry_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  3257. },
  3258. }
  3259. dialer := &openAIWSQueueDialer{
  3260. conns: []openAIWSClientConn{firstConn, secondConn},
  3261. }
  3262. pool := newOpenAIWSConnPool(cfg)
  3263. pool.setClientDialerForTest(dialer)
  3264. svc := &OpenAIGatewayService{
  3265. cfg: cfg,
  3266. httpUpstream: &httpUpstreamRecorder{},
  3267. cache: &stubGatewayCache{},
  3268. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  3269. toolCorrector: NewCodexToolCorrector(),
  3270. openaiWSPool: pool,
  3271. }
  3272. account := &Account{
  3273. ID: 117,
  3274. Name: "openai-ingress-write-retry",
  3275. Platform: PlatformOpenAI,
  3276. Type: AccountTypeAPIKey,
  3277. Status: StatusActive,
  3278. Schedulable: true,
  3279. Concurrency: 1,
  3280. Credentials: map[string]any{
  3281. "api_key": "sk-test",
  3282. },
  3283. Extra: map[string]any{
  3284. "responses_websockets_v2_enabled": true,
  3285. },
  3286. }
  3287. var hooksMu sync.Mutex
  3288. beforeTurnCalls := make(map[int]int)
  3289. afterTurnCalls := make(map[int]int)
  3290. hooks := &OpenAIWSIngressHooks{
  3291. BeforeTurn: func(turn int) error {
  3292. hooksMu.Lock()
  3293. beforeTurnCalls[turn]++
  3294. hooksMu.Unlock()
  3295. return nil
  3296. },
  3297. AfterTurn: func(turn int, _ *OpenAIForwardResult, _ error) {
  3298. hooksMu.Lock()
  3299. afterTurnCalls[turn]++
  3300. hooksMu.Unlock()
  3301. },
  3302. }
  3303. serverErrCh := make(chan error, 1)
  3304. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  3305. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  3306. CompressionMode: coderws.CompressionContextTakeover,
  3307. })
  3308. if err != nil {
  3309. serverErrCh <- err
  3310. return
  3311. }
  3312. defer func() {
  3313. _ = conn.CloseNow()
  3314. }()
  3315. rec := httptest.NewRecorder()
  3316. ginCtx, _ := gin.CreateTestContext(rec)
  3317. req := r.Clone(r.Context())
  3318. req.Header = req.Header.Clone()
  3319. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  3320. ginCtx.Request = req
  3321. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  3322. msgType, firstMessage, readErr := conn.Read(readCtx)
  3323. cancel()
  3324. if readErr != nil {
  3325. serverErrCh <- readErr
  3326. return
  3327. }
  3328. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  3329. serverErrCh <- errors.New("unsupported websocket client message type")
  3330. return
  3331. }
  3332. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, hooks)
  3333. }))
  3334. defer wsServer.Close()
  3335. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  3336. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  3337. cancelDial()
  3338. require.NoError(t, err)
  3339. defer func() {
  3340. _ = clientConn.CloseNow()
  3341. }()
  3342. writeMessage := func(payload string) {
  3343. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3344. defer cancel()
  3345. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  3346. }
  3347. readMessage := func() []byte {
  3348. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3349. defer cancel()
  3350. msgType, message, readErr := clientConn.Read(readCtx)
  3351. require.NoError(t, readErr)
  3352. require.Equal(t, coderws.MessageText, msgType)
  3353. return message
  3354. }
  3355. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false}`)
  3356. firstTurn := readMessage()
  3357. require.Equal(t, "resp_turn_write_retry_1", gjson.GetBytes(firstTurn, "response.id").String())
  3358. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_turn_write_retry_1"}`)
  3359. secondTurn := readMessage()
  3360. require.Equal(t, "resp_turn_write_retry_2", gjson.GetBytes(secondTurn, "response.id").String())
  3361. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  3362. select {
  3363. case serverErr := <-serverErrCh:
  3364. require.NoError(t, serverErr)
  3365. case <-time.After(5 * time.Second):
  3366. t.Fatal("等待 ingress websocket 结束超时")
  3367. }
  3368. require.Equal(t, 2, dialer.DialCount(), "第二轮 turn 上游写失败且未写下游时应自动重试并换连")
  3369. hooksMu.Lock()
  3370. beforeTurn1 := beforeTurnCalls[1]
  3371. beforeTurn2 := beforeTurnCalls[2]
  3372. afterTurn1 := afterTurnCalls[1]
  3373. afterTurn2 := afterTurnCalls[2]
  3374. hooksMu.Unlock()
  3375. require.Equal(t, 1, beforeTurn1, "首轮 turn BeforeTurn 应执行一次")
  3376. require.Equal(t, 1, beforeTurn2, "同一 turn 重试不应重复触发 BeforeTurn")
  3377. require.Equal(t, 1, afterTurn1, "首轮 turn AfterTurn 应执行一次")
  3378. require.Equal(t, 1, afterTurn2, "第二轮 turn AfterTurn 应执行一次")
  3379. }
  3380. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PreviousResponseNotFoundRecoversByDroppingPrevID(t *testing.T) {
  3381. gin.SetMode(gin.TestMode)
  3382. cfg := &config.Config{}
  3383. cfg.Security.URLAllowlist.Enabled = false
  3384. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  3385. cfg.Gateway.OpenAIWS.Enabled = true
  3386. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  3387. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  3388. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  3389. cfg.Gateway.OpenAIWS.IngressPreviousResponseRecoveryEnabled = true
  3390. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  3391. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  3392. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  3393. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  3394. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  3395. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  3396. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  3397. firstConn := &openAIWSCaptureConn{
  3398. events: [][]byte{
  3399. []byte(`{"type":"response.completed","response":{"id":"resp_turn_prev_recover_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  3400. []byte(`{"type":"error","error":{"type":"invalid_request_error","code":"previous_response_not_found","message":""}}`),
  3401. },
  3402. }
  3403. secondConn := &openAIWSCaptureConn{
  3404. events: [][]byte{
  3405. []byte(`{"type":"response.completed","response":{"id":"resp_turn_prev_recover_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  3406. },
  3407. }
  3408. dialer := &openAIWSQueueDialer{
  3409. conns: []openAIWSClientConn{firstConn, secondConn},
  3410. }
  3411. pool := newOpenAIWSConnPool(cfg)
  3412. pool.setClientDialerForTest(dialer)
  3413. svc := &OpenAIGatewayService{
  3414. cfg: cfg,
  3415. httpUpstream: &httpUpstreamRecorder{},
  3416. cache: &stubGatewayCache{},
  3417. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  3418. toolCorrector: NewCodexToolCorrector(),
  3419. openaiWSPool: pool,
  3420. }
  3421. account := &Account{
  3422. ID: 118,
  3423. Name: "openai-ingress-prev-recovery",
  3424. Platform: PlatformOpenAI,
  3425. Type: AccountTypeAPIKey,
  3426. Status: StatusActive,
  3427. Schedulable: true,
  3428. Concurrency: 1,
  3429. Credentials: map[string]any{
  3430. "api_key": "sk-test",
  3431. },
  3432. Extra: map[string]any{
  3433. "responses_websockets_v2_enabled": true,
  3434. },
  3435. }
  3436. serverErrCh := make(chan error, 1)
  3437. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  3438. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  3439. CompressionMode: coderws.CompressionContextTakeover,
  3440. })
  3441. if err != nil {
  3442. serverErrCh <- err
  3443. return
  3444. }
  3445. defer func() {
  3446. _ = conn.CloseNow()
  3447. }()
  3448. rec := httptest.NewRecorder()
  3449. ginCtx, _ := gin.CreateTestContext(rec)
  3450. req := r.Clone(r.Context())
  3451. req.Header = req.Header.Clone()
  3452. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  3453. ginCtx.Request = req
  3454. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  3455. msgType, firstMessage, readErr := conn.Read(readCtx)
  3456. cancel()
  3457. if readErr != nil {
  3458. serverErrCh <- readErr
  3459. return
  3460. }
  3461. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  3462. serverErrCh <- errors.New("unsupported websocket client message type")
  3463. return
  3464. }
  3465. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  3466. }))
  3467. defer wsServer.Close()
  3468. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  3469. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  3470. cancelDial()
  3471. require.NoError(t, err)
  3472. defer func() {
  3473. _ = clientConn.CloseNow()
  3474. }()
  3475. writeMessage := func(payload string) {
  3476. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3477. defer cancel()
  3478. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  3479. }
  3480. readMessage := func() []byte {
  3481. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3482. defer cancel()
  3483. msgType, message, readErr := clientConn.Read(readCtx)
  3484. require.NoError(t, readErr)
  3485. require.Equal(t, coderws.MessageText, msgType)
  3486. return message
  3487. }
  3488. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_seed_anchor"}`)
  3489. firstTurn := readMessage()
  3490. require.Equal(t, "resp_turn_prev_recover_1", gjson.GetBytes(firstTurn, "response.id").String())
  3491. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_turn_prev_recover_1"}`)
  3492. secondTurn := readMessage()
  3493. require.Equal(t, "response.completed", gjson.GetBytes(secondTurn, "type").String())
  3494. require.Equal(t, "resp_turn_prev_recover_2", gjson.GetBytes(secondTurn, "response.id").String())
  3495. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  3496. select {
  3497. case serverErr := <-serverErrCh:
  3498. require.NoError(t, serverErr)
  3499. case <-time.After(5 * time.Second):
  3500. t.Fatal("等待 ingress websocket 结束超时")
  3501. }
  3502. require.Equal(t, 2, dialer.DialCount(), "previous_response_not_found 恢复应触发换连重试")
  3503. firstConn.mu.Lock()
  3504. firstWrites := append([]map[string]any(nil), firstConn.writes...)
  3505. firstConn.mu.Unlock()
  3506. require.Len(t, firstWrites, 2, "首个连接应处理首轮与失败的第二轮请求")
  3507. require.True(t, gjson.Get(requestToJSONString(firstWrites[1]), "previous_response_id").Exists(), "失败轮次首发请求应包含 previous_response_id")
  3508. secondConn.mu.Lock()
  3509. secondWrites := append([]map[string]any(nil), secondConn.writes...)
  3510. secondConn.mu.Unlock()
  3511. require.Len(t, secondWrites, 1, "恢复重试应在第二个连接发送一次请求")
  3512. require.False(t, gjson.Get(requestToJSONString(secondWrites[0]), "previous_response_id").Exists(), "恢复重试应移除 previous_response_id")
  3513. }
  3514. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_StoreDisabledStrictAffinityPreviousResponseNotFoundLayer2Recovery(t *testing.T) {
  3515. gin.SetMode(gin.TestMode)
  3516. cfg := &config.Config{}
  3517. cfg.Security.URLAllowlist.Enabled = false
  3518. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  3519. cfg.Gateway.OpenAIWS.Enabled = true
  3520. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  3521. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  3522. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  3523. cfg.Gateway.OpenAIWS.IngressPreviousResponseRecoveryEnabled = true
  3524. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  3525. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  3526. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  3527. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  3528. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  3529. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  3530. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  3531. firstConn := &openAIWSCaptureConn{
  3532. events: [][]byte{
  3533. []byte(`{"type":"response.completed","response":{"id":"resp_turn_prev_strict_recover_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  3534. []byte(`{"type":"error","error":{"type":"invalid_request_error","code":"previous_response_not_found","message":"missing strict anchor"}}`),
  3535. },
  3536. }
  3537. secondConn := &openAIWSCaptureConn{
  3538. events: [][]byte{
  3539. []byte(`{"type":"response.completed","response":{"id":"resp_turn_prev_strict_recover_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  3540. },
  3541. }
  3542. dialer := &openAIWSQueueDialer{
  3543. conns: []openAIWSClientConn{firstConn, secondConn},
  3544. }
  3545. pool := newOpenAIWSConnPool(cfg)
  3546. pool.setClientDialerForTest(dialer)
  3547. svc := &OpenAIGatewayService{
  3548. cfg: cfg,
  3549. httpUpstream: &httpUpstreamRecorder{},
  3550. cache: &stubGatewayCache{},
  3551. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  3552. toolCorrector: NewCodexToolCorrector(),
  3553. openaiWSPool: pool,
  3554. }
  3555. account := &Account{
  3556. ID: 122,
  3557. Name: "openai-ingress-prev-strict-layer2",
  3558. Platform: PlatformOpenAI,
  3559. Type: AccountTypeAPIKey,
  3560. Status: StatusActive,
  3561. Schedulable: true,
  3562. Concurrency: 1,
  3563. Credentials: map[string]any{
  3564. "api_key": "sk-test",
  3565. },
  3566. Extra: map[string]any{
  3567. "responses_websockets_v2_enabled": true,
  3568. },
  3569. }
  3570. serverErrCh := make(chan error, 1)
  3571. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  3572. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  3573. CompressionMode: coderws.CompressionContextTakeover,
  3574. })
  3575. if err != nil {
  3576. serverErrCh <- err
  3577. return
  3578. }
  3579. defer func() {
  3580. _ = conn.CloseNow()
  3581. }()
  3582. rec := httptest.NewRecorder()
  3583. ginCtx, _ := gin.CreateTestContext(rec)
  3584. req := r.Clone(r.Context())
  3585. req.Header = req.Header.Clone()
  3586. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  3587. ginCtx.Request = req
  3588. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  3589. msgType, firstMessage, readErr := conn.Read(readCtx)
  3590. cancel()
  3591. if readErr != nil {
  3592. serverErrCh <- readErr
  3593. return
  3594. }
  3595. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  3596. serverErrCh <- errors.New("unsupported websocket client message type")
  3597. return
  3598. }
  3599. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  3600. }))
  3601. defer wsServer.Close()
  3602. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  3603. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  3604. cancelDial()
  3605. require.NoError(t, err)
  3606. defer func() {
  3607. _ = clientConn.CloseNow()
  3608. }()
  3609. writeMessage := func(payload string) {
  3610. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3611. defer cancel()
  3612. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  3613. }
  3614. readMessage := func() []byte {
  3615. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3616. defer cancel()
  3617. msgType, message, readErr := clientConn.Read(readCtx)
  3618. require.NoError(t, readErr)
  3619. require.Equal(t, coderws.MessageText, msgType)
  3620. return message
  3621. }
  3622. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"prompt_cache_key":"pk_strict_layer2","input":[{"type":"input_text","text":"hello"}]}`)
  3623. firstTurn := readMessage()
  3624. require.Equal(t, "resp_turn_prev_strict_recover_1", gjson.GetBytes(firstTurn, "response.id").String())
  3625. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false,"prompt_cache_key":"pk_strict_layer2","previous_response_id":"resp_turn_prev_strict_recover_1","input":[{"type":"input_text","text":"world"}]}`)
  3626. secondTurn := readMessage()
  3627. require.Equal(t, "resp_turn_prev_strict_recover_2", gjson.GetBytes(secondTurn, "response.id").String())
  3628. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  3629. select {
  3630. case serverErr := <-serverErrCh:
  3631. require.NoError(t, serverErr)
  3632. case <-time.After(5 * time.Second):
  3633. t.Fatal("等待 ingress websocket 严格亲和 Layer2 恢复结束超时")
  3634. }
  3635. require.Equal(t, 2, dialer.DialCount(), "严格亲和链路命中 previous_response_not_found 应触发 Layer2 恢复重试")
  3636. firstConn.mu.Lock()
  3637. firstWrites := append([]map[string]any(nil), firstConn.writes...)
  3638. firstConn.mu.Unlock()
  3639. require.Len(t, firstWrites, 2, "首连接应收到首轮请求和失败的续链请求")
  3640. require.True(t, gjson.Get(requestToJSONString(firstWrites[1]), "previous_response_id").Exists())
  3641. secondConn.mu.Lock()
  3642. secondWrites := append([]map[string]any(nil), secondConn.writes...)
  3643. secondConn.mu.Unlock()
  3644. require.Len(t, secondWrites, 1, "Layer2 恢复应仅重放一次")
  3645. secondWrite := requestToJSONString(secondWrites[0])
  3646. require.False(t, gjson.Get(secondWrite, "previous_response_id").Exists(), "Layer2 恢复重放应移除 previous_response_id")
  3647. require.True(t, gjson.Get(secondWrite, "store").Exists(), "Layer2 恢复不应改变 store 标志")
  3648. require.False(t, gjson.Get(secondWrite, "store").Bool())
  3649. require.Equal(t, 2, len(gjson.Get(secondWrite, "input").Array()), "Layer2 恢复应重放完整 input 上下文")
  3650. require.Equal(t, "hello", gjson.Get(secondWrite, "input.0.text").String())
  3651. require.Equal(t, "world", gjson.Get(secondWrite, "input.1.text").String())
  3652. }
  3653. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PreviousResponseNotFoundRecoveryRemovesDuplicatePrevID(t *testing.T) {
  3654. gin.SetMode(gin.TestMode)
  3655. cfg := &config.Config{}
  3656. cfg.Security.URLAllowlist.Enabled = false
  3657. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  3658. cfg.Gateway.OpenAIWS.Enabled = true
  3659. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  3660. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  3661. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  3662. cfg.Gateway.OpenAIWS.IngressPreviousResponseRecoveryEnabled = true
  3663. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  3664. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  3665. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  3666. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  3667. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  3668. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  3669. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  3670. firstConn := &openAIWSCaptureConn{
  3671. events: [][]byte{
  3672. []byte(`{"type":"response.completed","response":{"id":"resp_turn_prev_once_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  3673. []byte(`{"type":"error","error":{"type":"invalid_request_error","code":"previous_response_not_found","message":"first missing"}}`),
  3674. },
  3675. }
  3676. secondConn := &openAIWSCaptureConn{
  3677. events: [][]byte{
  3678. []byte(`{"type":"response.completed","response":{"id":"resp_turn_prev_once_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  3679. },
  3680. }
  3681. dialer := &openAIWSQueueDialer{
  3682. conns: []openAIWSClientConn{firstConn, secondConn},
  3683. }
  3684. pool := newOpenAIWSConnPool(cfg)
  3685. pool.setClientDialerForTest(dialer)
  3686. svc := &OpenAIGatewayService{
  3687. cfg: cfg,
  3688. httpUpstream: &httpUpstreamRecorder{},
  3689. cache: &stubGatewayCache{},
  3690. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  3691. toolCorrector: NewCodexToolCorrector(),
  3692. openaiWSPool: pool,
  3693. }
  3694. account := &Account{
  3695. ID: 120,
  3696. Name: "openai-ingress-prev-recovery-once",
  3697. Platform: PlatformOpenAI,
  3698. Type: AccountTypeAPIKey,
  3699. Status: StatusActive,
  3700. Schedulable: true,
  3701. Concurrency: 1,
  3702. Credentials: map[string]any{
  3703. "api_key": "sk-test",
  3704. },
  3705. Extra: map[string]any{
  3706. "responses_websockets_v2_enabled": true,
  3707. },
  3708. }
  3709. serverErrCh := make(chan error, 1)
  3710. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  3711. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  3712. CompressionMode: coderws.CompressionContextTakeover,
  3713. })
  3714. if err != nil {
  3715. serverErrCh <- err
  3716. return
  3717. }
  3718. defer func() {
  3719. _ = conn.CloseNow()
  3720. }()
  3721. rec := httptest.NewRecorder()
  3722. ginCtx, _ := gin.CreateTestContext(rec)
  3723. req := r.Clone(r.Context())
  3724. req.Header = req.Header.Clone()
  3725. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  3726. ginCtx.Request = req
  3727. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  3728. msgType, firstMessage, readErr := conn.Read(readCtx)
  3729. cancel()
  3730. if readErr != nil {
  3731. serverErrCh <- readErr
  3732. return
  3733. }
  3734. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  3735. serverErrCh <- errors.New("unsupported websocket client message type")
  3736. return
  3737. }
  3738. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  3739. }))
  3740. defer wsServer.Close()
  3741. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  3742. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  3743. cancelDial()
  3744. require.NoError(t, err)
  3745. defer func() {
  3746. _ = clientConn.CloseNow()
  3747. }()
  3748. writeMessage := func(payload string) {
  3749. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3750. defer cancel()
  3751. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  3752. }
  3753. readMessage := func() []byte {
  3754. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  3755. defer cancel()
  3756. msgType, message, readErr := clientConn.Read(readCtx)
  3757. require.NoError(t, readErr)
  3758. require.Equal(t, coderws.MessageText, msgType)
  3759. return message
  3760. }
  3761. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false}`)
  3762. firstTurn := readMessage()
  3763. require.Equal(t, "resp_turn_prev_once_1", gjson.GetBytes(firstTurn, "response.id").String())
  3764. // duplicate previous_response_id: 恢复重试时应删除所有重复键,避免再次 previous_response_not_found。
  3765. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_turn_prev_once_1","input":[],"previous_response_id":"resp_turn_prev_duplicate"}`)
  3766. secondTurn := readMessage()
  3767. require.Equal(t, "resp_turn_prev_once_2", gjson.GetBytes(secondTurn, "response.id").String())
  3768. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  3769. select {
  3770. case serverErr := <-serverErrCh:
  3771. require.NoError(t, serverErr)
  3772. case <-time.After(5 * time.Second):
  3773. t.Fatal("等待 ingress websocket 结束超时")
  3774. }
  3775. require.Equal(t, 2, dialer.DialCount(), "previous_response_not_found 恢复应只重试一次")
  3776. firstConn.mu.Lock()
  3777. firstWrites := append([]map[string]any(nil), firstConn.writes...)
  3778. firstConn.mu.Unlock()
  3779. require.Len(t, firstWrites, 2)
  3780. require.True(t, gjson.Get(requestToJSONString(firstWrites[1]), "previous_response_id").Exists())
  3781. secondConn.mu.Lock()
  3782. secondWrites := append([]map[string]any(nil), secondConn.writes...)
  3783. secondConn.mu.Unlock()
  3784. require.Len(t, secondWrites, 1)
  3785. require.False(t, gjson.Get(requestToJSONString(secondWrites[0]), "previous_response_id").Exists(), "重复键场景恢复重试后不应保留 previous_response_id")
  3786. }
  3787. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_RejectsMessageIDAsPreviousResponseID(t *testing.T) {
  3788. gin.SetMode(gin.TestMode)
  3789. cfg := &config.Config{}
  3790. cfg.Security.URLAllowlist.Enabled = false
  3791. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  3792. cfg.Gateway.OpenAIWS.Enabled = true
  3793. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  3794. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  3795. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  3796. svc := &OpenAIGatewayService{
  3797. cfg: cfg,
  3798. httpUpstream: &httpUpstreamRecorder{},
  3799. cache: &stubGatewayCache{},
  3800. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  3801. toolCorrector: NewCodexToolCorrector(),
  3802. }
  3803. account := &Account{
  3804. ID: 119,
  3805. Name: "openai-ingress-prev-validation",
  3806. Platform: PlatformOpenAI,
  3807. Type: AccountTypeAPIKey,
  3808. Status: StatusActive,
  3809. Schedulable: true,
  3810. Concurrency: 1,
  3811. Credentials: map[string]any{
  3812. "api_key": "sk-test",
  3813. },
  3814. Extra: map[string]any{
  3815. "responses_websockets_v2_enabled": true,
  3816. },
  3817. }
  3818. serverErrCh := make(chan error, 1)
  3819. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  3820. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  3821. CompressionMode: coderws.CompressionContextTakeover,
  3822. })
  3823. if err != nil {
  3824. serverErrCh <- err
  3825. return
  3826. }
  3827. defer func() {
  3828. _ = conn.CloseNow()
  3829. }()
  3830. rec := httptest.NewRecorder()
  3831. ginCtx, _ := gin.CreateTestContext(rec)
  3832. req := r.Clone(r.Context())
  3833. req.Header = req.Header.Clone()
  3834. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  3835. ginCtx.Request = req
  3836. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  3837. msgType, firstMessage, readErr := conn.Read(readCtx)
  3838. cancel()
  3839. if readErr != nil {
  3840. serverErrCh <- readErr
  3841. return
  3842. }
  3843. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  3844. serverErrCh <- errors.New("unsupported websocket client message type")
  3845. return
  3846. }
  3847. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  3848. }))
  3849. defer wsServer.Close()
  3850. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  3851. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  3852. cancelDial()
  3853. require.NoError(t, err)
  3854. defer func() {
  3855. _ = clientConn.CloseNow()
  3856. }()
  3857. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  3858. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"msg_123456"}`))
  3859. cancelWrite()
  3860. require.NoError(t, err)
  3861. select {
  3862. case serverErr := <-serverErrCh:
  3863. require.Error(t, serverErr)
  3864. var closeErr *OpenAIWSClientCloseError
  3865. require.ErrorAs(t, serverErr, &closeErr)
  3866. require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode())
  3867. require.Contains(t, closeErr.Reason(), "previous_response_id must be a response.id")
  3868. case <-time.After(5 * time.Second):
  3869. t.Fatal("等待 ingress websocket 结束超时")
  3870. }
  3871. }
  3872. type openAIWSQueueDialer struct {
  3873. mu sync.Mutex
  3874. conns []openAIWSClientConn
  3875. dialCount int
  3876. }
  3877. func (d *openAIWSQueueDialer) Dial(
  3878. ctx context.Context,
  3879. wsURL string,
  3880. headers http.Header,
  3881. proxyURL string,
  3882. ) (openAIWSClientConn, int, http.Header, error) {
  3883. _ = ctx
  3884. _ = wsURL
  3885. _ = headers
  3886. _ = proxyURL
  3887. d.mu.Lock()
  3888. defer d.mu.Unlock()
  3889. d.dialCount++
  3890. if len(d.conns) == 0 {
  3891. return nil, 503, nil, errors.New("no test conn")
  3892. }
  3893. conn := d.conns[0]
  3894. if len(d.conns) > 1 {
  3895. d.conns = d.conns[1:]
  3896. }
  3897. return conn, 0, nil, nil
  3898. }
  3899. func (d *openAIWSQueueDialer) DialCount() int {
  3900. d.mu.Lock()
  3901. defer d.mu.Unlock()
  3902. return d.dialCount
  3903. }
  3904. type openAIWSPreflightFailConn struct {
  3905. mu sync.Mutex
  3906. events [][]byte
  3907. pingFails bool
  3908. writeCount int
  3909. pingCount int
  3910. }
  3911. func (c *openAIWSPreflightFailConn) WriteJSON(context.Context, any) error {
  3912. c.mu.Lock()
  3913. c.writeCount++
  3914. c.mu.Unlock()
  3915. return nil
  3916. }
  3917. func (c *openAIWSPreflightFailConn) ReadMessage(context.Context) ([]byte, error) {
  3918. c.mu.Lock()
  3919. defer c.mu.Unlock()
  3920. if len(c.events) == 0 {
  3921. return nil, io.EOF
  3922. }
  3923. event := c.events[0]
  3924. c.events = c.events[1:]
  3925. if len(c.events) == 0 {
  3926. c.pingFails = true
  3927. }
  3928. return event, nil
  3929. }
  3930. func (c *openAIWSPreflightFailConn) Ping(context.Context) error {
  3931. c.mu.Lock()
  3932. defer c.mu.Unlock()
  3933. c.pingCount++
  3934. if c.pingFails {
  3935. return errors.New("preflight ping failed")
  3936. }
  3937. return nil
  3938. }
  3939. func (c *openAIWSPreflightFailConn) Close() error {
  3940. return nil
  3941. }
  3942. func (c *openAIWSPreflightFailConn) WriteCount() int {
  3943. c.mu.Lock()
  3944. defer c.mu.Unlock()
  3945. return c.writeCount
  3946. }
  3947. func (c *openAIWSPreflightFailConn) PingCount() int {
  3948. c.mu.Lock()
  3949. defer c.mu.Unlock()
  3950. return c.pingCount
  3951. }
  3952. type openAIWSWriteFailAfterFirstTurnConn struct {
  3953. mu sync.Mutex
  3954. events [][]byte
  3955. failOnWrite bool
  3956. }
  3957. func (c *openAIWSWriteFailAfterFirstTurnConn) WriteJSON(context.Context, any) error {
  3958. c.mu.Lock()
  3959. defer c.mu.Unlock()
  3960. if c.failOnWrite {
  3961. return errors.New("write failed on stale conn")
  3962. }
  3963. return nil
  3964. }
  3965. func (c *openAIWSWriteFailAfterFirstTurnConn) ReadMessage(context.Context) ([]byte, error) {
  3966. c.mu.Lock()
  3967. defer c.mu.Unlock()
  3968. if len(c.events) == 0 {
  3969. return nil, io.EOF
  3970. }
  3971. event := c.events[0]
  3972. c.events = c.events[1:]
  3973. if len(c.events) == 0 {
  3974. c.failOnWrite = true
  3975. }
  3976. return event, nil
  3977. }
  3978. func (c *openAIWSWriteFailAfterFirstTurnConn) Ping(context.Context) error {
  3979. return nil
  3980. }
  3981. func (c *openAIWSWriteFailAfterFirstTurnConn) Close() error {
  3982. return nil
  3983. }
  3984. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_ClientDisconnectStillDrainsUpstream(t *testing.T) {
  3985. gin.SetMode(gin.TestMode)
  3986. cfg := &config.Config{}
  3987. cfg.Security.URLAllowlist.Enabled = false
  3988. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  3989. cfg.Gateway.OpenAIWS.Enabled = true
  3990. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  3991. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  3992. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  3993. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  3994. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  3995. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  3996. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  3997. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  3998. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  3999. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  4000. // 多个上游事件:前几个为非 terminal 事件,最后一个为 terminal。
  4001. // 第一个事件延迟 250ms 让客户端 RST 有时间传播,使 writeClientMessage 可靠失败。
  4002. captureConn := &openAIWSCaptureConn{
  4003. readDelays: []time.Duration{250 * time.Millisecond, 0, 0},
  4004. events: [][]byte{
  4005. []byte(`{"type":"response.created","response":{"id":"resp_ingress_disconnect","model":"gpt-5.1"}}`),
  4006. []byte(`{"type":"response.output_item.added","response":{"id":"resp_ingress_disconnect"}}`),
  4007. []byte(`{"type":"response.completed","response":{"id":"resp_ingress_disconnect","model":"gpt-5.1","usage":{"input_tokens":2,"output_tokens":1}}}`),
  4008. },
  4009. }
  4010. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  4011. pool := newOpenAIWSConnPool(cfg)
  4012. pool.setClientDialerForTest(captureDialer)
  4013. svc := &OpenAIGatewayService{
  4014. cfg: cfg,
  4015. httpUpstream: &httpUpstreamRecorder{},
  4016. cache: &stubGatewayCache{},
  4017. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  4018. toolCorrector: NewCodexToolCorrector(),
  4019. openaiWSPool: pool,
  4020. }
  4021. account := &Account{
  4022. ID: 115,
  4023. Name: "openai-ingress-client-disconnect",
  4024. Platform: PlatformOpenAI,
  4025. Type: AccountTypeAPIKey,
  4026. Status: StatusActive,
  4027. Schedulable: true,
  4028. Concurrency: 1,
  4029. Credentials: map[string]any{
  4030. "api_key": "sk-test",
  4031. "model_mapping": map[string]any{
  4032. "custom-original-model": "gpt-5.1",
  4033. },
  4034. },
  4035. Extra: map[string]any{
  4036. "responses_websockets_v2_enabled": true,
  4037. },
  4038. }
  4039. serverErrCh := make(chan error, 1)
  4040. resultCh := make(chan *OpenAIForwardResult, 1)
  4041. hooks := &OpenAIWSIngressHooks{
  4042. AfterTurn: func(_ int, result *OpenAIForwardResult, turnErr error) {
  4043. if turnErr == nil && result != nil {
  4044. resultCh <- result
  4045. }
  4046. },
  4047. }
  4048. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  4049. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  4050. CompressionMode: coderws.CompressionContextTakeover,
  4051. })
  4052. if err != nil {
  4053. serverErrCh <- err
  4054. return
  4055. }
  4056. defer func() {
  4057. _ = conn.CloseNow()
  4058. }()
  4059. rec := httptest.NewRecorder()
  4060. ginCtx, _ := gin.CreateTestContext(rec)
  4061. req := r.Clone(r.Context())
  4062. req.Header = req.Header.Clone()
  4063. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  4064. ginCtx.Request = req
  4065. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  4066. msgType, firstMessage, readErr := conn.Read(readCtx)
  4067. cancel()
  4068. if readErr != nil {
  4069. serverErrCh <- readErr
  4070. return
  4071. }
  4072. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  4073. serverErrCh <- errors.New("unsupported websocket client message type")
  4074. return
  4075. }
  4076. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, hooks)
  4077. }))
  4078. defer wsServer.Close()
  4079. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  4080. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  4081. cancelDial()
  4082. require.NoError(t, err)
  4083. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  4084. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"custom-original-model","stream":false,"service_tier":"flex"}`))
  4085. cancelWrite()
  4086. require.NoError(t, err)
  4087. // 立即关闭客户端,模拟客户端在 relay 期间断连。
  4088. require.NoError(t, clientConn.CloseNow(), "模拟 ingress 客户端提前断连")
  4089. select {
  4090. case serverErr := <-serverErrCh:
  4091. require.NoError(t, serverErr, "客户端断连后应继续 drain 上游直到 terminal 或正常结束")
  4092. case <-time.After(5 * time.Second):
  4093. t.Fatal("等待 ingress websocket 结束超时")
  4094. }
  4095. select {
  4096. case result := <-resultCh:
  4097. require.Equal(t, "resp_ingress_disconnect", result.RequestID)
  4098. require.Equal(t, 2, result.Usage.InputTokens)
  4099. require.Equal(t, 1, result.Usage.OutputTokens)
  4100. require.NotNil(t, result.ServiceTier)
  4101. require.Equal(t, "flex", *result.ServiceTier)
  4102. case <-time.After(2 * time.Second):
  4103. t.Fatal("未收到断连后的 turn 结果回调")
  4104. }
  4105. }
  4106. func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_InvalidEncryptedContentLineageStripsNextTurn(t *testing.T) {
  4107. gin.SetMode(gin.TestMode)
  4108. cfg := &config.Config{}
  4109. cfg.Security.URLAllowlist.Enabled = false
  4110. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  4111. cfg.Gateway.OpenAIWS.Enabled = true
  4112. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  4113. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  4114. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  4115. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  4116. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  4117. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  4118. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  4119. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  4120. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  4121. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  4122. upstreamConn := &openAIWSCaptureConn{
  4123. events: [][]byte{
  4124. []byte(`{"type":"error","error":{"type":"invalid_request_error","code":"invalid_encrypted_content","message":"The encrypted content could not be verified"}}`),
  4125. []byte(`{"type":"response.failed","response":{"id":"resp_enc_lineage_1","model":"gpt-5.1","error":{"code":"invalid_encrypted_content","message":"The encrypted content could not be verified"}}}`),
  4126. []byte(`{"type":"response.completed","response":{"id":"resp_enc_lineage_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  4127. },
  4128. }
  4129. dialer := &openAIWSQueueDialer{
  4130. conns: []openAIWSClientConn{upstreamConn},
  4131. }
  4132. pool := newOpenAIWSConnPool(cfg)
  4133. pool.setClientDialerForTest(dialer)
  4134. svc := &OpenAIGatewayService{
  4135. cfg: cfg,
  4136. httpUpstream: &httpUpstreamRecorder{},
  4137. cache: &stubGatewayCache{},
  4138. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  4139. toolCorrector: NewCodexToolCorrector(),
  4140. openaiWSPool: pool,
  4141. }
  4142. account := &Account{
  4143. ID: 119,
  4144. Name: "openai-ingress-enc-lineage",
  4145. Platform: PlatformOpenAI,
  4146. Type: AccountTypeAPIKey,
  4147. Status: StatusActive,
  4148. Schedulable: true,
  4149. Concurrency: 1,
  4150. Credentials: map[string]any{
  4151. "api_key": "sk-test",
  4152. },
  4153. Extra: map[string]any{
  4154. "responses_websockets_v2_enabled": true,
  4155. },
  4156. }
  4157. serverErrCh := make(chan error, 1)
  4158. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  4159. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  4160. CompressionMode: coderws.CompressionContextTakeover,
  4161. })
  4162. if err != nil {
  4163. serverErrCh <- err
  4164. return
  4165. }
  4166. defer func() {
  4167. _ = conn.CloseNow()
  4168. }()
  4169. rec := httptest.NewRecorder()
  4170. ginCtx, _ := gin.CreateTestContext(rec)
  4171. req := r.Clone(r.Context())
  4172. req.Header = req.Header.Clone()
  4173. req.Header.Set("User-Agent", "unit-test-agent/1.0")
  4174. ginCtx.Request = req
  4175. readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  4176. msgType, firstMessage, readErr := conn.Read(readCtx)
  4177. cancel()
  4178. if readErr != nil {
  4179. serverErrCh <- readErr
  4180. return
  4181. }
  4182. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  4183. serverErrCh <- errors.New("unsupported websocket client message type")
  4184. return
  4185. }
  4186. serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  4187. }))
  4188. defer wsServer.Close()
  4189. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  4190. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  4191. cancelDial()
  4192. require.NoError(t, err)
  4193. defer func() {
  4194. _ = clientConn.CloseNow()
  4195. }()
  4196. writeMessage := func(payload string) {
  4197. writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  4198. defer cancel()
  4199. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  4200. }
  4201. readMessage := func() []byte {
  4202. readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  4203. defer cancel()
  4204. msgType, message, readErr := clientConn.Read(readCtx)
  4205. require.NoError(t, readErr)
  4206. require.Equal(t, coderws.MessageText, msgType)
  4207. return message
  4208. }
  4209. // turn1:携带失效密文,上游以 invalid_encrypted_content 拒绝(error + response.failed 透传给客户端)。
  4210. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"input":[{"type":"reasoning","id":"rs_1","encrypted_content":"stale-cipher","summary":[]},{"type":"input_text","text":"hi"}]}`)
  4211. firstEvent := readMessage()
  4212. require.Equal(t, "error", gjson.GetBytes(firstEvent, "type").String())
  4213. require.Equal(t, "invalid_encrypted_content", gjson.GetBytes(firstEvent, "error.code").String())
  4214. secondEvent := readMessage()
  4215. require.Equal(t, "response.failed", gjson.GetBytes(secondEvent, "type").String())
  4216. // turn2:客户端历史仍带同一失效密文,进场应被 lineage 预剥离后再发上游。
  4217. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"input":[{"type":"reasoning","id":"rs_1","encrypted_content":"stale-cipher","summary":[]},{"type":"input_text","text":"hi"},{"type":"input_text","text":"again"}]}`)
  4218. thirdEvent := readMessage()
  4219. require.Equal(t, "response.completed", gjson.GetBytes(thirdEvent, "type").String())
  4220. require.Equal(t, "resp_enc_lineage_2", gjson.GetBytes(thirdEvent, "response.id").String())
  4221. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  4222. select {
  4223. case serverErr := <-serverErrCh:
  4224. require.NoError(t, serverErr)
  4225. case <-time.After(5 * time.Second):
  4226. t.Fatal("等待 ingress websocket 结束超时")
  4227. }
  4228. upstreamConn.mu.Lock()
  4229. writes := append([]map[string]any(nil), upstreamConn.writes...)
  4230. upstreamConn.mu.Unlock()
  4231. require.Len(t, writes, 2, "两轮各应发送一次上游请求")
  4232. firstUpstream := requestToJSONString(writes[0])
  4233. require.Equal(t, "stale-cipher", gjson.Get(firstUpstream, "input.0.encrypted_content").String(), "首轮请求原样携带密文")
  4234. secondUpstream := requestToJSONString(writes[1])
  4235. secondInput := gjson.Get(secondUpstream, "input").Array()
  4236. require.Len(t, secondInput, 3, "剥离仅移除 encrypted_content 字段,reasoning 骨架保留")
  4237. for _, item := range secondInput {
  4238. require.False(t, item.Get("encrypted_content").Exists(), "第二轮请求不得再携带已失效密文: %s", item.Raw)
  4239. }
  4240. require.Equal(t, "rs_1", gjson.Get(secondUpstream, "input.0.id").String())
  4241. require.Equal(t, "again", gjson.Get(secondUpstream, "input.2.text").String())
  4242. }