openai_ws_forwarder_success_test.go 71 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052
  1. package service
  2. import (
  3. "context"
  4. "encoding/json"
  5. "errors"
  6. "io"
  7. "net/http"
  8. "net/http/httptest"
  9. "strconv"
  10. "strings"
  11. "sync"
  12. "sync/atomic"
  13. "testing"
  14. "time"
  15. "github.com/Wei-Shaw/sub2api/internal/config"
  16. "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
  17. coderws "github.com/coder/websocket"
  18. "github.com/gin-gonic/gin"
  19. "github.com/gorilla/websocket"
  20. "github.com/stretchr/testify/require"
  21. "github.com/tidwall/gjson"
  22. )
  23. func TestOpenAIGatewayService_Forward_WSv2_SuccessAndBindSticky(t *testing.T) {
  24. gin.SetMode(gin.TestMode)
  25. type receivedPayload struct {
  26. Type string
  27. PreviousResponseID string
  28. StreamExists bool
  29. Stream bool
  30. }
  31. receivedCh := make(chan receivedPayload, 1)
  32. upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
  33. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  34. conn, err := upgrader.Upgrade(w, r, nil)
  35. if err != nil {
  36. t.Errorf("upgrade websocket failed: %v", err)
  37. return
  38. }
  39. defer func() {
  40. _ = conn.Close()
  41. }()
  42. var request map[string]any
  43. if err := conn.ReadJSON(&request); err != nil {
  44. t.Errorf("read ws request failed: %v", err)
  45. return
  46. }
  47. requestJSON := requestToJSONString(request)
  48. receivedCh <- receivedPayload{
  49. Type: strings.TrimSpace(gjson.Get(requestJSON, "type").String()),
  50. PreviousResponseID: strings.TrimSpace(gjson.Get(requestJSON, "previous_response_id").String()),
  51. StreamExists: gjson.Get(requestJSON, "stream").Exists(),
  52. Stream: gjson.Get(requestJSON, "stream").Bool(),
  53. }
  54. if err := conn.WriteJSON(map[string]any{
  55. "type": "response.created",
  56. "response": map[string]any{
  57. "id": "resp_new_1",
  58. "model": "gpt-5.1",
  59. },
  60. }); err != nil {
  61. t.Errorf("write response.created failed: %v", err)
  62. return
  63. }
  64. if err := conn.WriteJSON(map[string]any{
  65. "type": "response.completed",
  66. "response": map[string]any{
  67. "id": "resp_new_1",
  68. "model": "gpt-5.1",
  69. "usage": map[string]any{
  70. "input_tokens": 12,
  71. "output_tokens": 7,
  72. "input_tokens_details": map[string]any{
  73. "cached_tokens": 3,
  74. },
  75. },
  76. },
  77. }); err != nil {
  78. t.Errorf("write response.completed failed: %v", err)
  79. return
  80. }
  81. }))
  82. defer wsServer.Close()
  83. rec := httptest.NewRecorder()
  84. c, _ := gin.CreateTestContext(rec)
  85. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  86. c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
  87. groupID := int64(1001)
  88. c.Set("api_key", &APIKey{GroupID: &groupID})
  89. cfg := &config.Config{}
  90. cfg.Security.URLAllowlist.Enabled = false
  91. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  92. cfg.Gateway.OpenAIWS.Enabled = true
  93. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  94. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  95. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  96. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  97. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  98. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  99. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 30
  100. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 10
  101. cfg.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = 3600
  102. upstream := &httpUpstreamRecorder{
  103. resp: &http.Response{
  104. StatusCode: http.StatusOK,
  105. Header: http.Header{"Content-Type": []string{"application/json"}},
  106. Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":1}}`)),
  107. },
  108. }
  109. cache := &stubGatewayCache{}
  110. svc := &OpenAIGatewayService{
  111. cfg: cfg,
  112. httpUpstream: upstream,
  113. cache: cache,
  114. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  115. toolCorrector: NewCodexToolCorrector(),
  116. }
  117. account := &Account{
  118. ID: 9,
  119. Name: "openai-ws",
  120. Platform: PlatformOpenAI,
  121. Type: AccountTypeAPIKey,
  122. Status: StatusActive,
  123. Schedulable: true,
  124. Concurrency: 2,
  125. Credentials: map[string]any{
  126. "api_key": "sk-test",
  127. "base_url": wsServer.URL,
  128. },
  129. Extra: map[string]any{
  130. "responses_websockets_v2_enabled": true,
  131. },
  132. }
  133. body := []byte(`{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_prev_1","input":[{"type":"input_text","text":"hello"}]}`)
  134. result, err := svc.Forward(context.Background(), c, account, body)
  135. require.NoError(t, err)
  136. require.NotNil(t, result)
  137. require.Equal(t, 12, result.Usage.InputTokens)
  138. require.Equal(t, 7, result.Usage.OutputTokens)
  139. require.Equal(t, 3, result.Usage.CacheReadInputTokens)
  140. require.Equal(t, "resp_new_1", result.RequestID)
  141. require.True(t, result.OpenAIWSMode)
  142. require.False(t, gjson.GetBytes(upstream.lastBody, "model").Exists(), "WSv2 成功时不应回落 HTTP 上游")
  143. received := <-receivedCh
  144. require.Equal(t, "response.create", received.Type)
  145. require.Equal(t, "resp_prev_1", received.PreviousResponseID)
  146. require.True(t, received.StreamExists, "WS 请求应携带 stream 字段")
  147. require.False(t, received.Stream, "应保持客户端 stream=false 的原始语义")
  148. store := svc.getOpenAIWSStateStore()
  149. mappedAccountID, getErr := store.GetResponseAccount(context.Background(), groupID, "resp_new_1")
  150. require.NoError(t, getErr)
  151. require.Equal(t, account.ID, mappedAccountID)
  152. connID, ok := store.GetResponseConn("resp_new_1")
  153. require.True(t, ok)
  154. require.NotEmpty(t, connID)
  155. responseBody := rec.Body.Bytes()
  156. require.Equal(t, "resp_new_1", gjson.GetBytes(responseBody, "id").String())
  157. }
  158. func TestOpenAIGatewayService_Forward_WSv2_UsesPatchedBodyAfterValidationDecode(t *testing.T) {
  159. gin.SetMode(gin.TestMode)
  160. type receivedPayload struct {
  161. MaxCompletionTokensExists bool
  162. }
  163. receivedCh := make(chan receivedPayload, 1)
  164. upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
  165. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  166. conn, err := upgrader.Upgrade(w, r, nil)
  167. if err != nil {
  168. t.Errorf("upgrade websocket failed: %v", err)
  169. return
  170. }
  171. defer func() { _ = conn.Close() }()
  172. var request map[string]any
  173. if err := conn.ReadJSON(&request); err != nil {
  174. t.Errorf("read ws request failed: %v", err)
  175. return
  176. }
  177. requestJSON := requestToJSONString(request)
  178. receivedCh <- receivedPayload{MaxCompletionTokensExists: gjson.Get(requestJSON, "max_completion_tokens").Exists()}
  179. if err := conn.WriteJSON(map[string]any{
  180. "type": "response.completed",
  181. "response": map[string]any{
  182. "id": "resp_patched_ws_1",
  183. "model": "gpt-5.3-codex-spark",
  184. "usage": map[string]any{"input_tokens": 1, "output_tokens": 1},
  185. },
  186. }); err != nil {
  187. t.Errorf("write response.completed failed: %v", err)
  188. return
  189. }
  190. }))
  191. defer wsServer.Close()
  192. rec := httptest.NewRecorder()
  193. c, _ := gin.CreateTestContext(rec)
  194. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  195. c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
  196. cfg := &config.Config{}
  197. cfg.Security.URLAllowlist.Enabled = false
  198. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  199. cfg.Gateway.OpenAIWS.Enabled = true
  200. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  201. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  202. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  203. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 30
  204. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 10
  205. svc := &OpenAIGatewayService{
  206. cfg: cfg,
  207. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  208. toolCorrector: NewCodexToolCorrector(),
  209. }
  210. account := &Account{
  211. ID: 10,
  212. Name: "openai-ws",
  213. Platform: PlatformOpenAI,
  214. Type: AccountTypeAPIKey,
  215. Status: StatusActive,
  216. Schedulable: true,
  217. Concurrency: 1,
  218. Credentials: map[string]any{
  219. "api_key": "sk-test",
  220. "base_url": wsServer.URL,
  221. },
  222. Extra: map[string]any{"responses_websockets_v2_enabled": true},
  223. }
  224. body := []byte(`{"model":"gpt-5.4","stream":false,"max_completion_tokens":12,"tools":[{"type":"image_generation"}],"input":[{"type":"input_text","text":"hello"}]}`)
  225. result, err := svc.Forward(context.Background(), c, account, body)
  226. require.NoError(t, err)
  227. require.NotNil(t, result)
  228. require.True(t, result.OpenAIWSMode)
  229. received := <-receivedCh
  230. require.False(t, received.MaxCompletionTokensExists)
  231. }
  232. func TestOpenAIGatewayService_Forward_WSv2_ImageGenerationCountsOutputs(t *testing.T) {
  233. gin.SetMode(gin.TestMode)
  234. upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
  235. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  236. conn, err := upgrader.Upgrade(w, r, nil)
  237. if err != nil {
  238. t.Errorf("upgrade websocket failed: %v", err)
  239. return
  240. }
  241. defer func() {
  242. _ = conn.Close()
  243. }()
  244. var request map[string]any
  245. if err := conn.ReadJSON(&request); err != nil {
  246. t.Errorf("read ws request failed: %v", err)
  247. return
  248. }
  249. if err := conn.WriteJSON(map[string]any{
  250. "type": "response.output_item.done",
  251. "item": map[string]any{
  252. "id": "ig_ws_1",
  253. "type": "image_generation_call",
  254. "status": "generating",
  255. "result": "final-image",
  256. },
  257. }); err != nil {
  258. t.Errorf("write response.output_item.done failed: %v", err)
  259. return
  260. }
  261. if err := conn.WriteJSON(map[string]any{
  262. "type": "response.completed",
  263. "response": map[string]any{
  264. "id": "resp_ws_image_1",
  265. "model": "gpt-5.4",
  266. "output": []any{
  267. map[string]any{
  268. "id": "ig_ws_1",
  269. "type": "image_generation_call",
  270. "status": "in_progress",
  271. "result": "final-image",
  272. },
  273. },
  274. "usage": map[string]any{
  275. "input_tokens": 9,
  276. "output_tokens": 4,
  277. },
  278. },
  279. }); err != nil {
  280. t.Errorf("write response.completed failed: %v", err)
  281. return
  282. }
  283. }))
  284. defer wsServer.Close()
  285. rec := httptest.NewRecorder()
  286. c, _ := gin.CreateTestContext(rec)
  287. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  288. groupID := int64(1010)
  289. c.Set("api_key", &APIKey{
  290. GroupID: &groupID,
  291. Group: &Group{
  292. ID: groupID,
  293. AllowImageGeneration: true,
  294. },
  295. })
  296. cfg := &config.Config{}
  297. cfg.Security.URLAllowlist.Enabled = false
  298. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  299. cfg.Gateway.OpenAIWS.Enabled = true
  300. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  301. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  302. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  303. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  304. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  305. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  306. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  307. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  308. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5
  309. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  310. svc := &OpenAIGatewayService{
  311. cfg: cfg,
  312. httpUpstream: &httpUpstreamRecorder{},
  313. cache: &stubGatewayCache{},
  314. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  315. toolCorrector: NewCodexToolCorrector(),
  316. }
  317. account := &Account{
  318. ID: 10,
  319. Name: "openai-ws-image",
  320. Platform: PlatformOpenAI,
  321. Type: AccountTypeAPIKey,
  322. Status: StatusActive,
  323. Schedulable: true,
  324. Concurrency: 1,
  325. Credentials: map[string]any{
  326. "api_key": "sk-test",
  327. "base_url": wsServer.URL,
  328. },
  329. Extra: map[string]any{
  330. "responses_websockets_v2_enabled": true,
  331. },
  332. }
  333. body := []byte(`{"model":"gpt-5.4","stream":false,"input":"draw","tools":[{"type":"image_generation","model":"gpt-image-2","size":"1024x1024"}],"tool_choice":{"type":"image_generation"}}`)
  334. result, err := svc.Forward(context.Background(), c, account, body)
  335. require.NoError(t, err)
  336. require.NotNil(t, result)
  337. require.Equal(t, "resp_ws_image_1", result.RequestID)
  338. require.Equal(t, 1, result.ImageCount)
  339. require.Equal(t, "1K", result.ImageSize)
  340. require.Equal(t, "gpt-image-2", result.BillingModel)
  341. require.Equal(t, 9, result.Usage.InputTokens)
  342. require.Equal(t, 4, result.Usage.OutputTokens)
  343. require.True(t, result.OpenAIWSMode)
  344. require.Equal(t, "resp_ws_image_1", gjson.GetBytes(rec.Body.Bytes(), "id").String())
  345. require.Equal(t, "completed", gjson.GetBytes(rec.Body.Bytes(), "output.0.status").String())
  346. }
  347. func requestToJSONString(payload map[string]any) string {
  348. if len(payload) == 0 {
  349. return "{}"
  350. }
  351. b, err := json.Marshal(payload)
  352. if err != nil {
  353. return "{}"
  354. }
  355. return string(b)
  356. }
  357. func TestOpenAIGatewayService_BuildOpenAIWSHeadersPreservesCodexIdentity(t *testing.T) {
  358. gin.SetMode(gin.TestMode)
  359. rec := httptest.NewRecorder()
  360. c, _ := gin.CreateTestContext(rec)
  361. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  362. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
  363. c.Request.Header.Set("X-Codex-Window-ID", "window-ws")
  364. c.Request.Header.Set("X-Codex-Installation-ID", "installation-ws")
  365. c.Request.Header.Set("session-id", "session-ws")
  366. c.Request.Header.Set("thread-id", "thread-ws")
  367. c.Request.Header.Set("x-client-request-id", "client-request-ws")
  368. c.Request.Header.Set("X-Test", "blocked")
  369. svc := &OpenAIGatewayService{}
  370. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  371. headers, _, err := svc.buildOpenAIWSHeaders(
  372. context.Background(),
  373. c,
  374. account,
  375. "token",
  376. OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2},
  377. true,
  378. "",
  379. "",
  380. "",
  381. "",
  382. "",
  383. )
  384. require.NoError(t, err)
  385. require.Equal(t, "window-ws", headers.Get("X-Codex-Window-ID"))
  386. require.Equal(t, "installation-ws", headers.Get("X-Codex-Installation-ID"))
  387. require.Equal(t, "session-ws", headers.Get("session-id"))
  388. require.Equal(t, "thread-ws", headers.Get("thread-id"))
  389. require.Equal(t, "client-request-ws", headers.Get("x-client-request-id"))
  390. require.Empty(t, headers.Get("X-Test"))
  391. }
  392. func TestOpenAIGatewayService_BuildOpenAIWSHeadersDeviceModePreservesNamespacedClientSessionIdentity(t *testing.T) {
  393. gin.SetMode(gin.TestMode)
  394. rec := httptest.NewRecorder()
  395. c, _ := gin.CreateTestContext(rec)
  396. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  397. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
  398. c.Request.Header.Set("X-Codex-Installation-ID", "client-installation")
  399. c.Request.Header.Set("X-Codex-Window-ID", "client-window")
  400. c.Request.Header.Set("session-id", "client-session")
  401. c.Request.Header.Set("thread-id", "client-thread")
  402. c.Request.Header.Set("x-client-request-id", "client-request")
  403. account := newTestOAuthAccount(1300, map[string]any{codexFingerprintModeExtraKey: "device"})
  404. ids := resolveCodexFingerprintIDsFromRequest(account, c.Request.Header)
  405. require.NotNil(t, ids)
  406. stageCodexFingerprintIDs(c, ids)
  407. svc := &OpenAIGatewayService{}
  408. headers, _, err := svc.buildOpenAIWSHeaders(
  409. context.Background(),
  410. c,
  411. account,
  412. "token",
  413. OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2},
  414. true,
  415. "",
  416. "",
  417. "",
  418. "",
  419. "",
  420. )
  421. require.NoError(t, err)
  422. require.Equal(t, ids.installationID, headers.Get("x-codex-installation-id"))
  423. require.NotEqual(t, "client-installation", headers.Get("x-codex-installation-id"))
  424. require.Equal(t, scopeCodexAccountIdentityValue(account, 0, "window", "client-window"), headers.Get("x-codex-window-id"))
  425. require.Equal(t, scopeCodexAccountIdentityValue(account, 0, "session", "client-session"), headers.Get("session-id"))
  426. require.Equal(t, scopeCodexAccountIdentityValue(account, 0, "thread", "client-thread"), headers.Get("thread-id"))
  427. require.Equal(t, scopeCodexAccountIdentityValue(account, 0, "request", "client-request"), headers.Get("x-client-request-id"))
  428. }
  429. func TestLogOpenAIWSBindResponseAccountWarn(t *testing.T) {
  430. require.NotPanics(t, func() {
  431. logOpenAIWSBindResponseAccountWarn(1, 2, "resp_ok", nil)
  432. })
  433. require.NotPanics(t, func() {
  434. logOpenAIWSBindResponseAccountWarn(1, 2, "resp_err", errors.New("bind failed"))
  435. })
  436. }
  437. func TestOpenAIGatewayService_Forward_WSv2_RewriteModelAndToolCallsOnCompletedEvent(t *testing.T) {
  438. gin.SetMode(gin.TestMode)
  439. rec := httptest.NewRecorder()
  440. c, _ := gin.CreateTestContext(rec)
  441. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  442. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
  443. groupID := int64(3001)
  444. c.Set("api_key", &APIKey{GroupID: &groupID})
  445. cfg := &config.Config{}
  446. cfg.Security.URLAllowlist.Enabled = false
  447. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  448. cfg.Gateway.OpenAIWS.Enabled = true
  449. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  450. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  451. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  452. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  453. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  454. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  455. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  456. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  457. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5
  458. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  459. captureConn := &openAIWSCaptureConn{
  460. events: [][]byte{
  461. []byte(`{"type":"response.completed","response":{"id":"resp_model_tool_1","model":"gpt-5.1","tool_calls":[{"function":{"name":"apply_patch","arguments":"{\"file_path\":\"/tmp/a.txt\",\"old_string\":\"a\",\"new_string\":\"b\"}"}}],"usage":{"input_tokens":2,"output_tokens":1}},"tool_calls":[{"function":{"name":"apply_patch","arguments":"{\"file_path\":\"/tmp/a.txt\",\"old_string\":\"a\",\"new_string\":\"b\"}"}}]}`),
  462. },
  463. }
  464. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  465. pool := newOpenAIWSConnPool(cfg)
  466. pool.setClientDialerForTest(captureDialer)
  467. svc := &OpenAIGatewayService{
  468. cfg: cfg,
  469. httpUpstream: &httpUpstreamRecorder{},
  470. cache: &stubGatewayCache{},
  471. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  472. toolCorrector: NewCodexToolCorrector(),
  473. openaiWSPool: pool,
  474. }
  475. account := &Account{
  476. ID: 1301,
  477. Name: "openai-rewrite",
  478. Platform: PlatformOpenAI,
  479. Type: AccountTypeAPIKey,
  480. Status: StatusActive,
  481. Schedulable: true,
  482. Concurrency: 1,
  483. Credentials: map[string]any{
  484. "api_key": "sk-test",
  485. "model_mapping": map[string]any{
  486. "custom-original-model": "gpt-5.1",
  487. },
  488. },
  489. Extra: map[string]any{
  490. "responses_websockets_v2_enabled": true,
  491. },
  492. }
  493. body := []byte(`{"model":"custom-original-model","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
  494. result, err := svc.Forward(context.Background(), c, account, body)
  495. require.NoError(t, err)
  496. require.NotNil(t, result)
  497. require.Equal(t, "resp_model_tool_1", result.RequestID)
  498. require.Equal(t, "response.completed", result.UpstreamTerminalEvent)
  499. require.True(t, result.SucceededForScheduling())
  500. require.Equal(t, "custom-original-model", gjson.GetBytes(rec.Body.Bytes(), "model").String(), "响应模型应回写为原始请求模型")
  501. require.Equal(t, "edit", gjson.GetBytes(rec.Body.Bytes(), "tool_calls.0.function.name").String(), "工具名称应被修正为 OpenCode 规范")
  502. }
  503. func TestOpenAIGatewayService_Forward_WSv2_ResponseFailedIsNotSchedulingSuccess(t *testing.T) {
  504. gin.SetMode(gin.TestMode)
  505. rec := httptest.NewRecorder()
  506. c, _ := gin.CreateTestContext(rec)
  507. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  508. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
  509. cfg := newOpenAIWSV2TestConfig()
  510. cfg.Security.URLAllowlist.Enabled = false
  511. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  512. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  513. captureConn := &openAIWSCaptureConn{events: [][]byte{
  514. []byte(`{"type":"response.failed","response":{"id":"resp_failed_1","model":"gpt-5.5","error":{"code":"server_error","message":"Internal error"}}}`),
  515. }}
  516. pool := newOpenAIWSConnPool(cfg)
  517. pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn})
  518. svc := &OpenAIGatewayService{
  519. cfg: cfg,
  520. rateLimitService: NewRateLimitService(transientCooldownAccountRepo{}, nil, cfg, nil, nil),
  521. httpUpstream: &httpUpstreamRecorder{},
  522. cache: &stubGatewayCache{},
  523. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  524. toolCorrector: NewCodexToolCorrector(),
  525. openaiWSPool: pool,
  526. }
  527. account := &Account{
  528. ID: 1302,
  529. Platform: PlatformOpenAI,
  530. Type: AccountTypeAPIKey,
  531. Status: StatusActive,
  532. Schedulable: true,
  533. Concurrency: 1,
  534. Credentials: map[string]any{"api_key": "sk-test"},
  535. Extra: map[string]any{"responses_websockets_v2_enabled": true},
  536. }
  537. svc.recordOpenAIAccountModelTransientFailure(account, "gpt-5.5", time.Now())
  538. result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.5","stream":false,"input":"hello"}`))
  539. require.NoError(t, err)
  540. require.NotNil(t, result)
  541. require.Equal(t, "response.failed", result.UpstreamTerminalEvent)
  542. require.False(t, result.SucceededForScheduling())
  543. require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5"))
  544. }
  545. func TestOpenAIWSPayloadString_OnlyAcceptsStringValues(t *testing.T) {
  546. payload := map[string]any{
  547. "type": nil,
  548. "model": 123,
  549. "prompt_cache_key": " cache-key ",
  550. "previous_response_id": []byte(" resp_1 "),
  551. }
  552. require.Equal(t, "", openAIWSPayloadString(payload, "type"))
  553. require.Equal(t, "", openAIWSPayloadString(payload, "model"))
  554. require.Equal(t, "cache-key", openAIWSPayloadString(payload, "prompt_cache_key"))
  555. require.Equal(t, "resp_1", openAIWSPayloadString(payload, "previous_response_id"))
  556. }
  557. func TestOpenAIGatewayService_Forward_WSv2_PoolReuseNotOneToOne(t *testing.T) {
  558. gin.SetMode(gin.TestMode)
  559. var upgradeCount atomic.Int64
  560. var sequence atomic.Int64
  561. upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
  562. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  563. upgradeCount.Add(1)
  564. conn, err := upgrader.Upgrade(w, r, nil)
  565. if err != nil {
  566. t.Errorf("upgrade websocket failed: %v", err)
  567. return
  568. }
  569. defer func() {
  570. _ = conn.Close()
  571. }()
  572. for {
  573. var request map[string]any
  574. if err := conn.ReadJSON(&request); err != nil {
  575. return
  576. }
  577. idx := sequence.Add(1)
  578. responseID := "resp_reuse_" + strconv.FormatInt(idx, 10)
  579. if err := conn.WriteJSON(map[string]any{
  580. "type": "response.created",
  581. "response": map[string]any{
  582. "id": responseID,
  583. "model": "gpt-5.1",
  584. },
  585. }); err != nil {
  586. return
  587. }
  588. if err := conn.WriteJSON(map[string]any{
  589. "type": "response.completed",
  590. "response": map[string]any{
  591. "id": responseID,
  592. "model": "gpt-5.1",
  593. "usage": map[string]any{
  594. "input_tokens": 2,
  595. "output_tokens": 1,
  596. },
  597. },
  598. }); err != nil {
  599. return
  600. }
  601. }
  602. }))
  603. defer wsServer.Close()
  604. cfg := &config.Config{}
  605. cfg.Security.URLAllowlist.Enabled = false
  606. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  607. cfg.Gateway.OpenAIWS.Enabled = true
  608. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  609. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  610. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  611. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  612. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  613. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  614. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  615. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  616. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 30
  617. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 10
  618. svc := &OpenAIGatewayService{
  619. cfg: cfg,
  620. httpUpstream: &httpUpstreamRecorder{},
  621. cache: &stubGatewayCache{},
  622. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  623. toolCorrector: NewCodexToolCorrector(),
  624. }
  625. account := &Account{
  626. ID: 19,
  627. Name: "openai-ws",
  628. Platform: PlatformOpenAI,
  629. Type: AccountTypeAPIKey,
  630. Status: StatusActive,
  631. Schedulable: true,
  632. Concurrency: 2,
  633. Credentials: map[string]any{
  634. "api_key": "sk-test",
  635. "base_url": wsServer.URL,
  636. },
  637. Extra: map[string]any{
  638. "responses_websockets_v2_enabled": true,
  639. },
  640. }
  641. for i := 0; i < 2; i++ {
  642. rec := httptest.NewRecorder()
  643. c, _ := gin.CreateTestContext(rec)
  644. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  645. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
  646. groupID := int64(2001)
  647. c.Set("api_key", &APIKey{GroupID: &groupID})
  648. body := []byte(`{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_prev_reuse","input":[{"type":"input_text","text":"hello"}]}`)
  649. result, err := svc.Forward(context.Background(), c, account, body)
  650. require.NoError(t, err)
  651. require.NotNil(t, result)
  652. require.True(t, strings.HasPrefix(result.RequestID, "resp_reuse_"))
  653. }
  654. // 条件式 MarkBroken:正常终端事件退出后连接归还复用,不再无条件销毁。
  655. require.Equal(t, int64(1), upgradeCount.Load(), "正常完成后连接应归还复用,不应每次新建")
  656. metrics := svc.SnapshotOpenAIWSPoolMetrics()
  657. require.GreaterOrEqual(t, metrics.AcquireReuseTotal, int64(1))
  658. require.GreaterOrEqual(t, metrics.ConnPickTotal, int64(1))
  659. }
  660. func TestOpenAIGatewayService_Forward_WSv2_OAuthStoreFalseByDefault(t *testing.T) {
  661. gin.SetMode(gin.TestMode)
  662. rec := httptest.NewRecorder()
  663. c, _ := gin.CreateTestContext(rec)
  664. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  665. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
  666. c.Request.Header.Set("session_id", "sess-oauth-1")
  667. c.Request.Header.Set("conversation_id", "conv-oauth-1")
  668. c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2")
  669. cfg := &config.Config{}
  670. cfg.Security.URLAllowlist.Enabled = false
  671. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  672. cfg.Gateway.OpenAIWS.Enabled = true
  673. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  674. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  675. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  676. cfg.Gateway.OpenAIWS.AllowStoreRecovery = false
  677. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  678. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  679. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  680. captureConn := &openAIWSCaptureConn{
  681. events: [][]byte{
  682. []byte(`{"type":"response.completed","response":{"id":"resp_oauth_1","model":"gpt-5.1","usage":{"input_tokens":3,"output_tokens":2}}}`),
  683. },
  684. }
  685. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  686. pool := newOpenAIWSConnPool(cfg)
  687. pool.setClientDialerForTest(captureDialer)
  688. svc := &OpenAIGatewayService{
  689. cfg: cfg,
  690. httpUpstream: &httpUpstreamRecorder{},
  691. cache: &stubGatewayCache{},
  692. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  693. toolCorrector: NewCodexToolCorrector(),
  694. openaiWSPool: pool,
  695. }
  696. account := &Account{
  697. ID: 29,
  698. Name: "openai-oauth",
  699. Platform: PlatformOpenAI,
  700. Type: AccountTypeOAuth,
  701. Status: StatusActive,
  702. Schedulable: true,
  703. Concurrency: 1,
  704. Credentials: map[string]any{
  705. "access_token": "oauth-token-1",
  706. },
  707. Extra: map[string]any{
  708. "responses_websockets_v2_enabled": true,
  709. },
  710. }
  711. body := []byte(`{"model":"gpt-5.1","stream":false,"store":true,"input":[{"type":"input_text","text":"hello","namespace":"native-wsv2"}]}`)
  712. result, err := svc.Forward(context.Background(), c, account, body)
  713. require.NoError(t, err)
  714. require.NotNil(t, result)
  715. require.Equal(t, "resp_oauth_1", result.RequestID)
  716. require.NotNil(t, captureConn.lastWrite)
  717. requestJSON := requestToJSONString(captureConn.lastWrite)
  718. require.True(t, gjson.Get(requestJSON, "store").Exists(), "OAuth WSv2 应显式写入 store 字段")
  719. require.False(t, gjson.Get(requestJSON, "store").Bool(), "默认策略应将 OAuth store 置为 false")
  720. require.True(t, gjson.Get(requestJSON, "stream").Exists(), "WSv2 payload 应保留 stream 字段")
  721. require.True(t, gjson.Get(requestJSON, "stream").Bool(), "OAuth Codex 规范化后应强制 stream=true")
  722. require.Equal(t, "native-wsv2", gjson.Get(requestJSON, "input.0.namespace").String(), "OAuth WSv2 应保留原生 namespace")
  723. require.Equal(t, openAIWSBetaV2Value, captureDialer.lastHeaders.Get("OpenAI-Beta"))
  724. require.Equal(t, "remote_compaction_v2", captureDialer.lastHeaders.Get("x-codex-beta-features"))
  725. // OAuth 账号的 session_id/conversation_id 应同时按 API key 和上游账号隔离,
  726. // 测试中未设置 api_key 到 context,apiKeyID=0。
  727. require.Equal(t, isolateOpenAIUpstreamSessionID(0, account, "sess-oauth-1"), captureDialer.lastHeaders.Get("session_id"))
  728. require.Equal(t, isolateOpenAIUpstreamSessionID(0, account, "conv-oauth-1"), captureDialer.lastHeaders.Get("conversation_id"))
  729. }
  730. func TestOpenAIGatewayService_Forward_WSv2_OAuthSanitizesInvalidNativeToolItemID(t *testing.T) {
  731. gin.SetMode(gin.TestMode)
  732. rec := httptest.NewRecorder()
  733. c, _ := gin.CreateTestContext(rec)
  734. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  735. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
  736. groupID := int64(5662)
  737. c.Set("api_key", &APIKey{GroupID: &groupID})
  738. cfg := newOpenAIWSV2TestConfig()
  739. cfg.Security.URLAllowlist.Enabled = false
  740. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  741. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  742. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  743. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  744. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  745. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  746. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  747. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  748. captureConn := &openAIWSCaptureConn{events: [][]byte{
  749. []byte(`{"type":"response.completed","response":{"id":"resp_oauth_tool_history","model":"gpt-5.6-sol","usage":{"input_tokens":1,"output_tokens":1}}}`),
  750. }}
  751. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  752. pool := newOpenAIWSConnPool(cfg)
  753. pool.setClientDialerForTest(captureDialer)
  754. svc := &OpenAIGatewayService{
  755. cfg: cfg,
  756. httpUpstream: &httpUpstreamRecorder{},
  757. cache: &stubGatewayCache{},
  758. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  759. toolCorrector: NewCodexToolCorrector(),
  760. openaiWSPool: pool,
  761. }
  762. account := &Account{
  763. ID: 5662,
  764. Name: "openai-oauth-ws-tool-history",
  765. Platform: PlatformOpenAI,
  766. Type: AccountTypeOAuth,
  767. Status: StatusActive,
  768. Schedulable: true,
  769. Concurrency: 1,
  770. Credentials: map[string]any{"access_token": "test-oauth-token"},
  771. Extra: map[string]any{
  772. "responses_websockets_v2_enabled": true,
  773. },
  774. }
  775. body := []byte(`{"model":"gpt-5.6-sol","stream":false,"instructions":"Continue the task.","input":[{"type":"custom_tool_call","id":"fc_hotfix_probe","call_id":"fc_hotfix","name":"exec","input":"pwd","status":"completed"},{"type":"custom_tool_call_output","call_id":"fc_hotfix","output":"done"}]}`)
  776. result, err := svc.Forward(context.Background(), c, account, body)
  777. require.NoError(t, err)
  778. require.NotNil(t, result)
  779. captureConn.mu.Lock()
  780. requestPayload := cloneMapStringAny(captureConn.lastWrite)
  781. captureConn.mu.Unlock()
  782. requestJSON := requestToJSONString(requestPayload)
  783. require.Equal(t, "response.create", gjson.Get(requestJSON, "type").String())
  784. require.False(t, gjson.Get(requestJSON, "input.0.id").Exists(), "stale fc_* ID must not be replayed as a native custom_tool_call ID")
  785. require.Equal(t, "ctc_hotfix", gjson.Get(requestJSON, "input.0.call_id").String())
  786. require.Equal(t, "custom_tool_call_output", gjson.Get(requestJSON, "input.1.type").String())
  787. require.Equal(t, "ctc_hotfix", gjson.Get(requestJSON, "input.1.call_id").String())
  788. }
  789. func TestOpenAIGatewayService_Forward_WSv2_OAuthOriginatorCompatibility(t *testing.T) {
  790. gin.SetMode(gin.TestMode)
  791. // WS 握手头与 HTTP 出站共用身份收口:同样强制统一为网关规范身份,
  792. // 客户端自报的 originator / user-agent 不参与构造(issue #3901 的配对不变式自然满足)。
  793. tests := []struct {
  794. name string
  795. userAgent string
  796. originator string
  797. }{
  798. {name: "official desktop ua", userAgent: "Codex Desktop/1.2.3"},
  799. {
  800. name: "mismatched originator",
  801. userAgent: "codex_vscode/0.140.2 (Mac OS X 14.0; arm64) vscode (codex_vscode; 0.140.2)",
  802. originator: "codex_cli_rs",
  803. },
  804. {
  805. name: "tui identity",
  806. userAgent: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
  807. originator: "codex-tui",
  808. },
  809. {name: "official originator without ua", originator: "codex_vscode"},
  810. }
  811. for _, tt := range tests {
  812. t.Run(tt.name, func(t *testing.T) {
  813. rec := httptest.NewRecorder()
  814. c, _ := gin.CreateTestContext(rec)
  815. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  816. if tt.userAgent != "" {
  817. c.Request.Header.Set("User-Agent", tt.userAgent)
  818. }
  819. if tt.originator != "" {
  820. c.Request.Header.Set("originator", tt.originator)
  821. }
  822. cfg := &config.Config{}
  823. cfg.Security.URLAllowlist.Enabled = false
  824. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  825. cfg.Gateway.OpenAIWS.Enabled = true
  826. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  827. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  828. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  829. cfg.Gateway.OpenAIWS.AllowStoreRecovery = false
  830. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  831. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  832. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  833. captureConn := &openAIWSCaptureConn{
  834. events: [][]byte{
  835. []byte(`{"type":"response.completed","response":{"id":"resp_oauth_originator","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  836. },
  837. }
  838. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  839. pool := newOpenAIWSConnPool(cfg)
  840. pool.setClientDialerForTest(captureDialer)
  841. svc := &OpenAIGatewayService{
  842. cfg: cfg,
  843. httpUpstream: &httpUpstreamRecorder{},
  844. cache: &stubGatewayCache{},
  845. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  846. toolCorrector: NewCodexToolCorrector(),
  847. openaiWSPool: pool,
  848. }
  849. account := &Account{
  850. ID: 129,
  851. Name: "openai-oauth",
  852. Platform: PlatformOpenAI,
  853. Type: AccountTypeOAuth,
  854. Status: StatusActive,
  855. Schedulable: true,
  856. Concurrency: 1,
  857. Credentials: map[string]any{
  858. "access_token": "oauth-token-1",
  859. },
  860. Extra: map[string]any{
  861. "responses_websockets_v2_enabled": true,
  862. },
  863. }
  864. body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
  865. result, err := svc.Forward(context.Background(), c, account, body)
  866. require.NoError(t, err)
  867. require.NotNil(t, result)
  868. require.Equal(t, openai.CodexDefaultOriginator, captureDialer.lastHeaders.Get("originator"))
  869. require.Equal(t, codexCLIUserAgent, captureDialer.lastHeaders.Get("user-agent"))
  870. require.Equal(t, codexCLIVersion, captureDialer.lastHeaders.Get("version"))
  871. })
  872. }
  873. }
  874. // 账号级自定义 UA 是管理员的显式配置,WS 握手与 HTTP 出站必须一视同仁地生效——
  875. // 否则同一个账号在两种传输上以不同身份出站。指纹保留、版本段重建、originator 配套。
  876. func TestOpenAIGatewayService_Forward_WSv2_OAuthHonorsAccountUserAgent(t *testing.T) {
  877. gin.SetMode(gin.TestMode)
  878. rec := httptest.NewRecorder()
  879. c, _ := gin.CreateTestContext(rec)
  880. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  881. c.Request.Header.Set("User-Agent", "luna/1.0.0")
  882. cfg := &config.Config{}
  883. cfg.Security.URLAllowlist.Enabled = false
  884. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  885. cfg.Gateway.OpenAIWS.Enabled = true
  886. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  887. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  888. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  889. cfg.Gateway.OpenAIWS.AllowStoreRecovery = false
  890. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  891. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  892. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  893. captureConn := &openAIWSCaptureConn{
  894. events: [][]byte{
  895. []byte(`{"type":"response.completed","response":{"id":"resp_oauth_account_ua","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  896. },
  897. }
  898. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  899. pool := newOpenAIWSConnPool(cfg)
  900. pool.setClientDialerForTest(captureDialer)
  901. svc := &OpenAIGatewayService{
  902. cfg: cfg,
  903. httpUpstream: &httpUpstreamRecorder{},
  904. cache: &stubGatewayCache{},
  905. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  906. toolCorrector: NewCodexToolCorrector(),
  907. openaiWSPool: pool,
  908. }
  909. account := &Account{
  910. ID: 130,
  911. Name: "openai-oauth-custom-ua",
  912. Platform: PlatformOpenAI,
  913. Type: AccountTypeOAuth,
  914. Status: StatusActive,
  915. Schedulable: true,
  916. Concurrency: 1,
  917. Credentials: map[string]any{
  918. "access_token": "oauth-token-1",
  919. // 填写于某个历史版本的账号级 UA:指纹要保留,版本段不能被逐字沿用。
  920. "user_agent": "codex-tui/0.125.0 (Mac OS X 15.1.0; arm64) iTerm.app",
  921. },
  922. Extra: map[string]any{
  923. "responses_websockets_v2_enabled": true,
  924. },
  925. }
  926. body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
  927. result, err := svc.Forward(context.Background(), c, account, body)
  928. require.NoError(t, err)
  929. require.NotNil(t, result)
  930. require.Equal(t, "codex-tui", captureDialer.lastHeaders.Get("originator"))
  931. require.Equal(t,
  932. "codex-tui/"+codexCLIVersion+" (Mac OS X 15.1.0; arm64) iTerm.app",
  933. captureDialer.lastHeaders.Get("user-agent"),
  934. )
  935. require.Equal(t, codexCLIVersion, captureDialer.lastHeaders.Get("version"))
  936. }
  937. func TestOpenAIGatewayService_Forward_WSv2_HeaderSessionFallbackFromPromptCacheKey(t *testing.T) {
  938. gin.SetMode(gin.TestMode)
  939. rec := httptest.NewRecorder()
  940. c, _ := gin.CreateTestContext(rec)
  941. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  942. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
  943. cfg := &config.Config{}
  944. cfg.Security.URLAllowlist.Enabled = false
  945. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  946. cfg.Gateway.OpenAIWS.Enabled = true
  947. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  948. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  949. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  950. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  951. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  952. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  953. captureConn := &openAIWSCaptureConn{
  954. events: [][]byte{
  955. []byte(`{"type":"response.completed","response":{"id":"resp_prompt_cache_key","model":"gpt-5.1","usage":{"input_tokens":2,"output_tokens":1}}}`),
  956. },
  957. }
  958. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  959. pool := newOpenAIWSConnPool(cfg)
  960. pool.setClientDialerForTest(captureDialer)
  961. svc := &OpenAIGatewayService{
  962. cfg: cfg,
  963. httpUpstream: &httpUpstreamRecorder{},
  964. cache: &stubGatewayCache{},
  965. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  966. toolCorrector: NewCodexToolCorrector(),
  967. openaiWSPool: pool,
  968. }
  969. account := &Account{
  970. ID: 31,
  971. Name: "openai-oauth",
  972. Platform: PlatformOpenAI,
  973. Type: AccountTypeOAuth,
  974. Status: StatusActive,
  975. Schedulable: true,
  976. Concurrency: 1,
  977. Credentials: map[string]any{
  978. "access_token": "oauth-token-1",
  979. },
  980. Extra: map[string]any{
  981. "responses_websockets_v2_enabled": true,
  982. },
  983. }
  984. body := []byte(`{"model":"gpt-5.1","stream":true,"prompt_cache_key":"pcache_123","input":[{"type":"input_text","text":"hi"}]}`)
  985. result, err := svc.Forward(context.Background(), c, account, body)
  986. require.NoError(t, err)
  987. require.NotNil(t, result)
  988. require.Equal(t, "resp_prompt_cache_key", result.RequestID)
  989. // OAuth 账号的 session_id 应同时按 API key 和上游账号隔离(apiKeyID=0)。
  990. require.Equal(t, isolateOpenAIUpstreamSessionID(0, account, "pcache_123"), captureDialer.lastHeaders.Get("session_id"))
  991. require.Empty(t, captureDialer.lastHeaders.Get("conversation_id"))
  992. require.NotNil(t, captureConn.lastWrite)
  993. require.True(t, gjson.Get(requestToJSONString(captureConn.lastWrite), "stream").Exists())
  994. }
  995. func TestOpenAIGatewayService_Forward_WSv2_CodexFingerprintHandshakeBodyParityAndDefaultCacheKey(t *testing.T) {
  996. gin.SetMode(gin.TestMode)
  997. rec := httptest.NewRecorder()
  998. c, _ := gin.CreateTestContext(rec)
  999. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1000. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
  1001. c.Request.Header.Set("originator", "codex_cli_rs")
  1002. c.Request.Header.Set("session-id", "header-session")
  1003. c.Request.Header.Set("x-codex-turn-metadata", `{"installation_id":"header-install","session_id":"header-session","thread_id":"header-thread","turn_id":"header-turn","window_id":"header-window","sandbox":"seatbelt"}`)
  1004. cfg := &config.Config{}
  1005. cfg.Security.URLAllowlist.Enabled = false
  1006. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1007. cfg.Gateway.OpenAIWS.Enabled = true
  1008. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1009. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1010. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1011. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1012. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1013. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1014. captureConn := &openAIWSCaptureConn{
  1015. events: [][]byte{
  1016. []byte(`{"type":"response.completed","response":{"id":"resp_ws_fingerprint","model":"gpt-5.2","usage":{"input_tokens":2,"output_tokens":1}}}`),
  1017. },
  1018. }
  1019. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1020. pool := newOpenAIWSConnPool(cfg)
  1021. pool.setClientDialerForTest(captureDialer)
  1022. svc := &OpenAIGatewayService{
  1023. cfg: cfg,
  1024. httpUpstream: &httpUpstreamRecorder{},
  1025. cache: &stubGatewayCache{},
  1026. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1027. toolCorrector: NewCodexToolCorrector(),
  1028. openaiWSPool: pool,
  1029. }
  1030. account := newTestOAuthAccount(4405, map[string]any{
  1031. codexFingerprintModeExtraKey: "session",
  1032. "responses_websockets_v2_enabled": true,
  1033. })
  1034. account.Name = "oauth-ws-fingerprint"
  1035. account.Status = StatusActive
  1036. account.Schedulable = true
  1037. account.Concurrency = 1
  1038. account.Credentials = map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"}
  1039. body := []byte(`{"model":"gpt-5.2","stream":true,"prompt_cache_key":"body-session","client_metadata":{"session_id":"body-session","x-codex-turn-metadata":"{\"installation_id\":\"body-install\",\"session_id\":\"body-session\",\"thread_id\":\"body-thread\",\"turn_id\":\"body-turn\",\"window_id\":\"body-window\",\"sandbox\":\"seatbelt\"}"},"input":[{"type":"input_text","text":"hi"}]}`)
  1040. result, err := svc.Forward(context.Background(), c, account, body)
  1041. require.NoError(t, err)
  1042. require.NotNil(t, result)
  1043. require.Equal(t, "resp_ws_fingerprint", result.RequestID)
  1044. require.NotNil(t, captureConn.lastWrite)
  1045. seed, ok := codexFingerprintSeed(account.Extra)
  1046. require.True(t, ok)
  1047. wantInstall := resolveConvergedInstallationID(account, seed)
  1048. wantSession := resolveConvergedSessionID(seed)
  1049. wantThread := resolveConvergedThreadID(seed, "header-session")
  1050. payloadJSON := requestToJSONString(captureConn.lastWrite)
  1051. require.Equal(t, wantInstall, captureDialer.lastHeaders.Get("x-codex-installation-id"))
  1052. require.Equal(t, wantSession, captureDialer.lastHeaders.Get("session-id"))
  1053. require.Equal(t, wantSession, captureDialer.lastHeaders.Get("session_id"))
  1054. require.Equal(t, wantThread, captureDialer.lastHeaders.Get("thread-id"))
  1055. require.Equal(t, wantThread, captureDialer.lastHeaders.Get("x-client-request-id"))
  1056. require.Equal(t, wantThread+":0", captureDialer.lastHeaders.Get("x-codex-window-id"))
  1057. require.Equal(t, wantSession, gjson.Get(payloadJSON, "prompt_cache_key").String())
  1058. require.Equal(t, wantInstall, gjson.Get(payloadJSON, "client_metadata.x-codex-installation-id").String())
  1059. require.Equal(t, wantSession, gjson.Get(payloadJSON, "client_metadata.session_id").String())
  1060. require.Equal(t, wantThread, gjson.Get(payloadJSON, "client_metadata.thread_id").String())
  1061. require.Equal(t, wantThread+":0", gjson.Get(payloadJSON, "client_metadata.x-codex-window-id").String())
  1062. bodyTurnMetadata := gjson.Get(payloadJSON, "client_metadata.x-codex-turn-metadata").String()
  1063. headerTurnMetadata := captureDialer.lastHeaders.Get("x-codex-turn-metadata")
  1064. require.Equal(t, wantInstall, gjson.Get(bodyTurnMetadata, "installation_id").String())
  1065. require.Equal(t, wantSession, gjson.Get(bodyTurnMetadata, "session_id").String())
  1066. require.Equal(t, wantThread, gjson.Get(bodyTurnMetadata, "thread_id").String())
  1067. require.Equal(t, wantSession, gjson.Get(headerTurnMetadata, "session_id").String())
  1068. require.Equal(t, gjson.Get(bodyTurnMetadata, "turn_id").String(), gjson.Get(headerTurnMetadata, "turn_id").String())
  1069. require.NotZero(t, gjson.Get(bodyTurnMetadata, "turn_started_at_unix_ms").Int())
  1070. require.Equal(t,
  1071. gjson.Get(bodyTurnMetadata, "turn_started_at_unix_ms").Int(),
  1072. gjson.Get(headerTurnMetadata, "turn_started_at_unix_ms").Int(),
  1073. )
  1074. }
  1075. func TestOpenAIGatewayService_Forward_WSv2_ResponseDoneUsageParsed(t *testing.T) {
  1076. gin.SetMode(gin.TestMode)
  1077. rec := httptest.NewRecorder()
  1078. c, _ := gin.CreateTestContext(rec)
  1079. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1080. c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
  1081. cfg := &config.Config{}
  1082. cfg.Security.URLAllowlist.Enabled = false
  1083. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1084. cfg.Gateway.OpenAIWS.Enabled = true
  1085. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1086. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1087. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1088. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1089. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1090. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1091. captureConn := &openAIWSCaptureConn{
  1092. events: [][]byte{
  1093. []byte(`{"type":"response.done","response":{"id":"resp_done_usage","model":"gpt-5.1","usage":{"input_tokens":13,"output_tokens":8,"input_tokens_details":{"cached_tokens":5},"cache_creation_input_tokens":2,"output_tokens_details":{"image_tokens":4}}}}`),
  1094. },
  1095. }
  1096. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1097. pool := newOpenAIWSConnPool(cfg)
  1098. pool.setClientDialerForTest(captureDialer)
  1099. svc := &OpenAIGatewayService{
  1100. cfg: cfg,
  1101. httpUpstream: &httpUpstreamRecorder{},
  1102. cache: &stubGatewayCache{},
  1103. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1104. toolCorrector: NewCodexToolCorrector(),
  1105. openaiWSPool: pool,
  1106. }
  1107. account := &Account{
  1108. ID: 32,
  1109. Name: "openai-ws-done",
  1110. Platform: PlatformOpenAI,
  1111. Type: AccountTypeAPIKey,
  1112. Status: StatusActive,
  1113. Schedulable: true,
  1114. Concurrency: 1,
  1115. Credentials: map[string]any{
  1116. "api_key": "sk-test",
  1117. },
  1118. Extra: map[string]any{
  1119. "responses_websockets_v2_enabled": true,
  1120. },
  1121. }
  1122. body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hi"}]}`)
  1123. result, err := svc.Forward(context.Background(), c, account, body)
  1124. require.NoError(t, err)
  1125. require.NotNil(t, result)
  1126. require.Equal(t, "resp_done_usage", result.RequestID)
  1127. require.Equal(t, 13, result.Usage.InputTokens)
  1128. require.Equal(t, 8, result.Usage.OutputTokens)
  1129. require.Equal(t, 5, result.Usage.CacheReadInputTokens)
  1130. require.Equal(t, 2, result.Usage.CacheCreationInputTokens)
  1131. require.Equal(t, 4, result.Usage.ImageOutputTokens)
  1132. }
  1133. func TestOpenAIGatewayService_Forward_WSv1_Unsupported(t *testing.T) {
  1134. gin.SetMode(gin.TestMode)
  1135. rec := httptest.NewRecorder()
  1136. c, _ := gin.CreateTestContext(rec)
  1137. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1138. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
  1139. cfg := &config.Config{}
  1140. cfg.Security.URLAllowlist.Enabled = false
  1141. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1142. cfg.Gateway.OpenAIWS.Enabled = true
  1143. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1144. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1145. cfg.Gateway.OpenAIWS.ResponsesWebsockets = true
  1146. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = false
  1147. upstream := &httpUpstreamRecorder{
  1148. resp: &http.Response{
  1149. StatusCode: http.StatusOK,
  1150. Header: http.Header{"Content-Type": []string{"application/json"}},
  1151. Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":1}}`)),
  1152. },
  1153. }
  1154. svc := &OpenAIGatewayService{
  1155. cfg: cfg,
  1156. httpUpstream: upstream,
  1157. cache: &stubGatewayCache{},
  1158. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1159. toolCorrector: NewCodexToolCorrector(),
  1160. }
  1161. account := &Account{
  1162. ID: 39,
  1163. Name: "openai-ws-v1",
  1164. Platform: PlatformOpenAI,
  1165. Type: AccountTypeAPIKey,
  1166. Status: StatusActive,
  1167. Schedulable: true,
  1168. Concurrency: 1,
  1169. Credentials: map[string]any{
  1170. "api_key": "sk-test",
  1171. "base_url": "https://api.openai.com/v1/responses",
  1172. },
  1173. Extra: map[string]any{
  1174. "responses_websockets_v2_enabled": true,
  1175. },
  1176. }
  1177. body := []byte(`{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_prev_v1","input":[{"type":"input_text","text":"hello"}]}`)
  1178. result, err := svc.Forward(context.Background(), c, account, body)
  1179. require.Error(t, err)
  1180. require.Nil(t, result)
  1181. require.Contains(t, err.Error(), "ws v1")
  1182. require.Equal(t, http.StatusBadRequest, rec.Code)
  1183. require.Contains(t, rec.Body.String(), "WSv1")
  1184. require.Nil(t, upstream.lastReq, "WSv1 不支持时不应触发 HTTP 上游请求")
  1185. }
  1186. func TestOpenAIGatewayService_Forward_WSv2_TurnStateAndMetadataReplayOnReconnect(t *testing.T) {
  1187. gin.SetMode(gin.TestMode)
  1188. var connIndex atomic.Int64
  1189. headersCh := make(chan http.Header, 4)
  1190. upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
  1191. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1192. idx := connIndex.Add(1)
  1193. headersCh <- cloneHeader(r.Header)
  1194. respHeader := http.Header{}
  1195. if idx == 1 {
  1196. respHeader.Set("x-codex-turn-state", "turn_state_first")
  1197. }
  1198. conn, err := upgrader.Upgrade(w, r, respHeader)
  1199. if err != nil {
  1200. t.Errorf("upgrade websocket failed: %v", err)
  1201. return
  1202. }
  1203. defer func() {
  1204. _ = conn.Close()
  1205. }()
  1206. var request map[string]any
  1207. if err := conn.ReadJSON(&request); err != nil {
  1208. t.Errorf("read ws request failed: %v", err)
  1209. return
  1210. }
  1211. responseID := "resp_turn_" + strconv.FormatInt(idx, 10)
  1212. if err := conn.WriteJSON(map[string]any{
  1213. "type": "response.completed",
  1214. "response": map[string]any{
  1215. "id": responseID,
  1216. "model": "gpt-5.1",
  1217. "usage": map[string]any{
  1218. "input_tokens": 2,
  1219. "output_tokens": 1,
  1220. },
  1221. },
  1222. }); err != nil {
  1223. t.Errorf("write response.completed failed: %v", err)
  1224. return
  1225. }
  1226. }))
  1227. defer wsServer.Close()
  1228. cfg := &config.Config{}
  1229. cfg.Security.URLAllowlist.Enabled = false
  1230. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1231. cfg.Gateway.OpenAIWS.Enabled = true
  1232. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1233. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1234. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1235. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1236. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1237. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 0
  1238. svc := &OpenAIGatewayService{
  1239. cfg: cfg,
  1240. httpUpstream: &httpUpstreamRecorder{},
  1241. cache: &stubGatewayCache{},
  1242. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1243. toolCorrector: NewCodexToolCorrector(),
  1244. }
  1245. account := &Account{
  1246. ID: 49,
  1247. Name: "openai-turn-state",
  1248. Platform: PlatformOpenAI,
  1249. Type: AccountTypeAPIKey,
  1250. Status: StatusActive,
  1251. Schedulable: true,
  1252. Concurrency: 1,
  1253. Credentials: map[string]any{
  1254. "api_key": "sk-test",
  1255. "base_url": wsServer.URL,
  1256. },
  1257. Extra: map[string]any{
  1258. "responses_websockets_v2_enabled": true,
  1259. },
  1260. }
  1261. reqBody := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
  1262. rec1 := httptest.NewRecorder()
  1263. c1, _ := gin.CreateTestContext(rec1)
  1264. c1.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1265. c1.Request.Header.Set("session_id", "session_turn_state")
  1266. c1.Request.Header.Set("x-codex-turn-metadata", "turn_meta_1")
  1267. result1, err := svc.Forward(context.Background(), c1, account, reqBody)
  1268. require.NoError(t, err)
  1269. require.NotNil(t, result1)
  1270. sessionHash := svc.GenerateSessionHash(c1, reqBody)
  1271. store := svc.getOpenAIWSStateStore()
  1272. turnState, ok := store.GetSessionTurnState(0, sessionHash)
  1273. require.True(t, ok)
  1274. require.Equal(t, "turn_state_first", turnState)
  1275. // 主动淘汰连接,模拟下一次请求发生重连。
  1276. connID, hasConn := store.GetResponseConn(result1.RequestID)
  1277. require.True(t, hasConn)
  1278. svc.getOpenAIWSConnPool().evictConn(account.ID, connID)
  1279. rec2 := httptest.NewRecorder()
  1280. c2, _ := gin.CreateTestContext(rec2)
  1281. c2.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1282. c2.Request.Header.Set("session_id", "session_turn_state")
  1283. c2.Request.Header.Set("x-codex-turn-metadata", "turn_meta_2")
  1284. result2, err := svc.Forward(context.Background(), c2, account, reqBody)
  1285. require.NoError(t, err)
  1286. require.NotNil(t, result2)
  1287. firstHandshakeHeaders := <-headersCh
  1288. secondHandshakeHeaders := <-headersCh
  1289. require.Equal(t, "turn_meta_1", firstHandshakeHeaders.Get("X-Codex-Turn-Metadata"))
  1290. require.Equal(t, "turn_meta_2", secondHandshakeHeaders.Get("X-Codex-Turn-Metadata"))
  1291. require.Equal(t, "turn_state_first", secondHandshakeHeaders.Get("X-Codex-Turn-State"))
  1292. }
  1293. func TestOpenAIGatewayService_Forward_WSv2_GeneratePrewarm(t *testing.T) {
  1294. gin.SetMode(gin.TestMode)
  1295. rec := httptest.NewRecorder()
  1296. c, _ := gin.CreateTestContext(rec)
  1297. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1298. c.Request.Header.Set("session_id", "session-prewarm")
  1299. cfg := &config.Config{}
  1300. cfg.Security.URLAllowlist.Enabled = false
  1301. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1302. cfg.Gateway.OpenAIWS.Enabled = true
  1303. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1304. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1305. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1306. cfg.Gateway.OpenAIWS.PrewarmGenerateEnabled = true
  1307. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1308. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1309. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1310. captureConn := &openAIWSCaptureConn{
  1311. events: [][]byte{
  1312. []byte(`{"type":"response.completed","response":{"id":"resp_prewarm_1","model":"gpt-5.1","usage":{"input_tokens":0,"output_tokens":0}}}`),
  1313. []byte(`{"type":"response.completed","response":{"id":"resp_main_1","model":"gpt-5.1","usage":{"input_tokens":4,"output_tokens":2}}}`),
  1314. },
  1315. }
  1316. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1317. pool := newOpenAIWSConnPool(cfg)
  1318. pool.setClientDialerForTest(captureDialer)
  1319. svc := &OpenAIGatewayService{
  1320. cfg: cfg,
  1321. httpUpstream: &httpUpstreamRecorder{},
  1322. cache: &stubGatewayCache{},
  1323. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1324. toolCorrector: NewCodexToolCorrector(),
  1325. openaiWSPool: pool,
  1326. }
  1327. account := &Account{
  1328. ID: 59,
  1329. Name: "openai-prewarm",
  1330. Platform: PlatformOpenAI,
  1331. Type: AccountTypeAPIKey,
  1332. Status: StatusActive,
  1333. Schedulable: true,
  1334. Concurrency: 1,
  1335. Credentials: map[string]any{
  1336. "api_key": "sk-test",
  1337. },
  1338. Extra: map[string]any{
  1339. "responses_websockets_v2_enabled": true,
  1340. },
  1341. }
  1342. body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
  1343. result, err := svc.Forward(context.Background(), c, account, body)
  1344. require.NoError(t, err)
  1345. require.NotNil(t, result)
  1346. require.Equal(t, "resp_main_1", result.RequestID)
  1347. require.Len(t, captureConn.writes, 2, "开启 generate=false 预热后应发送两次 WS 请求")
  1348. firstWrite := requestToJSONString(captureConn.writes[0])
  1349. secondWrite := requestToJSONString(captureConn.writes[1])
  1350. require.True(t, gjson.Get(firstWrite, "generate").Exists())
  1351. require.False(t, gjson.Get(firstWrite, "generate").Bool())
  1352. require.False(t, gjson.Get(secondWrite, "generate").Exists())
  1353. }
  1354. func TestOpenAIGatewayService_PrewarmReadHonorsParentContext(t *testing.T) {
  1355. cfg := &config.Config{}
  1356. cfg.Gateway.OpenAIWS.PrewarmGenerateEnabled = true
  1357. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5
  1358. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1359. svc := &OpenAIGatewayService{
  1360. cfg: cfg,
  1361. toolCorrector: NewCodexToolCorrector(),
  1362. }
  1363. account := &Account{
  1364. ID: 601,
  1365. Name: "openai-prewarm-timeout",
  1366. Platform: PlatformOpenAI,
  1367. Type: AccountTypeAPIKey,
  1368. Status: StatusActive,
  1369. Schedulable: true,
  1370. }
  1371. conn := newOpenAIWSConn("prewarm_ctx_conn", account.ID, &openAIWSBlockingConn{
  1372. readDelay: 200 * time.Millisecond,
  1373. }, nil)
  1374. lease := &openAIWSConnLease{
  1375. accountID: account.ID,
  1376. conn: conn,
  1377. }
  1378. payload := map[string]any{
  1379. "type": "response.create",
  1380. "model": "gpt-5.1",
  1381. }
  1382. ctx, cancel := context.WithTimeout(context.Background(), 40*time.Millisecond)
  1383. defer cancel()
  1384. start := time.Now()
  1385. err := svc.performOpenAIWSGeneratePrewarm(
  1386. ctx,
  1387. lease,
  1388. OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2},
  1389. payload,
  1390. "",
  1391. map[string]any{"model": "gpt-5.1"},
  1392. account,
  1393. nil,
  1394. 0,
  1395. )
  1396. elapsed := time.Since(start)
  1397. require.Error(t, err)
  1398. require.Contains(t, err.Error(), "prewarm_read_event")
  1399. require.Less(t, elapsed, 180*time.Millisecond, "预热读取应受父 context 取消控制,不应阻塞到 read_timeout")
  1400. }
  1401. func TestOpenAIGatewayService_Forward_WSv2_TurnMetadataInPayloadOnConnReuse(t *testing.T) {
  1402. gin.SetMode(gin.TestMode)
  1403. cfg := &config.Config{}
  1404. cfg.Security.URLAllowlist.Enabled = false
  1405. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1406. cfg.Gateway.OpenAIWS.Enabled = true
  1407. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1408. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1409. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1410. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1411. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1412. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1413. captureConn := &openAIWSCaptureConn{
  1414. events: [][]byte{
  1415. []byte(`{"type":"response.completed","response":{"id":"resp_meta_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1416. []byte(`{"type":"response.completed","response":{"id":"resp_meta_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
  1417. },
  1418. }
  1419. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1420. pool := newOpenAIWSConnPool(cfg)
  1421. pool.setClientDialerForTest(captureDialer)
  1422. svc := &OpenAIGatewayService{
  1423. cfg: cfg,
  1424. httpUpstream: &httpUpstreamRecorder{},
  1425. cache: &stubGatewayCache{},
  1426. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1427. toolCorrector: NewCodexToolCorrector(),
  1428. openaiWSPool: pool,
  1429. }
  1430. account := &Account{
  1431. ID: 69,
  1432. Name: "openai-turn-metadata",
  1433. Platform: PlatformOpenAI,
  1434. Type: AccountTypeAPIKey,
  1435. Status: StatusActive,
  1436. Schedulable: true,
  1437. Concurrency: 1,
  1438. Credentials: map[string]any{
  1439. "api_key": "sk-test",
  1440. },
  1441. Extra: map[string]any{
  1442. "responses_websockets_v2_enabled": true,
  1443. },
  1444. }
  1445. body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
  1446. rec1 := httptest.NewRecorder()
  1447. c1, _ := gin.CreateTestContext(rec1)
  1448. c1.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1449. c1.Request.Header.Set("session_id", "session-metadata-reuse")
  1450. c1.Request.Header.Set("x-codex-turn-metadata", "turn_meta_payload_1")
  1451. result1, err := svc.Forward(context.Background(), c1, account, body)
  1452. require.NoError(t, err)
  1453. require.NotNil(t, result1)
  1454. require.Equal(t, "resp_meta_1", result1.RequestID)
  1455. require.Len(t, captureConn.writes, 1)
  1456. firstWrite := requestToJSONString(captureConn.writes[0])
  1457. require.Equal(t, "turn_meta_payload_1", gjson.Get(firstWrite, "client_metadata.x-codex-turn-metadata").String())
  1458. rec2 := httptest.NewRecorder()
  1459. c2, _ := gin.CreateTestContext(rec2)
  1460. c2.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1461. c2.Request.Header.Set("session_id", "session-metadata-reuse")
  1462. c2.Request.Header.Set("x-codex-turn-metadata", "turn_meta_payload_2")
  1463. result2, err := svc.Forward(context.Background(), c2, account, body)
  1464. require.NoError(t, err)
  1465. require.NotNil(t, result2)
  1466. require.Equal(t, "resp_meta_2", result2.RequestID)
  1467. require.Equal(t, 1, captureDialer.DialCount(), "同一账号两轮请求应复用同一 WS 连接")
  1468. require.Len(t, captureConn.writes, 2)
  1469. firstWrite = requestToJSONString(captureConn.writes[0])
  1470. secondWrite := requestToJSONString(captureConn.writes[1])
  1471. require.Equal(t, "turn_meta_payload_1", gjson.Get(firstWrite, "client_metadata.x-codex-turn-metadata").String())
  1472. require.Equal(t, "turn_meta_payload_2", gjson.Get(secondWrite, "client_metadata.x-codex-turn-metadata").String())
  1473. }
  1474. func TestOpenAIGatewayService_Forward_WSv2StoreFalseSessionConnIsolation(t *testing.T) {
  1475. gin.SetMode(gin.TestMode)
  1476. var upgradeCount atomic.Int64
  1477. var sequence atomic.Int64
  1478. upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
  1479. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1480. upgradeCount.Add(1)
  1481. conn, err := upgrader.Upgrade(w, r, nil)
  1482. if err != nil {
  1483. t.Errorf("upgrade websocket failed: %v", err)
  1484. return
  1485. }
  1486. defer func() {
  1487. _ = conn.Close()
  1488. }()
  1489. for {
  1490. var request map[string]any
  1491. if err := conn.ReadJSON(&request); err != nil {
  1492. return
  1493. }
  1494. responseID := "resp_store_false_" + strconv.FormatInt(sequence.Add(1), 10)
  1495. if err := conn.WriteJSON(map[string]any{
  1496. "type": "response.completed",
  1497. "response": map[string]any{
  1498. "id": responseID,
  1499. "model": "gpt-5.1",
  1500. "usage": map[string]any{
  1501. "input_tokens": 1,
  1502. "output_tokens": 1,
  1503. },
  1504. },
  1505. }); err != nil {
  1506. return
  1507. }
  1508. }
  1509. }))
  1510. defer wsServer.Close()
  1511. cfg := &config.Config{}
  1512. cfg.Security.URLAllowlist.Enabled = false
  1513. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1514. cfg.Gateway.OpenAIWS.Enabled = true
  1515. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1516. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1517. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1518. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
  1519. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1520. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 4
  1521. cfg.Gateway.OpenAIWS.StoreDisabledForceNewConn = true
  1522. svc := &OpenAIGatewayService{
  1523. cfg: cfg,
  1524. httpUpstream: &httpUpstreamRecorder{},
  1525. cache: &stubGatewayCache{},
  1526. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1527. toolCorrector: NewCodexToolCorrector(),
  1528. }
  1529. account := &Account{
  1530. ID: 79,
  1531. Name: "openai-store-false",
  1532. Platform: PlatformOpenAI,
  1533. Type: AccountTypeAPIKey,
  1534. Status: StatusActive,
  1535. Schedulable: true,
  1536. Concurrency: 2,
  1537. Credentials: map[string]any{
  1538. "api_key": "sk-test",
  1539. "base_url": wsServer.URL,
  1540. },
  1541. Extra: map[string]any{
  1542. "responses_websockets_v2_enabled": true,
  1543. },
  1544. }
  1545. body := []byte(`{"model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  1546. rec1 := httptest.NewRecorder()
  1547. c1, _ := gin.CreateTestContext(rec1)
  1548. c1.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1549. c1.Request.Header.Set("session_id", "session_store_false_a")
  1550. result1, err := svc.Forward(context.Background(), c1, account, body)
  1551. require.NoError(t, err)
  1552. require.NotNil(t, result1)
  1553. require.Equal(t, int64(1), upgradeCount.Load())
  1554. rec2 := httptest.NewRecorder()
  1555. c2, _ := gin.CreateTestContext(rec2)
  1556. c2.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1557. c2.Request.Header.Set("session_id", "session_store_false_a")
  1558. result2, err := svc.Forward(context.Background(), c2, account, body)
  1559. require.NoError(t, err)
  1560. require.NotNil(t, result2)
  1561. require.Equal(t, int64(1), upgradeCount.Load(), "同一 session(store=false) 应复用同一 WS 连接")
  1562. rec3 := httptest.NewRecorder()
  1563. c3, _ := gin.CreateTestContext(rec3)
  1564. c3.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1565. c3.Request.Header.Set("session_id", "session_store_false_b")
  1566. result3, err := svc.Forward(context.Background(), c3, account, body)
  1567. require.NoError(t, err)
  1568. require.NotNil(t, result3)
  1569. require.Equal(t, int64(2), upgradeCount.Load(), "不同 session(store=false) 应隔离连接,避免续链状态互相覆盖")
  1570. }
  1571. func TestOpenAIGatewayService_Forward_WSv2StoreFalseDisableForceNewConnAllowsReuse(t *testing.T) {
  1572. gin.SetMode(gin.TestMode)
  1573. var upgradeCount atomic.Int64
  1574. var sequence atomic.Int64
  1575. upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
  1576. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1577. upgradeCount.Add(1)
  1578. conn, err := upgrader.Upgrade(w, r, nil)
  1579. if err != nil {
  1580. t.Errorf("upgrade websocket failed: %v", err)
  1581. return
  1582. }
  1583. defer func() {
  1584. _ = conn.Close()
  1585. }()
  1586. for {
  1587. var request map[string]any
  1588. if err := conn.ReadJSON(&request); err != nil {
  1589. return
  1590. }
  1591. responseID := "resp_store_false_reuse_" + strconv.FormatInt(sequence.Add(1), 10)
  1592. if err := conn.WriteJSON(map[string]any{
  1593. "type": "response.completed",
  1594. "response": map[string]any{
  1595. "id": responseID,
  1596. "model": "gpt-5.1",
  1597. "usage": map[string]any{
  1598. "input_tokens": 1,
  1599. "output_tokens": 1,
  1600. },
  1601. },
  1602. }); err != nil {
  1603. return
  1604. }
  1605. }
  1606. }))
  1607. defer wsServer.Close()
  1608. cfg := &config.Config{}
  1609. cfg.Security.URLAllowlist.Enabled = false
  1610. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1611. cfg.Gateway.OpenAIWS.Enabled = true
  1612. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1613. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1614. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1615. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1616. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1617. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1618. cfg.Gateway.OpenAIWS.StoreDisabledForceNewConn = false
  1619. svc := &OpenAIGatewayService{
  1620. cfg: cfg,
  1621. httpUpstream: &httpUpstreamRecorder{},
  1622. cache: &stubGatewayCache{},
  1623. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1624. toolCorrector: NewCodexToolCorrector(),
  1625. }
  1626. account := &Account{
  1627. ID: 80,
  1628. Name: "openai-store-false-reuse",
  1629. Platform: PlatformOpenAI,
  1630. Type: AccountTypeAPIKey,
  1631. Status: StatusActive,
  1632. Schedulable: true,
  1633. Concurrency: 2,
  1634. Credentials: map[string]any{
  1635. "api_key": "sk-test",
  1636. "base_url": wsServer.URL,
  1637. },
  1638. Extra: map[string]any{
  1639. "responses_websockets_v2_enabled": true,
  1640. },
  1641. }
  1642. body := []byte(`{"model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
  1643. rec1 := httptest.NewRecorder()
  1644. c1, _ := gin.CreateTestContext(rec1)
  1645. c1.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1646. c1.Request.Header.Set("session_id", "session_store_false_reuse_a")
  1647. result1, err := svc.Forward(context.Background(), c1, account, body)
  1648. require.NoError(t, err)
  1649. require.NotNil(t, result1)
  1650. require.Equal(t, int64(1), upgradeCount.Load())
  1651. rec2 := httptest.NewRecorder()
  1652. c2, _ := gin.CreateTestContext(rec2)
  1653. c2.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1654. c2.Request.Header.Set("session_id", "session_store_false_reuse_b")
  1655. result2, err := svc.Forward(context.Background(), c2, account, body)
  1656. require.NoError(t, err)
  1657. require.NotNil(t, result2)
  1658. require.Equal(t, int64(1), upgradeCount.Load(), "关闭强制新连后,不同 session(store=false) 可复用连接")
  1659. }
  1660. func TestOpenAIGatewayService_Forward_WSv2ReadTimeoutAppliesPerRead(t *testing.T) {
  1661. gin.SetMode(gin.TestMode)
  1662. rec := httptest.NewRecorder()
  1663. c, _ := gin.CreateTestContext(rec)
  1664. c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
  1665. c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
  1666. cfg := &config.Config{}
  1667. cfg.Security.URLAllowlist.Enabled = false
  1668. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1669. cfg.Gateway.OpenAIWS.Enabled = true
  1670. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1671. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1672. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1673. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1674. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1675. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1676. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  1677. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  1678. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 1
  1679. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1680. captureConn := &openAIWSCaptureConn{
  1681. readDelays: []time.Duration{
  1682. 700 * time.Millisecond,
  1683. 700 * time.Millisecond,
  1684. },
  1685. events: [][]byte{
  1686. []byte(`{"type":"response.created","response":{"id":"resp_timeout_ok","model":"gpt-5.1"}}`),
  1687. []byte(`{"type":"response.completed","response":{"id":"resp_timeout_ok","model":"gpt-5.1","usage":{"input_tokens":2,"output_tokens":1}}}`),
  1688. },
  1689. }
  1690. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1691. pool := newOpenAIWSConnPool(cfg)
  1692. pool.setClientDialerForTest(captureDialer)
  1693. upstream := &httpUpstreamRecorder{
  1694. resp: &http.Response{
  1695. StatusCode: http.StatusOK,
  1696. Header: http.Header{"Content-Type": []string{"application/json"}},
  1697. Body: io.NopCloser(strings.NewReader(`{"id":"resp_http_fallback","usage":{"input_tokens":1,"output_tokens":1}}`)),
  1698. },
  1699. }
  1700. svc := &OpenAIGatewayService{
  1701. cfg: cfg,
  1702. httpUpstream: upstream,
  1703. cache: &stubGatewayCache{},
  1704. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1705. toolCorrector: NewCodexToolCorrector(),
  1706. openaiWSPool: pool,
  1707. }
  1708. account := &Account{
  1709. ID: 81,
  1710. Name: "openai-read-timeout",
  1711. Platform: PlatformOpenAI,
  1712. Type: AccountTypeAPIKey,
  1713. Status: StatusActive,
  1714. Schedulable: true,
  1715. Concurrency: 1,
  1716. Credentials: map[string]any{
  1717. "api_key": "sk-test",
  1718. },
  1719. Extra: map[string]any{
  1720. "responses_websockets_v2_enabled": true,
  1721. },
  1722. }
  1723. body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
  1724. result, err := svc.Forward(context.Background(), c, account, body)
  1725. require.NoError(t, err)
  1726. require.NotNil(t, result)
  1727. require.Equal(t, "resp_timeout_ok", result.RequestID)
  1728. require.Nil(t, upstream.lastReq, "每次 Read 都应独立应用超时;总时长超过 read_timeout 不应误回退 HTTP")
  1729. }
  1730. type openAIWSCaptureDialer struct {
  1731. mu sync.Mutex
  1732. conn *openAIWSCaptureConn
  1733. lastHeaders http.Header
  1734. handshake http.Header
  1735. dialCount int
  1736. }
  1737. func (d *openAIWSCaptureDialer) Dial(
  1738. ctx context.Context,
  1739. wsURL string,
  1740. headers http.Header,
  1741. proxyURL string,
  1742. ) (openAIWSClientConn, int, http.Header, error) {
  1743. _ = ctx
  1744. _ = wsURL
  1745. _ = proxyURL
  1746. d.mu.Lock()
  1747. d.lastHeaders = cloneHeader(headers)
  1748. d.dialCount++
  1749. respHeaders := cloneHeader(d.handshake)
  1750. d.mu.Unlock()
  1751. return d.conn, 0, respHeaders, nil
  1752. }
  1753. func (d *openAIWSCaptureDialer) DialCount() int {
  1754. d.mu.Lock()
  1755. defer d.mu.Unlock()
  1756. return d.dialCount
  1757. }
  1758. type openAIWSCaptureConn struct {
  1759. mu sync.Mutex
  1760. readDelays []time.Duration
  1761. events [][]byte
  1762. lastWrite map[string]any
  1763. writes []map[string]any
  1764. closed bool
  1765. }
  1766. func (c *openAIWSCaptureConn) WriteJSON(ctx context.Context, value any) error {
  1767. _ = ctx
  1768. c.mu.Lock()
  1769. defer c.mu.Unlock()
  1770. if c.closed {
  1771. return errOpenAIWSConnClosed
  1772. }
  1773. switch payload := value.(type) {
  1774. case map[string]any:
  1775. c.lastWrite = cloneMapStringAny(payload)
  1776. c.writes = append(c.writes, cloneMapStringAny(payload))
  1777. case json.RawMessage:
  1778. var parsed map[string]any
  1779. if err := decodeOpenAIJSONUseNumber(payload, &parsed); err == nil {
  1780. c.lastWrite = cloneMapStringAny(parsed)
  1781. c.writes = append(c.writes, cloneMapStringAny(parsed))
  1782. }
  1783. case []byte:
  1784. var parsed map[string]any
  1785. if err := decodeOpenAIJSONUseNumber(payload, &parsed); err == nil {
  1786. c.lastWrite = cloneMapStringAny(parsed)
  1787. c.writes = append(c.writes, cloneMapStringAny(parsed))
  1788. }
  1789. }
  1790. return nil
  1791. }
  1792. func (c *openAIWSCaptureConn) ReadMessage(ctx context.Context) ([]byte, error) {
  1793. if ctx == nil {
  1794. ctx = context.Background()
  1795. }
  1796. c.mu.Lock()
  1797. if c.closed {
  1798. c.mu.Unlock()
  1799. return nil, errOpenAIWSConnClosed
  1800. }
  1801. if len(c.events) == 0 {
  1802. c.mu.Unlock()
  1803. return nil, io.EOF
  1804. }
  1805. delay := time.Duration(0)
  1806. if len(c.readDelays) > 0 {
  1807. delay = c.readDelays[0]
  1808. c.readDelays = c.readDelays[1:]
  1809. }
  1810. event := c.events[0]
  1811. c.events = c.events[1:]
  1812. c.mu.Unlock()
  1813. if delay > 0 {
  1814. timer := time.NewTimer(delay)
  1815. defer timer.Stop()
  1816. select {
  1817. case <-ctx.Done():
  1818. return nil, ctx.Err()
  1819. case <-timer.C:
  1820. }
  1821. }
  1822. return event, nil
  1823. }
  1824. func (c *openAIWSCaptureConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) {
  1825. payload, err := c.ReadMessage(ctx)
  1826. if err != nil {
  1827. return coderws.MessageText, nil, err
  1828. }
  1829. return coderws.MessageText, payload, nil
  1830. }
  1831. func (c *openAIWSCaptureConn) WriteFrame(ctx context.Context, _ coderws.MessageType, payload []byte) error {
  1832. return c.WriteJSON(ctx, json.RawMessage(payload))
  1833. }
  1834. func (c *openAIWSCaptureConn) Ping(ctx context.Context) error {
  1835. _ = ctx
  1836. return nil
  1837. }
  1838. func (c *openAIWSCaptureConn) Close() error {
  1839. c.mu.Lock()
  1840. defer c.mu.Unlock()
  1841. c.closed = true
  1842. return nil
  1843. }
  1844. func cloneMapStringAny(src map[string]any) map[string]any {
  1845. if src == nil {
  1846. return nil
  1847. }
  1848. dst := make(map[string]any, len(src))
  1849. for k, v := range src {
  1850. dst[k] = v
  1851. }
  1852. return dst
  1853. }