openai_gateway_handler_test.go 114 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205
  1. package handler
  2. import (
  3. "context"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "io"
  8. "net/http"
  9. "net/http/httptest"
  10. "strings"
  11. "sync"
  12. "sync/atomic"
  13. "testing"
  14. "time"
  15. "github.com/Wei-Shaw/sub2api/internal/config"
  16. pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
  17. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  18. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  19. "github.com/Wei-Shaw/sub2api/internal/server/middleware"
  20. "github.com/Wei-Shaw/sub2api/internal/service"
  21. coderws "github.com/coder/websocket"
  22. "github.com/gin-gonic/gin"
  23. "github.com/stretchr/testify/assert"
  24. "github.com/stretchr/testify/require"
  25. "github.com/tidwall/gjson"
  26. "github.com/tidwall/sjson"
  27. )
  28. func TestOpenAIHandleStreamingAwareError_JSONEscaping(t *testing.T) {
  29. tests := []struct {
  30. name string
  31. errType string
  32. message string
  33. }{
  34. {
  35. name: "包含双引号的消息",
  36. errType: "server_error",
  37. message: `upstream returned "invalid" response`,
  38. },
  39. {
  40. name: "包含反斜杠的消息",
  41. errType: "server_error",
  42. message: `path C:\Users\test\file.txt not found`,
  43. },
  44. {
  45. name: "包含双引号和反斜杠的消息",
  46. errType: "upstream_error",
  47. message: `error parsing "key\value": unexpected token`,
  48. },
  49. {
  50. name: "包含换行符的消息",
  51. errType: "server_error",
  52. message: "line1\nline2\ttab",
  53. },
  54. {
  55. name: "普通消息",
  56. errType: "upstream_error",
  57. message: "Upstream service temporarily unavailable",
  58. },
  59. }
  60. for _, tt := range tests {
  61. t.Run(tt.name, func(t *testing.T) {
  62. gin.SetMode(gin.TestMode)
  63. w := httptest.NewRecorder()
  64. c, _ := gin.CreateTestContext(w)
  65. c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
  66. h := &OpenAIGatewayHandler{}
  67. h.handleStreamingAwareError(c, http.StatusBadGateway, tt.errType, tt.message, true)
  68. body := w.Body.String()
  69. // 验证 SSE 格式:event: error\ndata: {JSON}\n\n
  70. assert.True(t, strings.HasPrefix(body, "event: error\n"), "应以 'event: error\\n' 开头")
  71. assert.True(t, strings.HasSuffix(body, "\n\n"), "应以 '\\n\\n' 结尾")
  72. // 提取 data 部分
  73. lines := strings.Split(strings.TrimSuffix(body, "\n\n"), "\n")
  74. require.Len(t, lines, 2, "应有 event 行和 data 行")
  75. dataLine := lines[1]
  76. require.True(t, strings.HasPrefix(dataLine, "data: "), "第二行应以 'data: ' 开头")
  77. jsonStr := strings.TrimPrefix(dataLine, "data: ")
  78. // 验证 JSON 合法性
  79. var parsed map[string]any
  80. err := json.Unmarshal([]byte(jsonStr), &parsed)
  81. require.NoError(t, err, "JSON 应能被成功解析,原始 JSON: %s", jsonStr)
  82. // 验证结构
  83. errorObj, ok := parsed["error"].(map[string]any)
  84. require.True(t, ok, "应包含 error 对象")
  85. assert.Equal(t, tt.errType, errorObj["type"])
  86. assert.Equal(t, tt.message, errorObj["message"])
  87. })
  88. }
  89. }
  90. func TestOpenAIHandleStreamingAwareErrorWithCode_EmitsStableClassification(t *testing.T) {
  91. gin.SetMode(gin.TestMode)
  92. w := httptest.NewRecorder()
  93. c, _ := gin.CreateTestContext(w)
  94. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
  95. h := &OpenAIGatewayHandler{}
  96. h.handleStreamingAwareErrorWithCode(
  97. c,
  98. http.StatusBadGateway,
  99. "upstream_error",
  100. service.OpenAIUpstreamHTTP2StreamErrorCode,
  101. "Upstream HTTP/2 stream failed",
  102. true,
  103. true,
  104. )
  105. body := w.Body.String()
  106. require.Contains(t, body, "event: error\n")
  107. require.Equal(t, "upstream_error", gjson.Get(body[strings.Index(body, "{"):], "error.type").String())
  108. require.Equal(t, service.OpenAIUpstreamHTTP2StreamErrorCode, gjson.Get(body[strings.Index(body, "{"):], "error.code").String())
  109. require.NotContains(t, body, "stream ID")
  110. streamErr, ok := service.GetOpsStreamError(c)
  111. require.True(t, ok)
  112. require.True(t, streamErr.CountTowardsSLA)
  113. require.Equal(t, http.StatusBadGateway, streamErr.IntendedStatus)
  114. }
  115. func TestOpenAIForwardSucceededForScheduling(t *testing.T) {
  116. require.True(t, openAIForwardSucceededForScheduling(nil))
  117. require.True(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{}))
  118. require.True(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{
  119. OpenAIWSMode: true,
  120. UpstreamTerminalEvent: "response.completed",
  121. }))
  122. require.False(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{
  123. OpenAIWSMode: true,
  124. UpstreamTerminalEvent: "response.failed",
  125. }))
  126. }
  127. func TestOpenAIResponsesRequiredCapability(t *testing.T) {
  128. tests := []struct {
  129. name string
  130. imageIntent bool
  131. platform string
  132. want service.OpenAIEndpointCapability
  133. }{
  134. {
  135. name: "OpenAI explicit image intent requires Responses",
  136. imageIntent: true,
  137. platform: service.PlatformOpenAI,
  138. want: service.OpenAIEndpointCapabilityResponses,
  139. },
  140. {
  141. name: "Grok explicit image intent keeps chat capability",
  142. imageIntent: true,
  143. platform: service.PlatformGrok,
  144. want: service.OpenAIEndpointCapabilityChatCompletions,
  145. },
  146. {
  147. name: "non-image intent keeps chat capability",
  148. platform: service.PlatformOpenAI,
  149. want: service.OpenAIEndpointCapabilityChatCompletions,
  150. },
  151. }
  152. for _, tt := range tests {
  153. t.Run(tt.name, func(t *testing.T) {
  154. require.Equal(t, tt.want, openAIResponsesRequiredCapability(tt.imageIntent, tt.platform))
  155. })
  156. }
  157. }
  158. func TestResolveOpenAIMessagesMetadataSession_DoesNotDerivePromptCacheKey(t *testing.T) {
  159. body := []byte(`{"model":"claude-sonnet-4-5","metadata":{"user_id":"claude-code-session"},"messages":[{"role":"user","content":"hello"}]}`)
  160. sessionHash, promptCacheKey := resolveOpenAIMessagesMetadataSession(nil, "", "", "claude-sonnet-4-5", body)
  161. require.NotEmpty(t, sessionHash)
  162. require.Empty(t, promptCacheKey)
  163. }
  164. func TestResolveOpenAIMessagesMetadataSession_PreservesExplicitPromptCacheKey(t *testing.T) {
  165. body := []byte(`{"metadata":{"user_id":"claude-code-session"}}`)
  166. sessionHash, promptCacheKey := resolveOpenAIMessagesMetadataSession(nil, "", "explicit-cache", "claude-sonnet-4-5", body)
  167. require.NotEmpty(t, sessionHash)
  168. require.Equal(t, "explicit-cache", promptCacheKey)
  169. }
  170. func TestResolveOpenAIMessagesMetadataSession_ClaudeCodeHeaderOverridesContentFallback(t *testing.T) {
  171. gin.SetMode(gin.TestMode)
  172. recorder := httptest.NewRecorder()
  173. c, _ := gin.CreateTestContext(recorder)
  174. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  175. c.Request.Header.Set("X-Claude-Code-Session-Id", "claude-session-001")
  176. body1 := []byte(`{"model":"gpt-5.6-sol","system":"parent","messages":[{"role":"user","content":"parent task"}]}`)
  177. body2 := []byte(`{"model":"gpt-5.6-sol","system":"subagent","messages":[{"role":"user","content":"child task"}]}`)
  178. contentHash1 := (&service.OpenAIGatewayService{}).GenerateSessionHash(c, body1)
  179. contentHash2 := (&service.OpenAIGatewayService{}).GenerateSessionHash(c, body2)
  180. require.NotEqual(t, contentHash1, contentHash2, "different bodies should prove the content fallback differs")
  181. hash1, cacheKey1 := resolveOpenAIMessagesMetadataSession(c, contentHash1, "", "gpt-5.6-sol", body1)
  182. hash2, cacheKey2 := resolveOpenAIMessagesMetadataSession(c, contentHash2, "", "gpt-5.6-sol", body2)
  183. require.Equal(t, service.DeriveSessionHashFromSeed("claude-session-001"), hash1)
  184. require.Equal(t, hash1, hash2, "the same Claude Code session must keep one sticky account across changed turn bodies")
  185. require.Empty(t, cacheKey1, "routing-only fix must not create an upstream prompt cache key")
  186. require.Empty(t, cacheKey2, "routing-only fix must not create an upstream prompt cache key")
  187. }
  188. func TestResolveOpenAIMessagesMetadataSession_OpenAISignalWinsOverClaudeHeader(t *testing.T) {
  189. gin.SetMode(gin.TestMode)
  190. recorder := httptest.NewRecorder()
  191. c, _ := gin.CreateTestContext(recorder)
  192. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  193. c.Request.Header.Set("X-Claude-Code-Session-Id", "claude-session-001")
  194. hash, cacheKey := resolveOpenAIMessagesMetadataSession(c, "content-hash", "explicit-openai-session", "gpt-5.6-sol", []byte(`{"metadata":{"user_id":"opaque"}}`))
  195. require.Equal(t, "content-hash", hash, "existing OpenAI session resolution must remain authoritative")
  196. require.Equal(t, "explicit-openai-session", cacheKey)
  197. }
  198. func TestResolveOpenAIMessagesMetadataSession_BlankClaudeHeaderKeepsContentFallback(t *testing.T) {
  199. gin.SetMode(gin.TestMode)
  200. recorder := httptest.NewRecorder()
  201. c, _ := gin.CreateTestContext(recorder)
  202. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  203. c.Request.Header.Set("X-Claude-Code-Session-Id", " ")
  204. hash, cacheKey := resolveOpenAIMessagesMetadataSession(c, "content-hash", "", "gpt-5.6-sol", []byte(`{"metadata":{"user_id":"opaque"}}`))
  205. require.Equal(t, "content-hash", hash)
  206. require.Empty(t, cacheKey)
  207. }
  208. func TestOpenAIHandleStreamingAwareError_NonStreaming(t *testing.T) {
  209. gin.SetMode(gin.TestMode)
  210. w := httptest.NewRecorder()
  211. c, _ := gin.CreateTestContext(w)
  212. c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
  213. h := &OpenAIGatewayHandler{}
  214. h.handleStreamingAwareError(c, http.StatusBadGateway, "upstream_error", "test error", false)
  215. // 非流式应返回 JSON 响应
  216. assert.Equal(t, http.StatusBadGateway, w.Code)
  217. var parsed map[string]any
  218. err := json.Unmarshal(w.Body.Bytes(), &parsed)
  219. require.NoError(t, err)
  220. errorObj, ok := parsed["error"].(map[string]any)
  221. require.True(t, ok)
  222. assert.Equal(t, "upstream_error", errorObj["type"])
  223. assert.Equal(t, "test error", errorObj["message"])
  224. }
  225. func TestReadRequestBodyWithPrealloc(t *testing.T) {
  226. payload := `{"model":"gpt-5","input":"hello"}`
  227. req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(payload))
  228. req.ContentLength = int64(len(payload))
  229. body, err := pkghttputil.ReadRequestBodyWithPrealloc(req)
  230. require.NoError(t, err)
  231. require.Equal(t, payload, string(body))
  232. }
  233. func TestReadRequestBodyWithPrealloc_MaxBytesError(t *testing.T) {
  234. rec := httptest.NewRecorder()
  235. req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(strings.Repeat("x", 8)))
  236. req.Body = http.MaxBytesReader(rec, req.Body, 4)
  237. _, err := pkghttputil.ReadRequestBodyWithPrealloc(req)
  238. require.Error(t, err)
  239. var maxErr *http.MaxBytesError
  240. require.ErrorAs(t, err, &maxErr)
  241. }
  242. func TestOpenAIEnsureForwardErrorResponse_WritesFallbackWhenNotWritten(t *testing.T) {
  243. gin.SetMode(gin.TestMode)
  244. w := httptest.NewRecorder()
  245. c, _ := gin.CreateTestContext(w)
  246. c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
  247. h := &OpenAIGatewayHandler{}
  248. wrote := h.ensureForwardErrorResponse(c, false)
  249. require.True(t, wrote)
  250. require.Equal(t, http.StatusBadGateway, w.Code)
  251. var parsed map[string]any
  252. err := json.Unmarshal(w.Body.Bytes(), &parsed)
  253. require.NoError(t, err)
  254. errorObj, ok := parsed["error"].(map[string]any)
  255. require.True(t, ok)
  256. assert.Equal(t, "upstream_error", errorObj["type"])
  257. assert.Equal(t, "Upstream request failed", errorObj["message"])
  258. }
  259. // Writer 已写后 ensureForwardErrorResponse 必须仍然把错误信息以 SSE
  260. // 形式追加给客户端(streamStarted 强制 true)。
  261. // 这是 case B 修复:旧实现遇到 Writer.Written 直接 return false,
  262. // 客户端只能拿到 silent EOF;Codex CLI 报 "stream closed before response.completed"。
  263. func TestOpenAIEnsureForwardErrorResponse_AppendsSSEAfterWritten(t *testing.T) {
  264. gin.SetMode(gin.TestMode)
  265. w := httptest.NewRecorder()
  266. c, _ := gin.CreateTestContext(w)
  267. c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
  268. c.String(http.StatusTeapot, "already written")
  269. h := &OpenAIGatewayHandler{}
  270. wrote := h.ensureForwardErrorResponse(c, false)
  271. require.True(t, wrote, "must attempt to communicate the failure to the client via SSE")
  272. // 状态码改不了(headers 已 flush),但 body 应该追加 SSE 错误事件。
  273. require.Equal(t, http.StatusTeapot, w.Code)
  274. assert.Contains(t, w.Body.String(), "already written")
  275. // 非 /responses 路径走 legacy event: error 分支。
  276. assert.Contains(t, w.Body.String(), "event: error\n")
  277. }
  278. // case B 回归测试:/responses 路径,Writer 已被写过(模拟 ping flushed),
  279. // ensureForwardErrorResponse 必须发 response.failed,让 Codex 收到合规终止事件。
  280. func TestOpenAIEnsureForwardErrorResponse_ResponsesRouteAfterWrittenEmitsResponseFailed(t *testing.T) {
  281. gin.SetMode(gin.TestMode)
  282. w := httptest.NewRecorder()
  283. c, _ := gin.CreateTestContext(w)
  284. c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
  285. // 模拟 ping 已 flush 的状态:Writer 已写过 1 个字节
  286. _, _ = c.Writer.WriteString(":\n\n")
  287. h := &OpenAIGatewayHandler{}
  288. wrote := h.ensureForwardErrorResponse(c, false)
  289. require.True(t, wrote)
  290. body := w.Body.String()
  291. assert.Contains(t, body, ":\n\n", "earlier ping bytes preserved")
  292. assert.Contains(t, body, "event: response.failed\n", "appended a Responses terminal event")
  293. assert.Contains(t, body, `"type":"response.failed"`)
  294. assert.Contains(t, body, `"code":"upstream_error"`)
  295. assert.Contains(t, body, "Upstream request failed")
  296. }
  297. func TestOpenAIEnsureForwardErrorResponse_AfterDeltaAppendsSingleValidResponseFailed(t *testing.T) {
  298. gin.SetMode(gin.TestMode)
  299. w := httptest.NewRecorder()
  300. c, _ := gin.CreateTestContext(w)
  301. c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
  302. delta := `{"type":"response.output_text.delta","delta":"ok","sequence_number":1}`
  303. _, err := c.Writer.WriteString("event: response.output_text.delta\ndata: " + delta + "\n\n")
  304. require.NoError(t, err)
  305. h := &OpenAIGatewayHandler{}
  306. require.True(t, h.ensureForwardErrorResponse(c, true))
  307. frames := strings.Split(strings.TrimSuffix(w.Body.String(), "\n\n"), "\n\n")
  308. require.Len(t, frames, 2)
  309. errorEvents := 0
  310. for _, frame := range frames {
  311. lines := strings.Split(frame, "\n")
  312. require.Len(t, lines, 2)
  313. require.True(t, strings.HasPrefix(lines[0], "event: "))
  314. require.True(t, strings.HasPrefix(lines[1], "data: "))
  315. eventType := strings.TrimPrefix(lines[0], "event: ")
  316. data := strings.TrimPrefix(lines[1], "data: ")
  317. require.True(t, json.Valid([]byte(data)), "each downstream SSE frame must contain valid JSON")
  318. var event struct {
  319. Type string `json:"type"`
  320. }
  321. require.NoError(t, json.Unmarshal([]byte(data), &event))
  322. require.Equal(t, eventType, event.Type)
  323. if eventType == "response.failed" {
  324. errorEvents++
  325. }
  326. }
  327. require.Equal(t, 1, errorEvents)
  328. }
  329. func TestOpenAIEnsureForwardErrorResponse_CompactKeepaliveOnlyWritesResponseFailed(t *testing.T) {
  330. gin.SetMode(gin.TestMode)
  331. w := httptest.NewRecorder()
  332. c, _ := gin.CreateTestContext(w)
  333. c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
  334. service.MarkOpenAICompactClientStream(c)
  335. stop := service.StartOpenAICompactSSEKeepalive(c, 5*time.Millisecond)
  336. defer stop()
  337. before := service.OpenAICompactKeepaliveAdjustedWrittenSize(c)
  338. require.Eventually(t, c.Writer.Written, time.Second, time.Millisecond)
  339. require.Equal(t, before, service.OpenAICompactKeepaliveAdjustedWrittenSize(c))
  340. h := &OpenAIGatewayHandler{}
  341. require.True(t, h.ensureForwardErrorResponse(c, false))
  342. require.Equal(t, http.StatusOK, w.Code)
  343. require.Contains(t, w.Body.String(), "event: response.failed\n")
  344. require.NotContains(t, w.Body.String(), "event: error\n")
  345. }
  346. func TestOpenAIEnsureForwardErrorResponse_ImageJSONKeepaliveWritesSingleJSONFallback(t *testing.T) {
  347. gin.SetMode(gin.TestMode)
  348. w := httptest.NewRecorder()
  349. c, _ := gin.CreateTestContext(w)
  350. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
  351. stop := service.StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond)
  352. defer stop()
  353. before := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)
  354. require.Eventually(t, c.Writer.Written, time.Second, time.Millisecond)
  355. require.Equal(t, before, service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c))
  356. require.False(t, openAIForwardErrorAlreadyCommunicated(c, before, errors.New("read upstream response: unexpected EOF")))
  357. h := &OpenAIGatewayHandler{}
  358. wrote := h.ensureForwardErrorResponse(c, false)
  359. require.True(t, wrote)
  360. require.Equal(t, http.StatusOK, w.Code, "heartbeat already committed the status")
  361. require.True(t, json.Valid(w.Body.Bytes()), w.Body.String())
  362. require.NotContains(t, w.Body.String(), "event:")
  363. require.NotContains(t, w.Body.String(), "data:")
  364. decoder := json.NewDecoder(strings.NewReader(w.Body.String()))
  365. var payload map[string]any
  366. require.NoError(t, decoder.Decode(&payload))
  367. require.ErrorIs(t, decoder.Decode(&payload), io.EOF)
  368. require.Equal(t, "upstream_error", gjson.Get(w.Body.String(), "error.type").String())
  369. require.Equal(t, "Upstream request failed", gjson.Get(w.Body.String(), "error.message").String())
  370. }
  371. func TestOpenAIEnsureForwardErrorResponse_ImageJSONKeepalivePreservesCompletedJSON(t *testing.T) {
  372. gin.SetMode(gin.TestMode)
  373. w := httptest.NewRecorder()
  374. c, _ := gin.CreateTestContext(w)
  375. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
  376. stop := service.StartOpenAIImagesJSONKeepalive(c, 5*time.Millisecond)
  377. defer stop()
  378. before := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)
  379. require.Eventually(t, c.Writer.Written, time.Second, time.Millisecond)
  380. c.JSON(http.StatusOK, gin.H{"data": []gin.H{{"b64_json": "aW1hZ2U="}}})
  381. completedBody := w.Body.String()
  382. require.True(t, json.Valid([]byte(completedBody)), completedBody)
  383. require.Greater(t, service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c), before)
  384. require.False(t, openAIForwardErrorAlreadyCommunicated(c, before, errors.New("read upstream trailer: unexpected EOF")))
  385. h := &OpenAIGatewayHandler{}
  386. wrote := h.ensureForwardErrorResponse(c, false)
  387. require.False(t, wrote, "the completed Images JSON already communicated the response")
  388. require.Equal(t, completedBody, w.Body.String())
  389. require.NotContains(t, w.Body.String(), "event:")
  390. require.NotContains(t, w.Body.String(), "data:")
  391. decoder := json.NewDecoder(strings.NewReader(w.Body.String()))
  392. var payload map[string]any
  393. require.NoError(t, decoder.Decode(&payload))
  394. require.ErrorIs(t, decoder.Decode(&payload), io.EOF)
  395. require.Equal(t, "aW1hZ2U=", gjson.Get(w.Body.String(), "data.0.b64_json").String())
  396. }
  397. func TestOpenAIEnsureForwardErrorResponse_FastImageJSONKeepalivePreservesCompletedJSON(t *testing.T) {
  398. gin.SetMode(gin.TestMode)
  399. w := httptest.NewRecorder()
  400. c, _ := gin.CreateTestContext(w)
  401. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
  402. stop := service.StartOpenAIImagesJSONKeepalive(c, time.Hour)
  403. defer stop()
  404. before := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)
  405. c.JSON(http.StatusOK, gin.H{"data": []gin.H{{"b64_json": "ZmFzdC1pbWFnZQ=="}}})
  406. completedBody := w.Body.String()
  407. require.True(t, json.Valid([]byte(completedBody)), completedBody)
  408. require.Greater(t, service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c), before)
  409. require.False(t, openAIForwardErrorAlreadyCommunicated(c, before, errors.New("read upstream trailer: unexpected EOF")))
  410. h := &OpenAIGatewayHandler{}
  411. wrote := h.ensureForwardErrorResponse(c, false)
  412. require.False(t, wrote, "fast completed Images JSON already communicated the response")
  413. require.Equal(t, completedBody, w.Body.String())
  414. require.NotContains(t, w.Body.String(), "event:")
  415. require.NotContains(t, w.Body.String(), "data:")
  416. decoder := json.NewDecoder(strings.NewReader(w.Body.String()))
  417. var payload map[string]any
  418. require.NoError(t, decoder.Decode(&payload))
  419. require.ErrorIs(t, decoder.Decode(&payload), io.EOF)
  420. require.Equal(t, "ZmFzdC1pbWFnZQ==", gjson.Get(w.Body.String(), "data.0.b64_json").String())
  421. }
  422. func TestShouldLogOpenAIForwardFailureAsWarn(t *testing.T) {
  423. gin.SetMode(gin.TestMode)
  424. t.Run("fallback_written_should_not_downgrade", func(t *testing.T) {
  425. w := httptest.NewRecorder()
  426. c, _ := gin.CreateTestContext(w)
  427. c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
  428. require.False(t, shouldLogOpenAIForwardFailureAsWarn(c, true))
  429. })
  430. t.Run("context_nil_should_not_downgrade", func(t *testing.T) {
  431. require.False(t, shouldLogOpenAIForwardFailureAsWarn(nil, false))
  432. })
  433. t.Run("response_not_written_should_not_downgrade", func(t *testing.T) {
  434. w := httptest.NewRecorder()
  435. c, _ := gin.CreateTestContext(w)
  436. c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
  437. require.False(t, shouldLogOpenAIForwardFailureAsWarn(c, false))
  438. })
  439. t.Run("response_already_written_should_downgrade", func(t *testing.T) {
  440. w := httptest.NewRecorder()
  441. c, _ := gin.CreateTestContext(w)
  442. c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
  443. c.String(http.StatusForbidden, "already written")
  444. require.True(t, shouldLogOpenAIForwardFailureAsWarn(c, false))
  445. })
  446. }
  447. func TestOpenAIRecoverResponsesPanic_WritesFallbackResponse(t *testing.T) {
  448. gin.SetMode(gin.TestMode)
  449. w := httptest.NewRecorder()
  450. c, _ := gin.CreateTestContext(w)
  451. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  452. h := &OpenAIGatewayHandler{}
  453. streamStarted := false
  454. require.NotPanics(t, func() {
  455. func() {
  456. defer h.recoverResponsesPanic(c, &streamStarted)
  457. panic("test panic")
  458. }()
  459. })
  460. require.Equal(t, http.StatusBadGateway, w.Code)
  461. var parsed map[string]any
  462. err := json.Unmarshal(w.Body.Bytes(), &parsed)
  463. require.NoError(t, err)
  464. errorObj, ok := parsed["error"].(map[string]any)
  465. require.True(t, ok)
  466. assert.Equal(t, "upstream_error", errorObj["type"])
  467. assert.Equal(t, "Upstream request failed", errorObj["message"])
  468. }
  469. func TestOpenAIRecoverResponsesPanic_NoPanicNoWrite(t *testing.T) {
  470. gin.SetMode(gin.TestMode)
  471. w := httptest.NewRecorder()
  472. c, _ := gin.CreateTestContext(w)
  473. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  474. h := &OpenAIGatewayHandler{}
  475. streamStarted := false
  476. require.NotPanics(t, func() {
  477. func() {
  478. defer h.recoverResponsesPanic(c, &streamStarted)
  479. }()
  480. })
  481. require.False(t, c.Writer.Written())
  482. assert.Equal(t, "", w.Body.String())
  483. }
  484. // Panic 在已 flush 的 /v1/responses 流中:状态码无法改(已 written),
  485. // 但 body 应追加 response.failed 让客户端识别为合规截断而不是 silent EOF。
  486. func TestOpenAIRecoverResponsesPanic_AppendsResponseFailedAfterWritten(t *testing.T) {
  487. gin.SetMode(gin.TestMode)
  488. w := httptest.NewRecorder()
  489. c, _ := gin.CreateTestContext(w)
  490. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  491. c.String(http.StatusTeapot, "already written")
  492. h := &OpenAIGatewayHandler{}
  493. streamStarted := false
  494. require.NotPanics(t, func() {
  495. func() {
  496. defer h.recoverResponsesPanic(c, &streamStarted)
  497. panic("test panic")
  498. }()
  499. })
  500. require.Equal(t, http.StatusTeapot, w.Code)
  501. body := w.Body.String()
  502. assert.Contains(t, body, "already written")
  503. assert.Contains(t, body, "event: response.failed\n")
  504. }
  505. func TestOpenAIMissingResponsesDependencies(t *testing.T) {
  506. t.Run("nil_handler", func(t *testing.T) {
  507. var h *OpenAIGatewayHandler
  508. require.Equal(t, []string{"handler"}, h.missingResponsesDependencies())
  509. })
  510. t.Run("all_dependencies_missing", func(t *testing.T) {
  511. h := &OpenAIGatewayHandler{}
  512. require.Equal(t,
  513. []string{"gatewayService", "billingCacheService", "apiKeyService", "concurrencyHelper"},
  514. h.missingResponsesDependencies(),
  515. )
  516. })
  517. t.Run("all_dependencies_present", func(t *testing.T) {
  518. h := &OpenAIGatewayHandler{
  519. gatewayService: &service.OpenAIGatewayService{},
  520. billingCacheService: &service.BillingCacheService{},
  521. apiKeyService: &service.APIKeyService{},
  522. concurrencyHelper: &ConcurrencyHelper{
  523. concurrencyService: &service.ConcurrencyService{},
  524. },
  525. }
  526. require.Empty(t, h.missingResponsesDependencies())
  527. })
  528. }
  529. func TestOpenAIEnsureResponsesDependencies(t *testing.T) {
  530. t.Run("missing_dependencies_returns_503", func(t *testing.T) {
  531. gin.SetMode(gin.TestMode)
  532. w := httptest.NewRecorder()
  533. c, _ := gin.CreateTestContext(w)
  534. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  535. h := &OpenAIGatewayHandler{}
  536. ok := h.ensureResponsesDependencies(c, nil)
  537. require.False(t, ok)
  538. require.Equal(t, http.StatusServiceUnavailable, w.Code)
  539. var parsed map[string]any
  540. err := json.Unmarshal(w.Body.Bytes(), &parsed)
  541. require.NoError(t, err)
  542. errorObj, exists := parsed["error"].(map[string]any)
  543. require.True(t, exists)
  544. assert.Equal(t, "api_error", errorObj["type"])
  545. assert.Equal(t, "Service temporarily unavailable", errorObj["message"])
  546. })
  547. t.Run("already_written_response_not_overridden", func(t *testing.T) {
  548. gin.SetMode(gin.TestMode)
  549. w := httptest.NewRecorder()
  550. c, _ := gin.CreateTestContext(w)
  551. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  552. c.String(http.StatusTeapot, "already written")
  553. h := &OpenAIGatewayHandler{}
  554. ok := h.ensureResponsesDependencies(c, nil)
  555. require.False(t, ok)
  556. require.Equal(t, http.StatusTeapot, w.Code)
  557. assert.Equal(t, "already written", w.Body.String())
  558. })
  559. t.Run("dependencies_ready_returns_true_and_no_write", func(t *testing.T) {
  560. gin.SetMode(gin.TestMode)
  561. w := httptest.NewRecorder()
  562. c, _ := gin.CreateTestContext(w)
  563. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  564. h := &OpenAIGatewayHandler{
  565. gatewayService: &service.OpenAIGatewayService{},
  566. billingCacheService: &service.BillingCacheService{},
  567. apiKeyService: &service.APIKeyService{},
  568. concurrencyHelper: &ConcurrencyHelper{
  569. concurrencyService: &service.ConcurrencyService{},
  570. },
  571. }
  572. ok := h.ensureResponsesDependencies(c, nil)
  573. require.True(t, ok)
  574. require.False(t, c.Writer.Written())
  575. assert.Equal(t, "", w.Body.String())
  576. })
  577. }
  578. func TestResolveOpenAIMessagesDispatchMappedModel(t *testing.T) {
  579. t.Run("exact_claude_model_override_wins", func(t *testing.T) {
  580. apiKey := &service.APIKey{
  581. Group: &service.Group{
  582. MessagesDispatchModelConfig: service.OpenAIMessagesDispatchModelConfig{
  583. SonnetMappedModel: "gpt-5.2",
  584. ExactModelMappings: map[string]string{
  585. "claude-sonnet-4-5-20250929": "gpt-5.4-mini-high",
  586. "claude-fable-5": "gpt-5.6-sol",
  587. },
  588. },
  589. },
  590. }
  591. require.Equal(t, "gpt-5.4-mini", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-sonnet-4-5-20250929"))
  592. require.Equal(t, "gpt-5.6-sol", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-fable-5"))
  593. })
  594. t.Run("uses_family_default_when_no_override", func(t *testing.T) {
  595. apiKey := &service.APIKey{Group: &service.Group{}}
  596. require.Equal(t, "gpt-5.4", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-opus-4-6"))
  597. require.Equal(t, "gpt-5.3-codex", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-sonnet-4-5-20250929"))
  598. require.Equal(t, "gpt-5.4-mini", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-haiku-4-5-20251001"))
  599. })
  600. t.Run("returns_empty_for_non_claude_or_missing_group", func(t *testing.T) {
  601. require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, nil, "claude-sonnet-4-5-20250929"))
  602. require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, &service.APIKey{}, "claude-sonnet-4-5-20250929"))
  603. require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, &service.APIKey{Group: &service.Group{}}, "gpt-5.4"))
  604. })
  605. t.Run("grok_group_maps_claude_cli_model_to_grok_default", func(t *testing.T) {
  606. original := xai.RuntimeModelMappingOptions()
  607. t.Cleanup(func() { xai.SetRuntimeModelMappingOptions(original) })
  608. xai.SetRuntimeModelMappingOptions(xai.ModelMappingOptions{EnableCrossClientMap: true})
  609. apiKey := &service.APIKey{
  610. Group: &service.Group{
  611. Platform: service.PlatformGrok,
  612. },
  613. }
  614. require.Equal(t, "grok-4.6", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-sonnet-4-5"))
  615. require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "grok"))
  616. })
  617. t.Run("does_not_fall_back_to_group_default_mapped_model", func(t *testing.T) {
  618. apiKey := &service.APIKey{
  619. Group: &service.Group{
  620. DefaultMappedModel: "gpt-5.4",
  621. },
  622. }
  623. require.Empty(t, resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "gpt-5.4"))
  624. require.Equal(t, "gpt-5.3-codex", resolveOpenAIMessagesDispatchMappedModel(nil, apiKey, "claude-sonnet-4-5-20250929"))
  625. })
  626. }
  627. func TestOpenAIGatewayMessagesDispatchGateAllowsGrokGroups(t *testing.T) {
  628. gin.SetMode(gin.TestMode)
  629. t.Run("openai_group_without_dispatch_flag_is_rejected", func(t *testing.T) {
  630. rec := httptest.NewRecorder()
  631. c, _ := gin.CreateTestContext(rec)
  632. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hi"}]}`))
  633. groupID := int64(4101)
  634. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  635. ID: 5101,
  636. GroupID: &groupID,
  637. User: &service.User{ID: 6101},
  638. Group: &service.Group{
  639. ID: groupID,
  640. Platform: service.PlatformOpenAI,
  641. AllowMessagesDispatch: false,
  642. },
  643. })
  644. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 6101, Concurrency: 1})
  645. h := &OpenAIGatewayHandler{}
  646. h.Messages(c)
  647. require.Equal(t, http.StatusForbidden, rec.Code)
  648. require.Equal(t, "permission_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
  649. require.Contains(t, rec.Body.String(), "This group does not allow /v1/messages dispatch")
  650. })
  651. t.Run("grok_group_without_dispatch_flag_reaches_gateway_dependencies", func(t *testing.T) {
  652. rec := httptest.NewRecorder()
  653. c, _ := gin.CreateTestContext(rec)
  654. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"grok-4.3","messages":[{"role":"user","content":"hi"}]}`))
  655. groupID := int64(4102)
  656. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  657. ID: 5102,
  658. GroupID: &groupID,
  659. User: &service.User{ID: 6102},
  660. Group: &service.Group{
  661. ID: groupID,
  662. Platform: service.PlatformGrok,
  663. AllowMessagesDispatch: false,
  664. },
  665. })
  666. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 6102, Concurrency: 1})
  667. h := &OpenAIGatewayHandler{}
  668. h.Messages(c)
  669. require.Equal(t, http.StatusServiceUnavailable, rec.Code)
  670. require.Equal(t, "api_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
  671. require.NotContains(t, rec.Body.String(), "This group does not allow /v1/messages dispatch")
  672. })
  673. }
  674. func TestOpenAIModelMappedBody(t *testing.T) {
  675. body := []byte(`{"model":"alias","input":"hello"}`)
  676. calls := 0
  677. forwardBody := openAIModelMappedBody(body, true, "gpt-5.4", func(body []byte, newModel string) []byte {
  678. calls++
  679. return service.ReplaceModelInBody(body, newModel)
  680. })
  681. require.Equal(t, 1, calls)
  682. require.Equal(t, "gpt-5.4", gjson.GetBytes(forwardBody, "model").String())
  683. require.Equal(t, "alias", gjson.GetBytes(body, "model").String())
  684. }
  685. func TestOpenAIModelMappedBodyCache(t *testing.T) {
  686. body := []byte(`{"model":"alias","input":"hello"}`)
  687. calls := 0
  688. mappedBody := newOpenAIModelMappedBodyCache(body, func(body []byte, newModel string) []byte {
  689. calls++
  690. return service.ReplaceModelInBody(body, newModel)
  691. })
  692. first := mappedBody(true, "gpt-5.4")
  693. second := mappedBody(true, "gpt-5.4")
  694. third := mappedBody(true, "gpt-5.3-codex")
  695. unmapped := mappedBody(false, "ignored")
  696. require.Equal(t, 2, calls)
  697. require.Equal(t, "gpt-5.4", gjson.GetBytes(first, "model").String())
  698. require.Equal(t, "gpt-5.4", gjson.GetBytes(second, "model").String())
  699. require.Equal(t, "gpt-5.3-codex", gjson.GetBytes(third, "model").String())
  700. require.Equal(t, body, unmapped)
  701. require.Same(t, &first[0], &second[0])
  702. }
  703. func TestOpenAIResponses_MissingDependencies_ReturnsServiceUnavailable(t *testing.T) {
  704. gin.SetMode(gin.TestMode)
  705. w := httptest.NewRecorder()
  706. c, _ := gin.CreateTestContext(w)
  707. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5","stream":false}`))
  708. c.Request.Header.Set("Content-Type", "application/json")
  709. groupID := int64(2)
  710. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  711. ID: 10,
  712. GroupID: &groupID,
  713. })
  714. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{
  715. UserID: 1,
  716. Concurrency: 1,
  717. })
  718. // 故意使用未初始化依赖,验证快速失败而不是崩溃。
  719. h := &OpenAIGatewayHandler{}
  720. require.NotPanics(t, func() {
  721. h.Responses(c)
  722. })
  723. require.Equal(t, http.StatusServiceUnavailable, w.Code)
  724. var parsed map[string]any
  725. err := json.Unmarshal(w.Body.Bytes(), &parsed)
  726. require.NoError(t, err)
  727. errorObj, ok := parsed["error"].(map[string]any)
  728. require.True(t, ok)
  729. assert.Equal(t, "api_error", errorObj["type"])
  730. assert.Equal(t, "Service temporarily unavailable", errorObj["message"])
  731. }
  732. func TestOpenAIResponses_SetsClientTransportHTTP(t *testing.T) {
  733. gin.SetMode(gin.TestMode)
  734. w := httptest.NewRecorder()
  735. c, _ := gin.CreateTestContext(w)
  736. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(`{"model":"gpt-5"}`))
  737. c.Request.Header.Set("Content-Type", "application/json")
  738. h := &OpenAIGatewayHandler{}
  739. h.Responses(c)
  740. require.Equal(t, http.StatusUnauthorized, w.Code)
  741. require.Equal(t, service.OpenAIClientTransportHTTP, service.GetOpenAIClientTransport(c))
  742. }
  743. func TestOpenAIResponses_RejectsMessageIDAsPreviousResponseID(t *testing.T) {
  744. gin.SetMode(gin.TestMode)
  745. w := httptest.NewRecorder()
  746. c, _ := gin.CreateTestContext(w)
  747. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(
  748. `{"model":"gpt-5.1","stream":false,"previous_response_id":"msg_123456","input":[{"type":"input_text","text":"hello"}]}`,
  749. ))
  750. c.Request.Header.Set("Content-Type", "application/json")
  751. groupID := int64(2)
  752. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  753. ID: 101,
  754. GroupID: &groupID,
  755. User: &service.User{ID: 1},
  756. })
  757. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{
  758. UserID: 1,
  759. Concurrency: 1,
  760. })
  761. h := newOpenAIHandlerForPreviousResponseIDValidation(t, nil)
  762. h.Responses(c)
  763. require.Equal(t, http.StatusBadRequest, w.Code)
  764. require.Contains(t, w.Body.String(), "previous_response_id must be a response.id")
  765. }
  766. func TestOpenAIResponses_AcceptsHTTPContinuationPreviousResponseIDBeforeRouting(t *testing.T) {
  767. gin.SetMode(gin.TestMode)
  768. w := httptest.NewRecorder()
  769. c, _ := gin.CreateTestContext(w)
  770. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(
  771. `{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_123456","input":[{"type":"input_text","text":"hello"}]}`,
  772. ))
  773. c.Request.Header.Set("Content-Type", "application/json")
  774. groupID := int64(2)
  775. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  776. ID: 101,
  777. GroupID: &groupID,
  778. User: &service.User{ID: 1},
  779. })
  780. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{
  781. UserID: 1,
  782. Concurrency: 1,
  783. })
  784. h := newOpenAIHandlerForPreviousResponseIDValidation(t, nil)
  785. require.NoError(t, h.gatewayService.BindOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_123456", 1, 101))
  786. h.Responses(c)
  787. require.NotEqual(t, http.StatusBadRequest, w.Code)
  788. require.NotContains(t, w.Body.String(), "Responses WebSocket v2")
  789. }
  790. func TestOpenAIResponses_RejectsHTTPContinuationOwnedByAnotherUser(t *testing.T) {
  791. gin.SetMode(gin.TestMode)
  792. w := httptest.NewRecorder()
  793. c, _ := gin.CreateTestContext(w)
  794. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(
  795. `{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_other_tenant","input":"hello"}`,
  796. ))
  797. c.Request.Header.Set("Content-Type", "application/json")
  798. groupID := int64(2)
  799. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  800. ID: 202,
  801. UserID: 2,
  802. GroupID: &groupID,
  803. User: &service.User{ID: 2},
  804. })
  805. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 2, Concurrency: 1})
  806. h := newOpenAIHandlerForPreviousResponseIDValidation(t, nil)
  807. require.NoError(t, h.gatewayService.BindOpenAIHTTPResponseOwner(context.Background(), groupID, "resp_other_tenant", 1, 101))
  808. h.Responses(c)
  809. require.Equal(t, http.StatusBadRequest, w.Code)
  810. require.Contains(t, w.Body.String(), "previous_response_id is not available for this user")
  811. }
  812. func TestOpenAIResponses_RejectsUnownedHTTPContinuation(t *testing.T) {
  813. gin.SetMode(gin.TestMode)
  814. w := httptest.NewRecorder()
  815. c, _ := gin.CreateTestContext(w)
  816. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(
  817. `{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_unknown","input":"hello"}`,
  818. ))
  819. c.Request.Header.Set("Content-Type", "application/json")
  820. groupID := int64(2)
  821. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{ID: 101, UserID: 1, GroupID: &groupID, User: &service.User{ID: 1}})
  822. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1, Concurrency: 1})
  823. h := newOpenAIHandlerForPreviousResponseIDValidation(t, nil)
  824. h.Responses(c)
  825. require.Equal(t, http.StatusBadRequest, w.Code)
  826. require.Contains(t, w.Body.String(), "previous_response_id is not available for this user")
  827. }
  828. func TestOpenAIResponses_FunctionCallOutputHTTPGuidanceDoesNotSuggestPreviousResponseReuse(t *testing.T) {
  829. gin.SetMode(gin.TestMode)
  830. w := httptest.NewRecorder()
  831. c, _ := gin.CreateTestContext(w)
  832. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(
  833. `{"model":"gpt-5.1","stream":false,"input":[{"type":"function_call_output","output":"{}"}]}`,
  834. ))
  835. c.Request.Header.Set("Content-Type", "application/json")
  836. groupID := int64(2)
  837. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  838. ID: 101,
  839. GroupID: &groupID,
  840. User: &service.User{ID: 1},
  841. })
  842. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{
  843. UserID: 1,
  844. Concurrency: 1,
  845. })
  846. h := newOpenAIHandlerForPreviousResponseIDValidation(t, nil)
  847. h.Responses(c)
  848. require.Equal(t, http.StatusBadRequest, w.Code)
  849. require.Contains(t, w.Body.String(), "Responses WebSocket v2")
  850. require.NotContains(t, w.Body.String(), "reuse previous_response_id")
  851. }
  852. func TestOpenAIResponsesWebSocket_SetsClientTransportWSWhenUpgradeValid(t *testing.T) {
  853. gin.SetMode(gin.TestMode)
  854. w := httptest.NewRecorder()
  855. c, _ := gin.CreateTestContext(w)
  856. c.Request = httptest.NewRequest(http.MethodGet, "/openai/v1/responses", nil)
  857. c.Request.Header.Set("Upgrade", "websocket")
  858. c.Request.Header.Set("Connection", "Upgrade")
  859. h := &OpenAIGatewayHandler{}
  860. h.ResponsesWebSocket(c)
  861. require.Equal(t, http.StatusUnauthorized, w.Code)
  862. require.Equal(t, service.OpenAIClientTransportWS, service.GetOpenAIClientTransport(c))
  863. }
  864. func TestOpenAIResponsesWebSocket_InvalidUpgradeDoesNotSetTransport(t *testing.T) {
  865. gin.SetMode(gin.TestMode)
  866. w := httptest.NewRecorder()
  867. c, _ := gin.CreateTestContext(w)
  868. c.Request = httptest.NewRequest(http.MethodGet, "/openai/v1/responses", nil)
  869. h := &OpenAIGatewayHandler{}
  870. h.ResponsesWebSocket(c)
  871. require.Equal(t, http.StatusUpgradeRequired, w.Code)
  872. require.Equal(t, service.OpenAIClientTransportUnknown, service.GetOpenAIClientTransport(c))
  873. }
  874. func TestOpenAIResponsesWebSocket_IngressCapacityRejected(t *testing.T) {
  875. gin.SetMode(gin.TestMode)
  876. cache := &concurrencyCacheMock{
  877. acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
  878. return false, nil
  879. },
  880. }
  881. h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
  882. h.cfg = &config.Config{}
  883. h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
  884. wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
  885. defer wsServer.Close()
  886. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  887. clientConn, response, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
  888. cancelDial()
  889. require.Error(t, err)
  890. require.Nil(t, clientConn)
  891. require.NotNil(t, response)
  892. require.Equal(t, http.StatusTooManyRequests, response.StatusCode)
  893. _ = response.Body.Close()
  894. }
  895. func TestOpenAIResponsesWebSocket_IngressLeaseBackendUnavailableBeforeUpgrade(t *testing.T) {
  896. gin.SetMode(gin.TestMode)
  897. cache := &concurrencyCacheMock{
  898. acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
  899. return false, errors.New("redis unavailable")
  900. },
  901. }
  902. h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
  903. h.cfg = &config.Config{}
  904. h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
  905. wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
  906. defer wsServer.Close()
  907. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  908. clientConn, response, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
  909. cancelDial()
  910. require.Error(t, err)
  911. require.Nil(t, clientConn)
  912. require.NotNil(t, response)
  913. require.Equal(t, http.StatusServiceUnavailable, response.StatusCode)
  914. _ = response.Body.Close()
  915. }
  916. func TestOpenAIResponsesWebSocket_FirstMessageTimeoutUsesConfig(t *testing.T) {
  917. gin.SetMode(gin.TestMode)
  918. logSink, restore := captureHandlerStructuredLog(t)
  919. defer restore()
  920. h := newOpenAIHandlerForPreviousResponseIDValidation(t, nil)
  921. h.cfg = &config.Config{}
  922. h.cfg.Gateway.OpenAIWS.ClientFirstMessageTimeoutSeconds = 1
  923. wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
  924. defer wsServer.Close()
  925. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  926. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
  927. cancelDial()
  928. require.NoError(t, err)
  929. defer func() { _ = clientConn.CloseNow() }()
  930. started := time.Now()
  931. readCtx, cancelRead := context.WithTimeout(context.Background(), 5*time.Second)
  932. _, _, err = clientConn.Read(readCtx)
  933. cancelRead()
  934. elapsed := time.Since(started)
  935. require.Error(t, err)
  936. require.NotErrorIs(t, err, context.DeadlineExceeded)
  937. require.GreaterOrEqual(t, elapsed, 500*time.Millisecond)
  938. require.Less(t, elapsed, 4*time.Second)
  939. require.Eventually(t, func() bool {
  940. readTimeout, ok := logSink.FieldValueForMessage("openai.websocket_read_first_message_failed", "read_timeout")
  941. return ok && readTimeout == time.Second &&
  942. logSink.ContainsMessageAtLevel("openai.websocket_read_first_message_failed", "warn")
  943. }, time.Second, 10*time.Millisecond)
  944. }
  945. func TestOpenAIResponsesWebSocket_IngressLeaseReleasedOnEarlyReturn(t *testing.T) {
  946. gin.SetMode(gin.TestMode)
  947. cache := &concurrencyCacheMock{
  948. acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
  949. return true, nil
  950. },
  951. }
  952. h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
  953. h.cfg = &config.Config{}
  954. h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
  955. wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
  956. defer wsServer.Close()
  957. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  958. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
  959. cancelDial()
  960. require.NoError(t, err)
  961. defer func() { _ = clientConn.CloseNow() }()
  962. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  963. err = clientConn.Write(writeCtx, coderws.MessageBinary, []byte("not a response.create frame"))
  964. cancelWrite()
  965. require.NoError(t, err)
  966. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  967. _, _, err = clientConn.Read(readCtx)
  968. cancelRead()
  969. var closeErr coderws.CloseError
  970. require.ErrorAs(t, err, &closeErr)
  971. require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
  972. require.Eventually(t, func() bool {
  973. return atomic.LoadInt32(&cache.releaseIngressCalled) == 1
  974. }, time.Second, 10*time.Millisecond)
  975. }
  976. func TestOpenAIResponsesWebSocket_IngressLeaseReleasedWhenUpgradeFails(t *testing.T) {
  977. gin.SetMode(gin.TestMode)
  978. cache := &concurrencyCacheMock{
  979. acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
  980. return true, nil
  981. },
  982. }
  983. h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
  984. h.cfg = &config.Config{}
  985. h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
  986. wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
  987. defer wsServer.Close()
  988. req, err := http.NewRequest(http.MethodGet, wsServer.URL+"/openai/v1/responses", nil)
  989. require.NoError(t, err)
  990. req.Header.Set("Upgrade", "websocket")
  991. req.Header.Set("Connection", "Upgrade")
  992. req.Header.Set("Sec-WebSocket-Version", "13")
  993. resp, err := http.DefaultClient.Do(req)
  994. require.NoError(t, err)
  995. _ = resp.Body.Close()
  996. require.NotEqual(t, http.StatusSwitchingProtocols, resp.StatusCode)
  997. require.Eventually(t, func() bool {
  998. return atomic.LoadInt32(&cache.releaseIngressCalled) == 1
  999. }, time.Second, 10*time.Millisecond)
  1000. }
  1001. func TestOpenAIResponsesWebSocket_RejectsMessageIDAsPreviousResponseID(t *testing.T) {
  1002. gin.SetMode(gin.TestMode)
  1003. h := newOpenAIHandlerForPreviousResponseIDValidation(t, nil)
  1004. wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
  1005. defer wsServer.Close()
  1006. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1007. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
  1008. cancelDial()
  1009. require.NoError(t, err)
  1010. defer func() {
  1011. _ = clientConn.CloseNow()
  1012. }()
  1013. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  1014. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(
  1015. `{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"msg_abc123"}`,
  1016. ))
  1017. cancelWrite()
  1018. require.NoError(t, err)
  1019. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  1020. _, _, err = clientConn.Read(readCtx)
  1021. cancelRead()
  1022. require.Error(t, err)
  1023. var closeErr coderws.CloseError
  1024. require.ErrorAs(t, err, &closeErr)
  1025. require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
  1026. require.Contains(t, strings.ToLower(closeErr.Reason), "previous_response_id")
  1027. }
  1028. func TestOpenAIResponsesWebSocket_PreviousResponseIDKindLoggedBeforeAcquireFailure(t *testing.T) {
  1029. gin.SetMode(gin.TestMode)
  1030. cache := &concurrencyCacheMock{
  1031. acquireUserSlotFn: func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
  1032. return false, errors.New("user slot unavailable")
  1033. },
  1034. }
  1035. h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
  1036. wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
  1037. defer wsServer.Close()
  1038. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1039. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
  1040. cancelDial()
  1041. require.NoError(t, err)
  1042. defer func() {
  1043. _ = clientConn.CloseNow()
  1044. }()
  1045. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  1046. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(
  1047. `{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_prev_123"}`,
  1048. ))
  1049. cancelWrite()
  1050. require.NoError(t, err)
  1051. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  1052. _, _, err = clientConn.Read(readCtx)
  1053. cancelRead()
  1054. require.Error(t, err)
  1055. var closeErr coderws.CloseError
  1056. require.ErrorAs(t, err, &closeErr)
  1057. require.Equal(t, coderws.StatusInternalError, closeErr.Code)
  1058. require.Contains(t, strings.ToLower(closeErr.Reason), "failed to acquire user concurrency slot")
  1059. }
  1060. type contentModerationHandlerSettingRepo struct {
  1061. values map[string]string
  1062. }
  1063. func (r *contentModerationHandlerSettingRepo) Get(ctx context.Context, key string) (*service.Setting, error) {
  1064. if value, ok := r.values[key]; ok {
  1065. return &service.Setting{Key: key, Value: value}, nil
  1066. }
  1067. return nil, service.ErrSettingNotFound
  1068. }
  1069. func (r *contentModerationHandlerSettingRepo) GetValue(ctx context.Context, key string) (string, error) {
  1070. if value, ok := r.values[key]; ok {
  1071. return value, nil
  1072. }
  1073. return "", service.ErrSettingNotFound
  1074. }
  1075. func (r *contentModerationHandlerSettingRepo) Set(ctx context.Context, key, value string) error {
  1076. if r.values == nil {
  1077. r.values = map[string]string{}
  1078. }
  1079. r.values[key] = value
  1080. return nil
  1081. }
  1082. func (r *contentModerationHandlerSettingRepo) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
  1083. out := map[string]string{}
  1084. for _, key := range keys {
  1085. if value, ok := r.values[key]; ok {
  1086. out[key] = value
  1087. }
  1088. }
  1089. return out, nil
  1090. }
  1091. func (r *contentModerationHandlerSettingRepo) SetMultiple(ctx context.Context, settings map[string]string) error {
  1092. if r.values == nil {
  1093. r.values = map[string]string{}
  1094. }
  1095. for key, value := range settings {
  1096. r.values[key] = value
  1097. }
  1098. return nil
  1099. }
  1100. func (r *contentModerationHandlerSettingRepo) GetAll(ctx context.Context) (map[string]string, error) {
  1101. out := make(map[string]string, len(r.values))
  1102. for key, value := range r.values {
  1103. out[key] = value
  1104. }
  1105. return out, nil
  1106. }
  1107. func (r *contentModerationHandlerSettingRepo) Delete(ctx context.Context, key string) error {
  1108. delete(r.values, key)
  1109. return nil
  1110. }
  1111. type contentModerationHandlerTestRepo struct {
  1112. mu sync.Mutex
  1113. logs []service.ContentModerationLog
  1114. }
  1115. func (r *contentModerationHandlerTestRepo) CreateLog(ctx context.Context, log *service.ContentModerationLog) error {
  1116. if log != nil {
  1117. r.mu.Lock()
  1118. defer r.mu.Unlock()
  1119. r.logs = append(r.logs, *log)
  1120. }
  1121. return nil
  1122. }
  1123. func (r *contentModerationHandlerTestRepo) resetLogs() {
  1124. r.mu.Lock()
  1125. defer r.mu.Unlock()
  1126. r.logs = nil
  1127. }
  1128. func (r *contentModerationHandlerTestRepo) logSnapshot() []service.ContentModerationLog {
  1129. r.mu.Lock()
  1130. defer r.mu.Unlock()
  1131. return append([]service.ContentModerationLog(nil), r.logs...)
  1132. }
  1133. func (r *contentModerationHandlerTestRepo) ListLogs(ctx context.Context, filter service.ContentModerationLogFilter) ([]service.ContentModerationLog, *pagination.PaginationResult, error) {
  1134. return nil, nil, nil
  1135. }
  1136. func (r *contentModerationHandlerTestRepo) CountFlaggedByUserSince(ctx context.Context, userID int64, since time.Time, excludeCyberPolicy bool) (int, error) {
  1137. return 0, nil
  1138. }
  1139. func (r *contentModerationHandlerTestRepo) CleanupExpiredLogs(ctx context.Context, hitBefore time.Time, nonHitBefore time.Time) (*service.ContentModerationCleanupResult, error) {
  1140. return &service.ContentModerationCleanupResult{}, nil
  1141. }
  1142. func (r *contentModerationHandlerTestRepo) UpdateLogEmailSent(ctx context.Context, id int64, sent bool) error {
  1143. return nil
  1144. }
  1145. func TestOpenAIResponsesWebSocket_ContentModerationBlocksFirstFrame(t *testing.T) {
  1146. gin.SetMode(gin.TestMode)
  1147. moderationServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1148. require.Equal(t, "/v1/moderations", r.URL.Path)
  1149. _, _ = w.Write([]byte(`{"results":[{"category_scores":{"sexual":0.9}}]}`))
  1150. }))
  1151. defer moderationServer.Close()
  1152. cfg := &service.ContentModerationConfig{
  1153. Enabled: true,
  1154. Mode: service.ContentModerationModePreBlock,
  1155. BaseURL: moderationServer.URL,
  1156. Model: "omni-moderation-latest",
  1157. APIKeys: []string{"sk-test"},
  1158. SampleRate: 100,
  1159. AllGroups: true,
  1160. BlockMessage: "内容审计测试阻断",
  1161. }
  1162. rawCfg, err := json.Marshal(cfg)
  1163. require.NoError(t, err)
  1164. repo := &contentModerationHandlerTestRepo{}
  1165. settingRepo := &contentModerationHandlerSettingRepo{values: map[string]string{
  1166. service.SettingKeyRiskControlEnabled: "true",
  1167. service.SettingKeyContentModerationConfig: string(rawCfg),
  1168. }}
  1169. moderationSvc := service.NewContentModerationService(
  1170. settingRepo,
  1171. repo,
  1172. nil,
  1173. nil,
  1174. nil,
  1175. nil,
  1176. nil,
  1177. nil,
  1178. )
  1179. decision, err := moderationSvc.Check(context.Background(), service.ContentModerationCheckInput{
  1180. UserID: 1,
  1181. Endpoint: "/v1/responses",
  1182. Provider: "openai",
  1183. Model: "gpt-5.5",
  1184. Protocol: service.ContentModerationProtocolOpenAIResponses,
  1185. Body: []byte(`{"model":"gpt-5.5","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"bad prompt"}]}]}`),
  1186. })
  1187. require.NoError(t, err)
  1188. require.True(t, decision.Blocked)
  1189. require.Eventually(t, func() bool {
  1190. return len(repo.logSnapshot()) == 1
  1191. }, time.Second, 10*time.Millisecond)
  1192. repo.resetLogs()
  1193. h := &OpenAIGatewayHandler{
  1194. gatewayService: &service.OpenAIGatewayService{},
  1195. billingCacheService: &service.BillingCacheService{},
  1196. apiKeyService: &service.APIKeyService{},
  1197. contentModerationService: moderationSvc,
  1198. concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&concurrencyCacheMock{}), SSEPingFormatNone, time.Second),
  1199. }
  1200. wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
  1201. defer wsServer.Close()
  1202. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1203. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
  1204. cancelDial()
  1205. require.NoError(t, err)
  1206. defer func() {
  1207. _ = clientConn.CloseNow()
  1208. }()
  1209. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  1210. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{
  1211. "type":"response.create",
  1212. "model":"gpt-5.5",
  1213. "input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"bad prompt"}]}]
  1214. }`))
  1215. cancelWrite()
  1216. require.NoError(t, err)
  1217. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  1218. _, payload, readErr := clientConn.Read(readCtx)
  1219. cancelRead()
  1220. if readErr == nil {
  1221. require.Contains(t, string(payload), "content_policy_violation")
  1222. require.Contains(t, string(payload), "内容审计测试阻断")
  1223. } else {
  1224. var closeErr coderws.CloseError
  1225. require.ErrorAs(t, readErr, &closeErr)
  1226. require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
  1227. require.Contains(t, closeErr.Reason, "内容审计测试阻断")
  1228. }
  1229. var logs []service.ContentModerationLog
  1230. require.Eventually(t, func() bool {
  1231. logs = repo.logSnapshot()
  1232. return len(logs) == 1
  1233. }, time.Second, 10*time.Millisecond)
  1234. require.True(t, logs[0].Flagged)
  1235. require.Equal(t, service.ContentModerationActionBlock, logs[0].Action)
  1236. require.Equal(t, "bad prompt", logs[0].InputExcerpt)
  1237. }
  1238. func TestOpenAIResponsesWebSocket_PassthroughUsageLogPersistsUserAgentAndReasoningEffort(t *testing.T) {
  1239. got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
  1240. firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false,"reasoning":{"effort":"HIGH"}}`,
  1241. userAgent: testStringPtr("codex_cli_rs/0.125.0 test"),
  1242. })
  1243. require.NotNil(t, got.log.UserAgent)
  1244. require.Equal(t, "codex_cli_rs/0.125.0 test", *got.log.UserAgent)
  1245. require.NotNil(t, got.log.ReasoningEffort)
  1246. require.Equal(t, "high", *got.log.ReasoningEffort)
  1247. require.True(t, got.log.OpenAIWSMode)
  1248. }
  1249. func TestOpenAIResponsesWebSocket_PassthroughUsageLogInfersReasoningFromInitialRequestModel(t *testing.T) {
  1250. got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
  1251. firstPayload: `{"type":"response.create","model":"gpt-5.4-xhigh","stream":false}`,
  1252. userAgent: testStringPtr("codex_cli_rs/0.125.0 mapped"),
  1253. channelMapping: map[string]string{
  1254. "gpt-5.4-xhigh": "gpt-5.4",
  1255. },
  1256. })
  1257. require.Equal(t, "gpt-5.4", gjson.GetBytes(got.upstreamFirstPayload, "model").String(),
  1258. "上游首帧应使用渠道映射后的模型")
  1259. require.NotNil(t, got.log.ReasoningEffort)
  1260. require.Equal(t, "xhigh", *got.log.ReasoningEffort,
  1261. "usage log reasoning effort 必须使用渠道映射前首帧模型后缀推导")
  1262. }
  1263. func TestOpenAIResponsesWebSocket_PassthroughUsageLogLeavesUserAgentNilWhenMissing(t *testing.T) {
  1264. got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
  1265. firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false,"reasoning":{"effort":"medium"}}`,
  1266. userAgent: testStringPtr(""),
  1267. })
  1268. require.Nil(t, got.log.UserAgent, "空入站 User-Agent 不应由上游握手 UA 或默认 UA 兜底")
  1269. require.NotNil(t, got.log.ReasoningEffort)
  1270. require.Equal(t, "medium", *got.log.ReasoningEffort)
  1271. }
  1272. func TestOpenAIResponsesWebSocket_PassthroughTracksModelPerTurn(t *testing.T) {
  1273. got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
  1274. firstPayload: `{"type":"response.create","model":"sol","stream":false}`,
  1275. secondPayload: `{"type":"response.create","model":"terra","stream":false}`,
  1276. channelMapping: map[string]string{
  1277. "sol": "sol-channel",
  1278. "terra": "terra-channel",
  1279. },
  1280. accountModelMapping: map[string]any{
  1281. "sol": "gpt-5.6-sol",
  1282. "terra": "gpt-5.6-terra",
  1283. "sol-channel": "gpt-5.6-sol",
  1284. "terra-channel": "gpt-5.6-terra",
  1285. },
  1286. })
  1287. require.Len(t, got.upstreamPayloads, 2)
  1288. require.Equal(t, "sol-channel", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
  1289. require.Equal(t, "terra-channel", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
  1290. require.Len(t, got.clientEvents, 2)
  1291. require.Equal(t, "sol", gjson.GetBytes(got.clientEvents[0], "response.model").String())
  1292. require.Equal(t, "terra", gjson.GetBytes(got.clientEvents[1], "response.model").String())
  1293. require.Len(t, got.logs, 2)
  1294. require.Equal(t, "sol", got.logs[0].Model)
  1295. require.Equal(t, "sol", got.logs[0].RequestedModel)
  1296. require.NotNil(t, got.logs[0].UpstreamModel)
  1297. require.Equal(t, "sol-channel", *got.logs[0].UpstreamModel)
  1298. require.NotNil(t, got.logs[0].ModelMappingChain)
  1299. require.Equal(t, "sol→sol-channel", *got.logs[0].ModelMappingChain)
  1300. require.Equal(t, "terra", got.logs[1].Model)
  1301. require.Equal(t, "terra", got.logs[1].RequestedModel)
  1302. require.NotNil(t, got.logs[1].UpstreamModel)
  1303. require.Equal(t, "terra-channel", *got.logs[1].UpstreamModel)
  1304. require.NotNil(t, got.logs[1].ModelMappingChain)
  1305. require.Equal(t, "terra→terra-channel", *got.logs[1].ModelMappingChain)
  1306. require.InDelta(t, got.logs[1].TotalCost*2.5, got.logs[0].TotalCost, 1e-12,
  1307. "each turn must be billed with its own channel-mapped model")
  1308. }
  1309. func TestOpenAIResponsesWebSocket_ChannelMappedTargetSelectsAccountWithoutRequestedAlias(t *testing.T) {
  1310. got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
  1311. firstPayload: `{"type":"response.create","model":"public-alias","stream":false}`,
  1312. secondPayload: `{"type":"response.create","stream":false}`,
  1313. channelMapping: map[string]string{
  1314. "public-alias": "gpt-5.6-sol",
  1315. },
  1316. accountModelMapping: map[string]any{
  1317. "gpt-5.6-sol": "gpt-5.6-sol",
  1318. },
  1319. })
  1320. require.Len(t, got.upstreamPayloads, 2)
  1321. require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
  1322. require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
  1323. require.Len(t, got.clientEvents, 2)
  1324. require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[0], "response.model").String())
  1325. require.Equal(t, "public-alias", gjson.GetBytes(got.clientEvents[1], "response.model").String())
  1326. require.Len(t, got.logs, 2)
  1327. for _, usageLog := range got.logs {
  1328. require.Equal(t, "public-alias", usageLog.RequestedModel)
  1329. require.NotNil(t, usageLog.UpstreamModel)
  1330. require.Equal(t, "gpt-5.6-sol", *usageLog.UpstreamModel)
  1331. require.NotNil(t, usageLog.ModelMappingChain)
  1332. require.Equal(t, "public-alias→gpt-5.6-sol", *usageLog.ModelMappingChain)
  1333. }
  1334. }
  1335. func TestOpenAIResponsesWebSocket_PassthroughKeepsTurnMappingSnapshot(t *testing.T) {
  1336. got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
  1337. firstPayload: `{"type":"response.create","model":"sol","stream":false}`,
  1338. secondPayload: `{"type":"response.create","model":"sol","stream":false}`,
  1339. channelMapping: map[string]string{
  1340. "sol": "gpt-5.6-sol",
  1341. },
  1342. afterFirstUpstreamRequest: func(channelSvc *service.ChannelService) error {
  1343. if channelSvc == nil {
  1344. return errors.New("channel service is nil")
  1345. }
  1346. _, err := channelSvc.Update(context.Background(), 7701, &service.UpdateChannelInput{
  1347. ModelMapping: map[string]map[string]string{
  1348. service.PlatformOpenAI: {"sol": "gpt-5.6-terra"},
  1349. },
  1350. })
  1351. return err
  1352. },
  1353. })
  1354. require.Len(t, got.upstreamPayloads, 2)
  1355. require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
  1356. require.Equal(t, "gpt-5.6-terra", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
  1357. require.Len(t, got.logs, 2)
  1358. require.Equal(t, "sol", got.logs[0].Model)
  1359. require.NotNil(t, got.logs[0].UpstreamModel)
  1360. require.Equal(t, "gpt-5.6-sol", *got.logs[0].UpstreamModel)
  1361. require.NotNil(t, got.logs[0].ModelMappingChain)
  1362. require.Equal(t, "sol→gpt-5.6-sol", *got.logs[0].ModelMappingChain)
  1363. require.InDelta(t, 40e-6, got.logs[0].TotalCost, 1e-12,
  1364. "the in-flight turn must retain the channel-mapped billing model used when it was sent")
  1365. require.Equal(t, "sol", got.logs[1].Model)
  1366. require.NotNil(t, got.logs[1].UpstreamModel)
  1367. require.Equal(t, "gpt-5.6-terra", *got.logs[1].UpstreamModel)
  1368. require.NotNil(t, got.logs[1].ModelMappingChain)
  1369. require.Equal(t, "sol→gpt-5.6-terra", *got.logs[1].ModelMappingChain)
  1370. require.InDelta(t, got.logs[1].TotalCost*2.5, got.logs[0].TotalCost, 1e-12,
  1371. "the next turn must use the updated channel mapping")
  1372. }
  1373. func TestOpenAIResponsesWebSocket_CtxPoolAppliesPerTurnMappingAndPreservesRequestedModel(t *testing.T) {
  1374. got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
  1375. firstPayload: `{"type":"response.create","model":"gpt-5.6-sol","stream":false}`,
  1376. secondPayload: `{"type":"response.create","model":"gpt-5.6-terra","stream":false}`,
  1377. ingressMode: service.OpenAIWSIngressModeCtxPool,
  1378. billingModelSource: service.BillingModelSourceRequested,
  1379. channelMapping: map[string]string{
  1380. "gpt-5.6-terra": "gpt-5.6-sol",
  1381. },
  1382. })
  1383. require.Len(t, got.upstreamPayloads, 2)
  1384. require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[0], "model").String())
  1385. require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.upstreamPayloads[1], "model").String())
  1386. require.Len(t, got.clientEvents, 2)
  1387. require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(got.clientEvents[0], "response.model").String())
  1388. require.Equal(t, "gpt-5.6-terra", gjson.GetBytes(got.clientEvents[1], "response.model").String())
  1389. require.Len(t, got.logs, 2)
  1390. require.Equal(t, "gpt-5.6-sol", got.logs[0].RequestedModel)
  1391. require.Nil(t, got.logs[0].ModelMappingChain)
  1392. require.InDelta(t, 40e-6, got.logs[0].TotalCost, 1e-12)
  1393. require.Equal(t, "gpt-5.6-terra", got.logs[1].RequestedModel)
  1394. require.NotNil(t, got.logs[1].ModelMappingChain)
  1395. require.Equal(t, "gpt-5.6-terra→gpt-5.6-sol", *got.logs[1].ModelMappingChain)
  1396. require.InDelta(t, 16e-6, got.logs[1].TotalCost, 1e-12,
  1397. "BillingModelSourceRequested must use the client model before channel mapping")
  1398. }
  1399. func TestOpenAIWSTurnBillingModelPreservesImagePricingModel(t *testing.T) {
  1400. tests := []struct {
  1401. name string
  1402. resultModel string
  1403. mapping service.ChannelMappingResult
  1404. requestedModel string
  1405. upstreamModel string
  1406. wantBillingModel string
  1407. }{
  1408. {
  1409. name: "upstream billing preserves image model",
  1410. resultModel: "gpt-image-2",
  1411. mapping: service.ChannelMappingResult{BillingModelSource: service.BillingModelSourceUpstream},
  1412. requestedModel: "gpt-5.6-sol",
  1413. upstreamModel: "gpt-5.6-sol",
  1414. wantBillingModel: "gpt-image-2",
  1415. },
  1416. {
  1417. name: "unmapped channel preserves image model",
  1418. resultModel: "gpt-image-2",
  1419. mapping: service.ChannelMappingResult{MappedModel: "gpt-5.6-sol", BillingModelSource: service.BillingModelSourceChannelMapped},
  1420. requestedModel: "gpt-5.6-sol",
  1421. upstreamModel: "gpt-5.6-sol",
  1422. wantBillingModel: "gpt-image-2",
  1423. },
  1424. {
  1425. name: "requested source overrides image model",
  1426. resultModel: "gpt-image-2",
  1427. mapping: service.ChannelMappingResult{BillingModelSource: service.BillingModelSourceRequested},
  1428. requestedModel: "public-image-alias",
  1429. upstreamModel: "gpt-5.6-sol",
  1430. wantBillingModel: "public-image-alias",
  1431. },
  1432. {
  1433. name: "mapped channel source overrides image model",
  1434. resultModel: "gpt-image-2",
  1435. mapping: service.ChannelMappingResult{MappedModel: "priced-channel-model", BillingModelSource: service.BillingModelSourceChannelMapped},
  1436. requestedModel: "public-image-alias",
  1437. upstreamModel: "gpt-5.6-sol",
  1438. wantBillingModel: "priced-channel-model",
  1439. },
  1440. {
  1441. name: "text turn falls back to upstream model",
  1442. mapping: service.ChannelMappingResult{BillingModelSource: service.BillingModelSourceUpstream},
  1443. requestedModel: "public-alias",
  1444. upstreamModel: "gpt-5.6-sol",
  1445. wantBillingModel: "gpt-5.6-sol",
  1446. },
  1447. }
  1448. for _, tt := range tests {
  1449. t.Run(tt.name, func(t *testing.T) {
  1450. result := &service.OpenAIForwardResult{BillingModel: tt.resultModel}
  1451. require.Equal(t, tt.wantBillingModel, openAIWSTurnBillingModel(result, tt.mapping, tt.requestedModel, tt.upstreamModel))
  1452. })
  1453. }
  1454. }
  1455. func TestOpenAIAccountScheduleModelUsesActualOrSharedResolver(t *testing.T) {
  1456. account := &service.Account{
  1457. Platform: service.PlatformOpenAI,
  1458. Type: service.AccountTypeOAuth,
  1459. Credentials: map[string]any{
  1460. "model_mapping": map[string]any{"public": "billing"},
  1461. "compact_model_mapping": map[string]any{"public": "compact-actual"},
  1462. },
  1463. }
  1464. reported := &service.OpenAIForwardResult{UpstreamModel: "observed-actual"}
  1465. require.Equal(t, "observed-actual", openAIAccountScheduleModel(nil, account, "public", true, reported))
  1466. require.Equal(t, "compact-actual", openAIAccountScheduleModel(nil, account, "public", true, nil))
  1467. require.Equal(t, "billing", openAIAccountScheduleModel(nil, account, "public", false, nil))
  1468. c, _ := gin.CreateTestContext(nil)
  1469. service.SetOpsUpstreamModel(c, "attempt-actual")
  1470. require.Equal(t, "attempt-actual", openAIAccountScheduleModel(c, account, "public", true, nil))
  1471. setOpsSelectedAccount(c, account.ID, account.Platform)
  1472. require.Equal(t, "attempt-actual", openAIAccountScheduleModel(c, account, "public", true, nil))
  1473. }
  1474. func TestOpenAIChannelForwardModelForScheduler(t *testing.T) {
  1475. tests := []struct {
  1476. name string
  1477. mapping service.ChannelMappingResult
  1478. requested string
  1479. want string
  1480. }{
  1481. {
  1482. name: "mapped model is used for capability filtering",
  1483. mapping: service.ChannelMappingResult{Mapped: true, MappedModel: " gpt-forward "},
  1484. requested: "client-alias",
  1485. want: "gpt-forward",
  1486. },
  1487. {
  1488. name: "unmapped request preserves client model",
  1489. mapping: service.ChannelMappingResult{MappedModel: "client-model"},
  1490. requested: "client-model",
  1491. want: "client-model",
  1492. },
  1493. {
  1494. name: "empty mapped model safely preserves client model",
  1495. mapping: service.ChannelMappingResult{Mapped: true, MappedModel: " "},
  1496. requested: "client-model",
  1497. want: "client-model",
  1498. },
  1499. }
  1500. for _, tt := range tests {
  1501. t.Run(tt.name, func(t *testing.T) {
  1502. require.Equal(t, tt.want, openAIChannelForwardModel(tt.mapping, tt.requested))
  1503. })
  1504. }
  1505. }
  1506. func TestShouldReportOpenAIWSProxyAccountFailure(t *testing.T) {
  1507. t.Run("unsupported client model switch does not penalize account", func(t *testing.T) {
  1508. err := fmt.Errorf("wrapped ingress turn: %w", newOpenAIWSUnsupportedModelSwitchError("gpt-unsupported"))
  1509. require.False(t, shouldReportOpenAIWSProxyAccountFailure(err))
  1510. var closeErr *service.OpenAIWSClientCloseError
  1511. require.ErrorAs(t, err, &closeErr)
  1512. require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode())
  1513. require.Equal(t, "model switch requires reconnect", closeErr.Reason())
  1514. })
  1515. t.Run("upstream policy violation still penalizes account", func(t *testing.T) {
  1516. err := service.NewOpenAIWSClientCloseError(
  1517. coderws.StatusPolicyViolation,
  1518. "upstream websocket authentication failed",
  1519. errors.New("upstream rejected credentials"),
  1520. )
  1521. require.True(t, shouldReportOpenAIWSProxyAccountFailure(err))
  1522. })
  1523. t.Run("generic proxy failure still penalizes account", func(t *testing.T) {
  1524. require.True(t, shouldReportOpenAIWSProxyAccountFailure(errors.New("upstream websocket read failed")))
  1525. })
  1526. }
  1527. func TestSetOpenAIClientTransportHTTP(t *testing.T) {
  1528. gin.SetMode(gin.TestMode)
  1529. w := httptest.NewRecorder()
  1530. c, _ := gin.CreateTestContext(w)
  1531. setOpenAIClientTransportHTTP(c)
  1532. require.Equal(t, service.OpenAIClientTransportHTTP, service.GetOpenAIClientTransport(c))
  1533. }
  1534. func TestSetOpenAIClientTransportWS(t *testing.T) {
  1535. gin.SetMode(gin.TestMode)
  1536. w := httptest.NewRecorder()
  1537. c, _ := gin.CreateTestContext(w)
  1538. setOpenAIClientTransportWS(c)
  1539. require.Equal(t, service.OpenAIClientTransportWS, service.GetOpenAIClientTransport(c))
  1540. }
  1541. // TestOpenAIHandler_GjsonExtraction 验证 gjson 从请求体中提取 model/stream 的正确性
  1542. func TestOpenAIHandler_GjsonExtraction(t *testing.T) {
  1543. tests := []struct {
  1544. name string
  1545. body string
  1546. wantModel string
  1547. wantStream bool
  1548. }{
  1549. {"正常提取", `{"model":"gpt-4","stream":true,"input":"hello"}`, "gpt-4", true},
  1550. {"stream false", `{"model":"gpt-4","stream":false}`, "gpt-4", false},
  1551. {"无 stream 字段", `{"model":"gpt-4"}`, "gpt-4", false},
  1552. {"model 缺失", `{"stream":true}`, "", true},
  1553. }
  1554. for _, tt := range tests {
  1555. t.Run(tt.name, func(t *testing.T) {
  1556. body := []byte(tt.body)
  1557. modelResult := gjson.GetBytes(body, "model")
  1558. model := ""
  1559. if modelResult.Type == gjson.String {
  1560. model = modelResult.String()
  1561. }
  1562. stream := gjson.GetBytes(body, "stream").Bool()
  1563. require.Equal(t, tt.wantModel, model)
  1564. require.Equal(t, tt.wantStream, stream)
  1565. })
  1566. }
  1567. }
  1568. // TestOpenAIHandler_GjsonValidation 验证修复后的 JSON 合法性和类型校验
  1569. func TestOpenAIHandler_GjsonValidation(t *testing.T) {
  1570. // 非法 JSON 被 gjson.ValidBytes 拦截
  1571. require.False(t, gjson.ValidBytes([]byte(`{invalid json`)))
  1572. // model 为数字 → 类型不是 gjson.String,应被拒绝
  1573. body := []byte(`{"model":123}`)
  1574. modelResult := gjson.GetBytes(body, "model")
  1575. require.True(t, modelResult.Exists())
  1576. require.NotEqual(t, gjson.String, modelResult.Type)
  1577. // model 为 null → 类型不是 gjson.String,应被拒绝
  1578. body2 := []byte(`{"model":null}`)
  1579. modelResult2 := gjson.GetBytes(body2, "model")
  1580. require.True(t, modelResult2.Exists())
  1581. require.NotEqual(t, gjson.String, modelResult2.Type)
  1582. // stream 为 string → 类型既不是 True 也不是 False,应被拒绝
  1583. body3 := []byte(`{"model":"gpt-4","stream":"true"}`)
  1584. streamResult := gjson.GetBytes(body3, "stream")
  1585. require.True(t, streamResult.Exists())
  1586. require.NotEqual(t, gjson.True, streamResult.Type)
  1587. require.NotEqual(t, gjson.False, streamResult.Type)
  1588. // stream 为 int → 同上
  1589. body4 := []byte(`{"model":"gpt-4","stream":1}`)
  1590. streamResult2 := gjson.GetBytes(body4, "stream")
  1591. require.True(t, streamResult2.Exists())
  1592. require.NotEqual(t, gjson.True, streamResult2.Type)
  1593. require.NotEqual(t, gjson.False, streamResult2.Type)
  1594. }
  1595. // TestOpenAIHandler_InstructionsInjection 验证 instructions 的 gjson/sjson 注入逻辑
  1596. func TestOpenAIHandler_InstructionsInjection(t *testing.T) {
  1597. // 测试 1:无 instructions → 注入
  1598. body := []byte(`{"model":"gpt-4"}`)
  1599. existing := gjson.GetBytes(body, "instructions").String()
  1600. require.Empty(t, existing)
  1601. newBody, err := sjson.SetBytes(body, "instructions", "test instruction")
  1602. require.NoError(t, err)
  1603. require.Equal(t, "test instruction", gjson.GetBytes(newBody, "instructions").String())
  1604. // 测试 2:已有 instructions → 不覆盖
  1605. body2 := []byte(`{"model":"gpt-4","instructions":"existing"}`)
  1606. existing2 := gjson.GetBytes(body2, "instructions").String()
  1607. require.Equal(t, "existing", existing2)
  1608. // 测试 3:空白 instructions → 注入
  1609. body3 := []byte(`{"model":"gpt-4","instructions":" "}`)
  1610. existing3 := strings.TrimSpace(gjson.GetBytes(body3, "instructions").String())
  1611. require.Empty(t, existing3)
  1612. // 测试 4:sjson.SetBytes 返回错误时不应 panic
  1613. // 正常 JSON 不会产生 sjson 错误,验证返回值被正确处理
  1614. validBody := []byte(`{"model":"gpt-4"}`)
  1615. result, setErr := sjson.SetBytes(validBody, "instructions", "hello")
  1616. require.NoError(t, setErr)
  1617. require.True(t, gjson.ValidBytes(result))
  1618. }
  1619. func newOpenAIHandlerForPreviousResponseIDValidation(t *testing.T, cache *concurrencyCacheMock) *OpenAIGatewayHandler {
  1620. t.Helper()
  1621. if cache == nil {
  1622. cache = &concurrencyCacheMock{
  1623. acquireUserSlotFn: func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
  1624. return true, nil
  1625. },
  1626. acquireAccountSlotFn: func(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
  1627. return true, nil
  1628. },
  1629. }
  1630. }
  1631. return &OpenAIGatewayHandler{
  1632. gatewayService: &service.OpenAIGatewayService{},
  1633. billingCacheService: &service.BillingCacheService{},
  1634. apiKeyService: &service.APIKeyService{},
  1635. concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(cache), SSEPingFormatNone, time.Second),
  1636. }
  1637. }
  1638. func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject middleware.AuthSubject) *httptest.Server {
  1639. t.Helper()
  1640. groupID := int64(2)
  1641. apiKey := &service.APIKey{
  1642. ID: 101,
  1643. GroupID: &groupID,
  1644. User: &service.User{ID: subject.UserID},
  1645. }
  1646. router := gin.New()
  1647. router.Use(func(c *gin.Context) {
  1648. c.Set(string(middleware.ContextKeyAPIKey), apiKey)
  1649. c.Set(string(middleware.ContextKeyUser), subject)
  1650. c.Next()
  1651. })
  1652. router.GET("/openai/v1/responses", h.ResponsesWebSocket)
  1653. return httptest.NewServer(router)
  1654. }
  1655. type openAIResponsesWSUsageLogCase struct {
  1656. firstPayload string
  1657. // midPayload 在首个 turn 完成后发送(如 session.update),上游桩会为它
  1658. // 回一个 response.completed,客户端按普通事件读取。
  1659. midPayload string
  1660. secondPayload string
  1661. userAgent *string
  1662. ingressMode string
  1663. channelMapping map[string]string
  1664. billingModelSource string
  1665. accountModelMapping map[string]any
  1666. afterFirstUpstreamRequest func(channelSvc *service.ChannelService) error
  1667. // group 覆盖 apiKey.Group(分组级模型白名单测试用);nil 保持原有无分组行为。
  1668. group *service.Group
  1669. // firstFrameCloseExpected:首帧即被拒(连接被 1008 关闭),不期待任何响应帧。
  1670. firstFrameCloseExpected bool
  1671. // secondTurnCloseExpected:第二个 turn 被拒(连接被 1008 关闭)。
  1672. secondTurnCloseExpected bool
  1673. }
  1674. type openAIResponsesWSUsageLogResult struct {
  1675. log *service.UsageLog
  1676. logs []*service.UsageLog
  1677. upstreamFirstPayload []byte
  1678. upstreamPayloads [][]byte
  1679. clientEvents [][]byte
  1680. }
  1681. type openAIWSUsageHandlerAccountRepoStub struct {
  1682. service.AccountRepository
  1683. account service.Account
  1684. }
  1685. func (s *openAIWSUsageHandlerAccountRepoStub) ListSchedulableByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
  1686. if s.account.Platform != platform {
  1687. return nil, nil
  1688. }
  1689. return []service.Account{s.account}, nil
  1690. }
  1691. func (s *openAIWSUsageHandlerAccountRepoStub) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]service.Account, error) {
  1692. return s.ListSchedulableByPlatform(ctx, platform)
  1693. }
  1694. func (s *openAIWSUsageHandlerAccountRepoStub) GetByID(ctx context.Context, id int64) (*service.Account, error) {
  1695. if s.account.ID != id {
  1696. return nil, nil
  1697. }
  1698. account := s.account
  1699. return &account, nil
  1700. }
  1701. type openAIWSFailoverHandlerAccountRepoStub struct {
  1702. service.AccountRepository
  1703. accounts []service.Account
  1704. rateLimitedIDs []int64
  1705. }
  1706. type openAIHTTPPassthroughFailoverUpstream struct {
  1707. service.HTTPUpstream
  1708. mu sync.Mutex
  1709. accountIDs []int64
  1710. }
  1711. func (u *openAIHTTPPassthroughFailoverUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
  1712. u.mu.Lock()
  1713. u.accountIDs = append(u.accountIDs, accountID)
  1714. u.mu.Unlock()
  1715. return &http.Response{
  1716. StatusCode: http.StatusBadGateway,
  1717. Header: http.Header{"Content-Type": []string{"application/json"}},
  1718. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"temporary upstream failure"}}`)),
  1719. }, nil
  1720. }
  1721. func (u *openAIHTTPPassthroughFailoverUpstream) calls() []int64 {
  1722. u.mu.Lock()
  1723. defer u.mu.Unlock()
  1724. return append([]int64(nil), u.accountIDs...)
  1725. }
  1726. type openAIHTTPPassthroughAuthFailoverUpstream struct {
  1727. service.HTTPUpstream
  1728. mu sync.Mutex
  1729. accountIDs []int64
  1730. statusCode int
  1731. }
  1732. func (u *openAIHTTPPassthroughAuthFailoverUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
  1733. u.mu.Lock()
  1734. u.accountIDs = append(u.accountIDs, accountID)
  1735. u.mu.Unlock()
  1736. if accountID == 9911 {
  1737. return &http.Response{
  1738. StatusCode: http.StatusOK,
  1739. Header: http.Header{"Content-Type": []string{"application/json"}},
  1740. Body: io.NopCloser(strings.NewReader(`{"id":"resp_healthy","object":"response","model":"gpt-5.2","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`)),
  1741. }, nil
  1742. }
  1743. return &http.Response{
  1744. StatusCode: u.statusCode,
  1745. Header: http.Header{"Content-Type": []string{"application/json"}},
  1746. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"upstream credential rejected"}}`)),
  1747. }, nil
  1748. }
  1749. func (u *openAIHTTPPassthroughAuthFailoverUpstream) calls() []int64 {
  1750. u.mu.Lock()
  1751. defer u.mu.Unlock()
  1752. return append([]int64(nil), u.accountIDs...)
  1753. }
  1754. type openAIHTTPPassthroughSSERateLimitUpstream struct {
  1755. service.HTTPUpstream
  1756. mu sync.Mutex
  1757. accountIDs []int64
  1758. }
  1759. func (u *openAIHTTPPassthroughSSERateLimitUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
  1760. u.mu.Lock()
  1761. u.accountIDs = append(u.accountIDs, accountID)
  1762. u.mu.Unlock()
  1763. body := strings.Join([]string{
  1764. "event: response.created",
  1765. `data: {"type":"response.created","response":{"id":"resp_rate_limited"}}`,
  1766. "",
  1767. "event: response.failed",
  1768. `data: {"type":"response.failed","response":{"id":"resp_rate_limited","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
  1769. "",
  1770. }, "\n")
  1771. return &http.Response{
  1772. StatusCode: http.StatusOK,
  1773. Header: http.Header{
  1774. "Content-Type": []string{"text/event-stream"},
  1775. "Retry-After": []string{"1"},
  1776. },
  1777. Body: io.NopCloser(strings.NewReader(body)),
  1778. }, nil
  1779. }
  1780. func (u *openAIHTTPPassthroughSSERateLimitUpstream) calls() []int64 {
  1781. u.mu.Lock()
  1782. defer u.mu.Unlock()
  1783. return append([]int64(nil), u.accountIDs...)
  1784. }
  1785. func (s *openAIWSFailoverHandlerAccountRepoStub) ListSchedulableByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
  1786. out := make([]service.Account, 0, len(s.accounts))
  1787. for _, account := range s.accounts {
  1788. if account.Platform == platform && account.IsSchedulable() {
  1789. out = append(out, account)
  1790. }
  1791. }
  1792. return out, nil
  1793. }
  1794. func (s *openAIWSFailoverHandlerAccountRepoStub) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]service.Account, error) {
  1795. return s.ListSchedulableByPlatform(ctx, platform)
  1796. }
  1797. func (s *openAIWSFailoverHandlerAccountRepoStub) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
  1798. return s.ListSchedulableByPlatform(ctx, platform)
  1799. }
  1800. func (s *openAIWSFailoverHandlerAccountRepoStub) GetByID(ctx context.Context, id int64) (*service.Account, error) {
  1801. for _, account := range s.accounts {
  1802. if account.ID == id {
  1803. acc := account
  1804. return &acc, nil
  1805. }
  1806. }
  1807. return nil, nil
  1808. }
  1809. func (s *openAIWSFailoverHandlerAccountRepoStub) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
  1810. s.rateLimitedIDs = append(s.rateLimitedIDs, id)
  1811. for i := range s.accounts {
  1812. if s.accounts[i].ID == id {
  1813. reset := resetAt
  1814. s.accounts[i].RateLimitResetAt = &reset
  1815. break
  1816. }
  1817. }
  1818. return nil
  1819. }
  1820. type openAIWSUsageHandlerUsageLogRepoStub struct {
  1821. service.UsageLogRepository
  1822. created chan *service.UsageLog
  1823. }
  1824. func (s *openAIWSUsageHandlerUsageLogRepoStub) Create(ctx context.Context, log *service.UsageLog) (bool, error) {
  1825. if s.created != nil {
  1826. s.created <- log
  1827. }
  1828. return true, nil
  1829. }
  1830. type openAIWSUsageHandlerChannelRepoStub struct {
  1831. service.ChannelRepository
  1832. mu sync.Mutex
  1833. channels []service.Channel
  1834. groupPlatforms map[int64]string
  1835. }
  1836. func (s *openAIWSUsageHandlerChannelRepoStub) ListAll(ctx context.Context) ([]service.Channel, error) {
  1837. s.mu.Lock()
  1838. defer s.mu.Unlock()
  1839. out := make([]service.Channel, 0, len(s.channels))
  1840. for i := range s.channels {
  1841. out = append(out, *s.channels[i].Clone())
  1842. }
  1843. return out, nil
  1844. }
  1845. func (s *openAIWSUsageHandlerChannelRepoStub) GetByID(ctx context.Context, id int64) (*service.Channel, error) {
  1846. s.mu.Lock()
  1847. defer s.mu.Unlock()
  1848. for i := range s.channels {
  1849. if s.channels[i].ID == id {
  1850. return s.channels[i].Clone(), nil
  1851. }
  1852. }
  1853. return nil, service.ErrChannelNotFound
  1854. }
  1855. func (s *openAIWSUsageHandlerChannelRepoStub) Update(ctx context.Context, channel *service.Channel) error {
  1856. s.mu.Lock()
  1857. defer s.mu.Unlock()
  1858. for i := range s.channels {
  1859. if s.channels[i].ID == channel.ID {
  1860. s.channels[i] = *channel.Clone()
  1861. return nil
  1862. }
  1863. }
  1864. return service.ErrChannelNotFound
  1865. }
  1866. func (s *openAIWSUsageHandlerChannelRepoStub) GetGroupIDs(ctx context.Context, channelID int64) ([]int64, error) {
  1867. channel, err := s.GetByID(ctx, channelID)
  1868. if err != nil {
  1869. return nil, err
  1870. }
  1871. return append([]int64(nil), channel.GroupIDs...), nil
  1872. }
  1873. func (s *openAIWSUsageHandlerChannelRepoStub) GetGroupPlatforms(ctx context.Context, groupIDs []int64) (map[int64]string, error) {
  1874. out := make(map[int64]string, len(groupIDs))
  1875. for _, groupID := range groupIDs {
  1876. if platform := strings.TrimSpace(s.groupPlatforms[groupID]); platform != "" {
  1877. out[groupID] = platform
  1878. }
  1879. }
  1880. return out, nil
  1881. }
  1882. func TestOpenAIResponses_APIKeyPassthroughPool5xxRetriesThenExhaustsMaxSwitches(t *testing.T) {
  1883. gin.SetMode(gin.TestMode)
  1884. groupID := int64(4203)
  1885. accounts := []service.Account{
  1886. {
  1887. ID: 9910, Name: "pool-api-key", Platform: service.PlatformOpenAI,
  1888. Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 1,
  1889. Credentials: map[string]any{
  1890. "api_key": "sk-pool",
  1891. "base_url": "https://api.example.test",
  1892. "pool_mode": true,
  1893. "pool_mode_retry_count": float64(1),
  1894. "pool_mode_retry_status_codes": []any{float64(http.StatusBadGateway)},
  1895. },
  1896. Extra: map[string]any{"openai_passthrough": true},
  1897. },
  1898. {
  1899. ID: 9911, Name: "fallback-api-key", Platform: service.PlatformOpenAI,
  1900. Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 2,
  1901. Credentials: map[string]any{
  1902. "api_key": "sk-fallback",
  1903. "base_url": "https://api.example.test",
  1904. },
  1905. Extra: map[string]any{"openai_passthrough": true},
  1906. },
  1907. }
  1908. cfg := &config.Config{RunMode: config.RunModeSimple}
  1909. cfg.Default.RateMultiplier = 1
  1910. cfg.Security.URLAllowlist.Enabled = false
  1911. cfg.Gateway.MaxAccountSwitches = 1
  1912. accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
  1913. upstream := &openAIHTTPPassthroughFailoverUpstream{}
  1914. billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
  1915. t.Cleanup(billingCacheSvc.Stop)
  1916. gatewaySvc := service.NewOpenAIGatewayService(
  1917. accountRepo,
  1918. nil,
  1919. nil,
  1920. nil,
  1921. nil,
  1922. nil,
  1923. nil,
  1924. cfg,
  1925. nil,
  1926. nil,
  1927. service.NewBillingService(cfg, nil),
  1928. nil,
  1929. billingCacheSvc,
  1930. upstream,
  1931. &service.DeferredService{},
  1932. nil,
  1933. nil,
  1934. nil,
  1935. nil,
  1936. nil,
  1937. nil,
  1938. nil,
  1939. )
  1940. h := NewOpenAIGatewayHandler(
  1941. gatewaySvc,
  1942. service.NewConcurrencyService(nil),
  1943. billingCacheSvc,
  1944. service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
  1945. nil,
  1946. nil,
  1947. nil,
  1948. nil,
  1949. cfg,
  1950. )
  1951. rec := httptest.NewRecorder()
  1952. c, _ := gin.CreateTestContext(rec)
  1953. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(`{"model":"gpt-5.2","input":"hello","stream":false}`))
  1954. c.Request.Header.Set("Content-Type", "application/json")
  1955. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  1956. ID: 1803, GroupID: &groupID,
  1957. User: &service.User{ID: 1703, Status: service.StatusActive},
  1958. Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
  1959. })
  1960. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1703, Concurrency: 0})
  1961. h.Responses(c)
  1962. require.Equal(t, []int64{9910, 9910, 9911}, upstream.calls())
  1963. require.Equal(t, http.StatusBadGateway, rec.Code)
  1964. require.Equal(t, "upstream_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
  1965. require.Equal(t, "Upstream service temporarily unavailable", gjson.GetBytes(rec.Body.Bytes(), "error.message").String())
  1966. }
  1967. func TestOpenAIResponses_APIKeyPassthroughPoolAuthFailureRetriesThenSwitchesToHealthyAccount(t *testing.T) {
  1968. tests := []struct {
  1969. name string
  1970. statusCode int
  1971. }{
  1972. {name: "401", statusCode: http.StatusUnauthorized},
  1973. {name: "403", statusCode: http.StatusForbidden},
  1974. }
  1975. for _, tt := range tests {
  1976. t.Run(tt.name, func(t *testing.T) {
  1977. gin.SetMode(gin.TestMode)
  1978. groupID := int64(4203)
  1979. accounts := []service.Account{
  1980. {
  1981. ID: 9910, Name: "pool-api-key", Platform: service.PlatformOpenAI,
  1982. Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 1,
  1983. Credentials: map[string]any{
  1984. "api_key": "sk-pool",
  1985. "base_url": "https://api.example.test",
  1986. "pool_mode": true,
  1987. "pool_mode_retry_count": float64(1),
  1988. "pool_mode_retry_status_codes": []any{float64(tt.statusCode)},
  1989. },
  1990. Extra: map[string]any{"openai_passthrough": true},
  1991. },
  1992. {
  1993. ID: 9911, Name: "fallback-api-key", Platform: service.PlatformOpenAI,
  1994. Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 2,
  1995. Credentials: map[string]any{
  1996. "api_key": "sk-fallback",
  1997. "base_url": "https://api.example.test",
  1998. },
  1999. Extra: map[string]any{"openai_passthrough": true},
  2000. },
  2001. }
  2002. cfg := &config.Config{RunMode: config.RunModeSimple}
  2003. cfg.Default.RateMultiplier = 1
  2004. cfg.Security.URLAllowlist.Enabled = false
  2005. cfg.Gateway.MaxAccountSwitches = 1
  2006. accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
  2007. upstream := &openAIHTTPPassthroughAuthFailoverUpstream{statusCode: tt.statusCode}
  2008. rateLimitSvc := service.NewRateLimitService(accountRepo, nil, cfg, nil, nil)
  2009. billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
  2010. t.Cleanup(billingCacheSvc.Stop)
  2011. gatewaySvc := service.NewOpenAIGatewayService(
  2012. accountRepo,
  2013. nil,
  2014. nil,
  2015. nil,
  2016. nil,
  2017. nil,
  2018. nil,
  2019. cfg,
  2020. nil,
  2021. nil,
  2022. service.NewBillingService(cfg, nil),
  2023. rateLimitSvc,
  2024. billingCacheSvc,
  2025. upstream,
  2026. &service.DeferredService{},
  2027. nil,
  2028. nil,
  2029. nil,
  2030. nil,
  2031. nil,
  2032. nil,
  2033. nil,
  2034. )
  2035. h := NewOpenAIGatewayHandler(
  2036. gatewaySvc,
  2037. service.NewConcurrencyService(nil),
  2038. billingCacheSvc,
  2039. service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
  2040. nil,
  2041. nil,
  2042. nil,
  2043. nil,
  2044. cfg,
  2045. )
  2046. rec := httptest.NewRecorder()
  2047. c, _ := gin.CreateTestContext(rec)
  2048. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(`{"model":"gpt-5.2","input":"hello","stream":false}`))
  2049. c.Request.Header.Set("Content-Type", "application/json")
  2050. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  2051. ID: 1803, GroupID: &groupID,
  2052. User: &service.User{ID: 1703, Status: service.StatusActive},
  2053. Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
  2054. })
  2055. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1703, Concurrency: 0})
  2056. h.Responses(c)
  2057. require.Equal(t, []int64{9910, 9910, 9911}, upstream.calls())
  2058. require.Equal(t, http.StatusOK, rec.Code)
  2059. require.Equal(t, "resp_healthy", gjson.GetBytes(rec.Body.Bytes(), "id").String())
  2060. })
  2061. }
  2062. }
  2063. func TestOpenAIResponses_APIKeyPassthroughSSERateLimitUsesConfiguredPoolRetry(t *testing.T) {
  2064. gin.SetMode(gin.TestMode)
  2065. groupID := int64(4204)
  2066. accounts := []service.Account{
  2067. {
  2068. ID: 9912, Name: "pool-sse-rate-limit", Platform: service.PlatformOpenAI,
  2069. Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Priority: 1,
  2070. Credentials: map[string]any{
  2071. "api_key": "sk-pool",
  2072. "base_url": "https://api.example.test",
  2073. "pool_mode": true,
  2074. "pool_mode_retry_count": float64(1),
  2075. "pool_mode_retry_status_codes": []any{float64(http.StatusTooManyRequests)},
  2076. },
  2077. Extra: map[string]any{"openai_passthrough": true},
  2078. },
  2079. }
  2080. cfg := &config.Config{RunMode: config.RunModeSimple}
  2081. cfg.Default.RateMultiplier = 1
  2082. cfg.Security.URLAllowlist.Enabled = false
  2083. cfg.Gateway.MaxAccountSwitches = 1
  2084. accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
  2085. upstream := &openAIHTTPPassthroughSSERateLimitUpstream{}
  2086. billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
  2087. t.Cleanup(billingCacheSvc.Stop)
  2088. gatewaySvc := service.NewOpenAIGatewayService(
  2089. accountRepo,
  2090. nil,
  2091. nil,
  2092. nil,
  2093. nil,
  2094. nil,
  2095. nil,
  2096. cfg,
  2097. nil,
  2098. nil,
  2099. service.NewBillingService(cfg, nil),
  2100. nil,
  2101. billingCacheSvc,
  2102. upstream,
  2103. &service.DeferredService{},
  2104. nil,
  2105. nil,
  2106. nil,
  2107. nil,
  2108. nil,
  2109. nil,
  2110. nil,
  2111. )
  2112. h := NewOpenAIGatewayHandler(
  2113. gatewaySvc,
  2114. service.NewConcurrencyService(nil),
  2115. billingCacheSvc,
  2116. service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
  2117. nil,
  2118. nil,
  2119. nil,
  2120. nil,
  2121. cfg,
  2122. )
  2123. rec := httptest.NewRecorder()
  2124. c, _ := gin.CreateTestContext(rec)
  2125. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader(`{"model":"gpt-5.6-sol","input":"hello","stream":true}`))
  2126. c.Request.Header.Set("Content-Type", "application/json")
  2127. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  2128. ID: 1804, GroupID: &groupID,
  2129. User: &service.User{ID: 1704, Status: service.StatusActive},
  2130. Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
  2131. })
  2132. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1704, Concurrency: 0})
  2133. h.Responses(c)
  2134. require.Equal(t, []int64{9912, 9912}, upstream.calls())
  2135. require.Equal(t, http.StatusTooManyRequests, rec.Code)
  2136. require.Equal(t, "1", rec.Header().Get("Retry-After"))
  2137. require.Equal(t, "rate_limit_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
  2138. require.Equal(t, "Upstream rate limit exceeded, please retry later", gjson.GetBytes(rec.Body.Bytes(), "error.message").String())
  2139. }
  2140. func TestOpenAIResponsesWebSocket_FailoverOnUpstreamUsageLimitEvent(t *testing.T) {
  2141. gin.SetMode(gin.TestMode)
  2142. firstHitCh := make(chan []byte, 1)
  2143. secondHitCh := make(chan []byte, 1)
  2144. firstUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2145. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  2146. if err != nil {
  2147. return
  2148. }
  2149. defer func() { _ = conn.CloseNow() }()
  2150. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  2151. _, payload, readErr := conn.Read(readCtx)
  2152. cancelRead()
  2153. if readErr == nil {
  2154. firstHitCh <- payload
  2155. }
  2156. writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
  2157. _ = conn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"error","error":{"code":"rate_limit_exceeded","type":"usage_limit_reached","message":"The usage limit has been reached"}}`))
  2158. cancelWrite()
  2159. }))
  2160. defer firstUpstream.Close()
  2161. secondUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2162. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  2163. if err != nil {
  2164. return
  2165. }
  2166. defer func() { _ = conn.CloseNow() }()
  2167. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  2168. _, payload, readErr := conn.Read(readCtx)
  2169. cancelRead()
  2170. if readErr == nil {
  2171. secondHitCh <- payload
  2172. }
  2173. writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
  2174. _ = conn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.completed","response":{"id":"resp_ws_failover_ok","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`))
  2175. cancelWrite()
  2176. _ = conn.Close(coderws.StatusNormalClosure, "done")
  2177. }))
  2178. defer secondUpstream.Close()
  2179. groupID := int64(4202)
  2180. accounts := []service.Account{
  2181. {
  2182. ID: 9902,
  2183. Name: "openai-ws-rate-limited",
  2184. Platform: service.PlatformOpenAI,
  2185. Type: service.AccountTypeAPIKey,
  2186. Status: service.StatusActive,
  2187. Schedulable: true,
  2188. Concurrency: 1,
  2189. Priority: 1,
  2190. Credentials: map[string]any{
  2191. "api_key": "sk-first",
  2192. "base_url": firstUpstream.URL,
  2193. },
  2194. Extra: map[string]any{
  2195. "openai_apikey_responses_websockets_v2_enabled": true,
  2196. "openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
  2197. },
  2198. },
  2199. {
  2200. ID: 9903,
  2201. Name: "openai-ws-healthy",
  2202. Platform: service.PlatformOpenAI,
  2203. Type: service.AccountTypeAPIKey,
  2204. Status: service.StatusActive,
  2205. Schedulable: true,
  2206. Concurrency: 1,
  2207. Priority: 2,
  2208. Credentials: map[string]any{
  2209. "api_key": "sk-second",
  2210. "base_url": secondUpstream.URL,
  2211. },
  2212. Extra: map[string]any{
  2213. "openai_apikey_responses_websockets_v2_enabled": true,
  2214. "openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
  2215. },
  2216. },
  2217. }
  2218. cfg := &config.Config{}
  2219. cfg.RunMode = config.RunModeSimple
  2220. cfg.Default.RateMultiplier = 1
  2221. cfg.Security.URLAllowlist.Enabled = false
  2222. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2223. cfg.Gateway.OpenAIWS.Enabled = true
  2224. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2225. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2226. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  2227. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2228. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2229. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2230. cfg.Gateway.MaxAccountSwitches = 3
  2231. accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
  2232. rateLimitSvc := service.NewRateLimitService(accountRepo, nil, cfg, nil, nil)
  2233. billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
  2234. gatewaySvc := service.NewOpenAIGatewayService(
  2235. accountRepo,
  2236. nil,
  2237. nil,
  2238. nil,
  2239. nil,
  2240. nil,
  2241. nil,
  2242. cfg,
  2243. nil,
  2244. nil,
  2245. service.NewBillingService(cfg, nil),
  2246. rateLimitSvc,
  2247. billingCacheSvc,
  2248. nil,
  2249. &service.DeferredService{},
  2250. nil,
  2251. nil,
  2252. nil,
  2253. nil,
  2254. nil,
  2255. nil,
  2256. nil,
  2257. )
  2258. cache := &concurrencyCacheMock{
  2259. acquireUserSlotFn: func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
  2260. return true, nil
  2261. },
  2262. acquireAccountSlotFn: func(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
  2263. return true, nil
  2264. },
  2265. }
  2266. h := &OpenAIGatewayHandler{
  2267. gatewayService: gatewaySvc,
  2268. billingCacheService: billingCacheSvc,
  2269. apiKeyService: &service.APIKeyService{},
  2270. concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(cache), SSEPingFormatNone, time.Second),
  2271. maxAccountSwitches: 3,
  2272. }
  2273. apiKey := &service.APIKey{
  2274. ID: 1802,
  2275. GroupID: &groupID,
  2276. User: &service.User{ID: 1702, Status: service.StatusActive},
  2277. Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
  2278. }
  2279. router := gin.New()
  2280. router.Use(func(c *gin.Context) {
  2281. c.Set(string(middleware.ContextKeyAPIKey), apiKey)
  2282. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
  2283. c.Next()
  2284. })
  2285. router.GET("/openai/v1/responses", h.ResponsesWebSocket)
  2286. handlerServer := httptest.NewServer(router)
  2287. defer handlerServer.Close()
  2288. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2289. clientConn, _, err := coderws.Dial(
  2290. dialCtx,
  2291. "ws"+strings.TrimPrefix(handlerServer.URL, "http")+"/openai/v1/responses",
  2292. &coderws.DialOptions{CompressionMode: coderws.CompressionContextTakeover},
  2293. )
  2294. cancelDial()
  2295. require.NoError(t, err)
  2296. defer func() { _ = clientConn.CloseNow() }()
  2297. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  2298. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`))
  2299. cancelWrite()
  2300. require.NoError(t, err)
  2301. readCtx, cancelRead := context.WithTimeout(context.Background(), 5*time.Second)
  2302. _, event, err := clientConn.Read(readCtx)
  2303. cancelRead()
  2304. require.NoError(t, err)
  2305. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  2306. require.Equal(t, "resp_ws_failover_ok", gjson.GetBytes(event, "response.id").String())
  2307. select {
  2308. case <-firstHitCh:
  2309. case <-time.After(3 * time.Second):
  2310. t.Fatal("等待第一个上游收到首帧超时")
  2311. }
  2312. select {
  2313. case <-secondHitCh:
  2314. case <-time.After(3 * time.Second):
  2315. t.Fatal("等待第二个上游收到重放首帧超时")
  2316. }
  2317. require.Equal(t, []int64{int64(9902)}, accountRepo.rateLimitedIDs)
  2318. }
  2319. func TestOpenAIResponsesWebSocket_FirstOutputTimeoutWithoutDownstreamReusesClientForOneFailover(t *testing.T) {
  2320. gin.SetMode(gin.TestMode)
  2321. firstHitCh := make(chan []byte, 1)
  2322. secondHitCh := make(chan []byte, 1)
  2323. var firstConnections atomic.Int32
  2324. var secondConnections atomic.Int32
  2325. firstUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2326. firstConnections.Add(1)
  2327. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  2328. if err != nil {
  2329. return
  2330. }
  2331. defer func() { _ = conn.CloseNow() }()
  2332. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  2333. _, payload, readErr := conn.Read(readCtx)
  2334. cancelRead()
  2335. if readErr == nil {
  2336. firstHitCh <- payload
  2337. }
  2338. select {
  2339. case <-r.Context().Done():
  2340. case <-time.After(3 * time.Second):
  2341. }
  2342. }))
  2343. defer firstUpstream.Close()
  2344. secondUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2345. secondConnections.Add(1)
  2346. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  2347. if err != nil {
  2348. return
  2349. }
  2350. defer func() { _ = conn.CloseNow() }()
  2351. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  2352. _, payload, readErr := conn.Read(readCtx)
  2353. cancelRead()
  2354. if readErr == nil {
  2355. secondHitCh <- payload
  2356. }
  2357. for _, event := range []string{
  2358. `{"type":"response.created","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1"}}`,
  2359. `{"type":"response.output_text.delta","response_id":"resp_ws_timeout_b","delta":"recovered"}`,
  2360. `{"type":"response.completed","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`,
  2361. } {
  2362. writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
  2363. writeErr := conn.Write(writeCtx, coderws.MessageText, []byte(event))
  2364. cancelWrite()
  2365. if writeErr != nil {
  2366. return
  2367. }
  2368. }
  2369. readCtx, cancelRead = context.WithTimeout(r.Context(), 3*time.Second)
  2370. _, _, _ = conn.Read(readCtx)
  2371. cancelRead()
  2372. }))
  2373. defer secondUpstream.Close()
  2374. groupID := int64(4212)
  2375. accounts := []service.Account{
  2376. {
  2377. ID: 9912,
  2378. Name: "openai-ws-first-semantic-timeout",
  2379. Platform: service.PlatformOpenAI,
  2380. Type: service.AccountTypeAPIKey,
  2381. Status: service.StatusActive,
  2382. Schedulable: true,
  2383. Concurrency: 1,
  2384. Priority: 1,
  2385. Credentials: map[string]any{"api_key": "sk-first", "base_url": firstUpstream.URL},
  2386. Extra: map[string]any{
  2387. "openai_apikey_responses_websockets_v2_enabled": true,
  2388. "openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
  2389. },
  2390. },
  2391. {
  2392. ID: 9913,
  2393. Name: "openai-ws-failover-healthy",
  2394. Platform: service.PlatformOpenAI,
  2395. Type: service.AccountTypeAPIKey,
  2396. Status: service.StatusActive,
  2397. Schedulable: true,
  2398. Concurrency: 1,
  2399. Priority: 2,
  2400. Credentials: map[string]any{"api_key": "sk-second", "base_url": secondUpstream.URL},
  2401. Extra: map[string]any{
  2402. "openai_apikey_responses_websockets_v2_enabled": true,
  2403. "openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
  2404. },
  2405. },
  2406. }
  2407. cfg := &config.Config{}
  2408. cfg.RunMode = config.RunModeSimple
  2409. cfg.Default.RateMultiplier = 1
  2410. cfg.Security.URLAllowlist.Enabled = false
  2411. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2412. cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 1
  2413. cfg.Gateway.OpenAIWS.Enabled = true
  2414. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2415. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2416. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  2417. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2418. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2419. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2420. cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 3
  2421. cfg.Gateway.MaxAccountSwitches = 3
  2422. accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
  2423. rateLimitSvc := service.NewRateLimitService(accountRepo, nil, cfg, nil, nil)
  2424. billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
  2425. gatewaySvc := service.NewOpenAIGatewayService(
  2426. accountRepo, nil, nil, nil, nil, nil, nil, cfg, nil, nil,
  2427. service.NewBillingService(cfg, nil), rateLimitSvc, billingCacheSvc,
  2428. nil, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil,
  2429. )
  2430. cache := &concurrencyCacheMock{
  2431. acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil },
  2432. acquireAccountSlotFn: func(context.Context, int64, int, string) (bool, error) {
  2433. return true, nil
  2434. },
  2435. }
  2436. h := &OpenAIGatewayHandler{
  2437. gatewayService: gatewaySvc,
  2438. billingCacheService: billingCacheSvc,
  2439. apiKeyService: &service.APIKeyService{},
  2440. concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(cache), SSEPingFormatNone, time.Second),
  2441. maxAccountSwitches: 3,
  2442. }
  2443. apiKey := &service.APIKey{
  2444. ID: 1812,
  2445. GroupID: &groupID,
  2446. User: &service.User{ID: 1712, Status: service.StatusActive},
  2447. Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
  2448. }
  2449. handlerDone := make(chan struct{})
  2450. router := gin.New()
  2451. router.Use(func(c *gin.Context) {
  2452. c.Set(string(middleware.ContextKeyAPIKey), apiKey)
  2453. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
  2454. c.Next()
  2455. })
  2456. router.GET("/openai/v1/responses", func(c *gin.Context) {
  2457. h.ResponsesWebSocket(c)
  2458. close(handlerDone)
  2459. })
  2460. handlerServer := httptest.NewServer(router)
  2461. defer handlerServer.Close()
  2462. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2463. clientConn, _, err := coderws.Dial(
  2464. dialCtx,
  2465. "ws"+strings.TrimPrefix(handlerServer.URL, "http")+"/openai/v1/responses",
  2466. &coderws.DialOptions{CompressionMode: coderws.CompressionContextTakeover},
  2467. )
  2468. cancelDial()
  2469. require.NoError(t, err)
  2470. defer func() { _ = clientConn.CloseNow() }()
  2471. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  2472. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`))
  2473. cancelWrite()
  2474. require.NoError(t, err)
  2475. var eventTypes []string
  2476. readCtx, cancelRead := context.WithTimeout(context.Background(), 6*time.Second)
  2477. for {
  2478. _, event, readErr := clientConn.Read(readCtx)
  2479. require.NoError(t, readErr)
  2480. eventType := gjson.GetBytes(event, "type").String()
  2481. eventTypes = append(eventTypes, eventType)
  2482. if eventType == "response.completed" {
  2483. require.Equal(t, "resp_ws_timeout_b", gjson.GetBytes(event, "response.id").String())
  2484. break
  2485. }
  2486. }
  2487. cancelRead()
  2488. require.Contains(t, eventTypes, "response.output_text.delta")
  2489. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  2490. select {
  2491. case <-handlerDone:
  2492. case <-time.After(3 * time.Second):
  2493. t.Fatal("websocket handler did not finish after healthy failover turn")
  2494. }
  2495. select {
  2496. case <-firstHitCh:
  2497. case <-time.After(3 * time.Second):
  2498. t.Fatal("first upstream did not receive replayable request")
  2499. }
  2500. select {
  2501. case <-secondHitCh:
  2502. case <-time.After(3 * time.Second):
  2503. t.Fatal("second upstream did not receive replayed request")
  2504. }
  2505. require.Equal(t, int32(1), firstConnections.Load())
  2506. require.Equal(t, int32(1), secondConnections.Load())
  2507. require.NotContains(t, accountRepo.rateLimitedIDs, int64(9913), "healthy failover account must not be penalized")
  2508. }
  2509. func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSUsageLogCase) openAIResponsesWSUsageLogResult {
  2510. t.Helper()
  2511. gin.SetMode(gin.TestMode)
  2512. turnCount := 1
  2513. if strings.TrimSpace(tc.midPayload) != "" {
  2514. turnCount++
  2515. }
  2516. if strings.TrimSpace(tc.secondPayload) != "" {
  2517. turnCount++
  2518. }
  2519. upstreamPayloadCh := make(chan []byte, turnCount)
  2520. upstreamErrCh := make(chan error, 1)
  2521. var channelSvc *service.ChannelService
  2522. upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  2523. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{
  2524. CompressionMode: coderws.CompressionContextTakeover,
  2525. })
  2526. if err != nil {
  2527. upstreamErrCh <- err
  2528. return
  2529. }
  2530. defer func() {
  2531. _ = conn.CloseNow()
  2532. }()
  2533. for turn := 1; turn <= turnCount; turn++ {
  2534. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  2535. msgType, payload, readErr := conn.Read(readCtx)
  2536. cancelRead()
  2537. if readErr != nil {
  2538. upstreamErrCh <- readErr
  2539. return
  2540. }
  2541. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  2542. upstreamErrCh <- errors.New("unexpected upstream websocket message type")
  2543. return
  2544. }
  2545. upstreamPayloadCh <- payload
  2546. if turn == 1 && tc.afterFirstUpstreamRequest != nil {
  2547. if callbackErr := tc.afterFirstUpstreamRequest(channelSvc); callbackErr != nil {
  2548. upstreamErrCh <- callbackErr
  2549. return
  2550. }
  2551. }
  2552. response := fmt.Sprintf(
  2553. `{"type":"response.completed","response":{"id":"resp_usage_e2e_%d","model":%q,"usage":{"input_tokens":2,"output_tokens":1}}}`,
  2554. turn,
  2555. gjson.GetBytes(payload, "model").String(),
  2556. )
  2557. writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
  2558. writeErr := conn.Write(writeCtx, coderws.MessageText, []byte(response))
  2559. cancelWrite()
  2560. if writeErr != nil {
  2561. upstreamErrCh <- writeErr
  2562. return
  2563. }
  2564. }
  2565. upstreamErrCh <- nil
  2566. }))
  2567. defer upstreamServer.Close()
  2568. groupID := int64(4201)
  2569. account := service.Account{
  2570. ID: 9901,
  2571. Name: "openai-ws-passthrough-usage-e2e",
  2572. Platform: service.PlatformOpenAI,
  2573. Type: service.AccountTypeAPIKey,
  2574. Status: service.StatusActive,
  2575. Schedulable: true,
  2576. Concurrency: 1,
  2577. Credentials: map[string]any{
  2578. "api_key": "sk-test",
  2579. "base_url": upstreamServer.URL,
  2580. "model_mapping": tc.accountModelMapping,
  2581. },
  2582. Extra: map[string]any{
  2583. "openai_apikey_responses_websockets_v2_enabled": true,
  2584. "openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
  2585. },
  2586. }
  2587. if strings.TrimSpace(tc.ingressMode) != "" {
  2588. account.Extra["openai_apikey_responses_websockets_v2_mode"] = tc.ingressMode
  2589. }
  2590. cfg := &config.Config{}
  2591. cfg.RunMode = config.RunModeSimple
  2592. cfg.Default.RateMultiplier = 1
  2593. cfg.Security.URLAllowlist.Enabled = false
  2594. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  2595. cfg.Gateway.OpenAIWS.Enabled = true
  2596. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2597. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2598. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  2599. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  2600. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  2601. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  2602. accountRepo := &openAIWSUsageHandlerAccountRepoStub{account: account}
  2603. usageRepo := &openAIWSUsageHandlerUsageLogRepoStub{created: make(chan *service.UsageLog, turnCount)}
  2604. if len(tc.channelMapping) > 0 {
  2605. channelSvc = service.NewChannelService(&openAIWSUsageHandlerChannelRepoStub{
  2606. channels: []service.Channel{{
  2607. ID: 7701,
  2608. Name: "openai-ws-e2e-channel",
  2609. Status: service.StatusActive,
  2610. GroupIDs: []int64{groupID},
  2611. ModelMapping: map[string]map[string]string{service.PlatformOpenAI: tc.channelMapping},
  2612. BillingModelSource: tc.billingModelSource,
  2613. }},
  2614. groupPlatforms: map[int64]string{groupID: service.PlatformOpenAI},
  2615. }, nil, nil, nil, nil)
  2616. }
  2617. billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
  2618. gatewaySvc := service.NewOpenAIGatewayService(
  2619. accountRepo,
  2620. usageRepo,
  2621. nil,
  2622. nil,
  2623. nil,
  2624. nil,
  2625. nil,
  2626. cfg,
  2627. nil,
  2628. nil,
  2629. service.NewBillingService(cfg, nil),
  2630. nil,
  2631. billingCacheSvc,
  2632. nil,
  2633. &service.DeferredService{},
  2634. nil,
  2635. nil,
  2636. nil,
  2637. channelSvc,
  2638. nil,
  2639. nil,
  2640. nil, // userPlatformQuotaRepo
  2641. )
  2642. cache := &concurrencyCacheMock{
  2643. acquireUserSlotFn: func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
  2644. return true, nil
  2645. },
  2646. acquireAccountSlotFn: func(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
  2647. return true, nil
  2648. },
  2649. }
  2650. h := &OpenAIGatewayHandler{
  2651. gatewayService: gatewaySvc,
  2652. billingCacheService: billingCacheSvc,
  2653. apiKeyService: &service.APIKeyService{},
  2654. concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(cache), SSEPingFormatNone, time.Second),
  2655. }
  2656. apiKey := &service.APIKey{
  2657. ID: 1801,
  2658. GroupID: &groupID,
  2659. User: &service.User{ID: 1701, Status: service.StatusActive},
  2660. }
  2661. if tc.group != nil {
  2662. apiKey.Group = tc.group
  2663. }
  2664. router := gin.New()
  2665. router.Use(func(c *gin.Context) {
  2666. c.Set(string(middleware.ContextKeyAPIKey), apiKey)
  2667. c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
  2668. c.Next()
  2669. })
  2670. router.GET("/openai/v1/responses", h.ResponsesWebSocket)
  2671. handlerServer := httptest.NewServer(router)
  2672. defer handlerServer.Close()
  2673. headers := http.Header{}
  2674. if tc.userAgent != nil {
  2675. headers.Set("User-Agent", *tc.userAgent)
  2676. }
  2677. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  2678. clientConn, _, err := coderws.Dial(
  2679. dialCtx,
  2680. "ws"+strings.TrimPrefix(handlerServer.URL, "http")+"/openai/v1/responses",
  2681. &coderws.DialOptions{HTTPHeader: headers, CompressionMode: coderws.CompressionContextTakeover},
  2682. )
  2683. cancelDial()
  2684. require.NoError(t, err)
  2685. defer func() {
  2686. _ = clientConn.CloseNow()
  2687. }()
  2688. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  2689. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.firstPayload))
  2690. cancelWrite()
  2691. require.NoError(t, err)
  2692. if tc.firstFrameCloseExpected {
  2693. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  2694. _, _, readErr := clientConn.Read(readCtx)
  2695. cancelRead()
  2696. require.Error(t, readErr, "first frame should have been rejected with a close")
  2697. var closeErr coderws.CloseError
  2698. require.ErrorAs(t, readErr, &closeErr)
  2699. require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
  2700. require.Contains(t, closeErr.Reason, "not available for this group")
  2701. _ = clientConn.CloseNow()
  2702. return openAIResponsesWSUsageLogResult{}
  2703. }
  2704. clientEvents := make([][]byte, 0, turnCount)
  2705. readCompleted := func() {
  2706. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  2707. _, event, readErr := clientConn.Read(readCtx)
  2708. cancelRead()
  2709. require.NoError(t, readErr)
  2710. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  2711. clientEvents = append(clientEvents, append([]byte(nil), event...))
  2712. }
  2713. readCompleted()
  2714. if tc.midPayload != "" {
  2715. writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
  2716. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.midPayload))
  2717. cancelWrite()
  2718. require.NoError(t, err)
  2719. readCompleted()
  2720. }
  2721. if strings.TrimSpace(tc.secondPayload) != "" && (turnCount >= 2) {
  2722. writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
  2723. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.secondPayload))
  2724. cancelWrite()
  2725. require.NoError(t, err)
  2726. if tc.secondTurnCloseExpected {
  2727. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  2728. _, _, readErr := clientConn.Read(readCtx)
  2729. cancelRead()
  2730. require.Error(t, readErr, "second turn should have been rejected with a close")
  2731. var closeErr coderws.CloseError
  2732. require.ErrorAs(t, readErr, &closeErr)
  2733. require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
  2734. require.Contains(t, closeErr.Reason, "not available for this group")
  2735. _ = clientConn.CloseNow()
  2736. return openAIResponsesWSUsageLogResult{}
  2737. }
  2738. readCompleted()
  2739. }
  2740. _ = clientConn.Close(coderws.StatusNormalClosure, "done")
  2741. usageLogs := make([]*service.UsageLog, 0, turnCount)
  2742. for len(usageLogs) < turnCount {
  2743. select {
  2744. case usageLog := <-usageRepo.created:
  2745. require.NotNil(t, usageLog)
  2746. usageLogs = append(usageLogs, usageLog)
  2747. case <-time.After(3 * time.Second):
  2748. t.Fatal("等待 WebSocket usage log 写入超时")
  2749. }
  2750. }
  2751. upstreamPayloads := make([][]byte, 0, turnCount)
  2752. for len(upstreamPayloads) < turnCount {
  2753. select {
  2754. case payload := <-upstreamPayloadCh:
  2755. upstreamPayloads = append(upstreamPayloads, payload)
  2756. case <-time.After(3 * time.Second):
  2757. t.Fatal("等待上游 WebSocket 请求帧超时")
  2758. }
  2759. }
  2760. select {
  2761. case upstreamErr := <-upstreamErrCh:
  2762. require.NoError(t, upstreamErr)
  2763. case <-time.After(3 * time.Second):
  2764. t.Fatal("等待上游 WebSocket 结束超时")
  2765. }
  2766. return openAIResponsesWSUsageLogResult{
  2767. log: usageLogs[0],
  2768. logs: usageLogs,
  2769. upstreamFirstPayload: upstreamPayloads[0],
  2770. upstreamPayloads: upstreamPayloads,
  2771. clientEvents: clientEvents,
  2772. }
  2773. }
  2774. func testStringPtr(v string) *string {
  2775. return &v
  2776. }
  2777. func TestOpenAIForwardErrorAlreadyCommunicated(t *testing.T) {
  2778. gin.SetMode(gin.TestMode)
  2779. t.Run("upstream response failed after write", func(t *testing.T) {
  2780. w := httptest.NewRecorder()
  2781. c, _ := gin.CreateTestContext(w)
  2782. c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
  2783. before := c.Writer.Size()
  2784. _, _ = c.Writer.WriteString(`event: response.failed
  2785. data: {"type":"response.failed","error":{"message":"This content was flagged"}}
  2786. `)
  2787. reported := openAIForwardErrorAlreadyCommunicated(c, before, errors.New("upstream response failed: This content was flagged"))
  2788. require.True(t, reported)
  2789. })
  2790. t.Run("no write still needs fallback", func(t *testing.T) {
  2791. w := httptest.NewRecorder()
  2792. c, _ := gin.CreateTestContext(w)
  2793. c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
  2794. reported := openAIForwardErrorAlreadyCommunicated(c, c.Writer.Size(), errors.New("upstream response failed: This content was flagged"))
  2795. require.False(t, reported)
  2796. })
  2797. t.Run("generic error after write still needs fallback", func(t *testing.T) {
  2798. w := httptest.NewRecorder()
  2799. c, _ := gin.CreateTestContext(w)
  2800. c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
  2801. before := c.Writer.Size()
  2802. _, _ = c.Writer.WriteString(":\n\n")
  2803. reported := openAIForwardErrorAlreadyCommunicated(c, before, errors.New("stream read error: unexpected EOF"))
  2804. require.False(t, reported)
  2805. })
  2806. // H-2: cyber_policy 命中且响应已写出时,即便 err 前缀不在白名单(非流式 400 cyber
  2807. // 返回 "openai cyber_policy:"、透传账号返回 "upstream error:"),也须判定已透传,避免
  2808. // ensureForwardErrorResponse 在已写出的完整响应尾部追加 SSE 污染响应体。
  2809. t.Run("cyber policy hit after write is already communicated", func(t *testing.T) {
  2810. w := httptest.NewRecorder()
  2811. c, _ := gin.CreateTestContext(w)
  2812. c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
  2813. service.MarkOpsCyberPolicy(c, service.CyberPolicyMark{Message: "blocked", UpstreamStatus: 400})
  2814. before := c.Writer.Size()
  2815. _, _ = c.Writer.WriteString(`{"error":{"code":"cyber_policy","message":"blocked"}}`)
  2816. require.True(t, openAIForwardErrorAlreadyCommunicated(c, before, errors.New("openai cyber_policy: blocked")))
  2817. })
  2818. // Size 守卫优先于 cyber 短路:cyber 命中但未写出任何响应时仍需补写错误。
  2819. t.Run("cyber policy without write still needs fallback", func(t *testing.T) {
  2820. w := httptest.NewRecorder()
  2821. c, _ := gin.CreateTestContext(w)
  2822. c.Request = httptest.NewRequest(http.MethodPost, EndpointResponses, nil)
  2823. service.MarkOpsCyberPolicy(c, service.CyberPolicyMark{Message: "blocked", UpstreamStatus: 400})
  2824. require.False(t, openAIForwardErrorAlreadyCommunicated(c, c.Writer.Size(), errors.New("openai cyber_policy: blocked")))
  2825. })
  2826. }
  2827. func (r *contentModerationHandlerTestRepo) UpdateLogOverturned(ctx context.Context, id int64) error {
  2828. return nil
  2829. }