| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052 |
- package service
- import (
- "context"
- "encoding/json"
- "errors"
- "io"
- "net/http"
- "net/http/httptest"
- "strconv"
- "strings"
- "sync"
- "sync/atomic"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
- coderws "github.com/coder/websocket"
- "github.com/gin-gonic/gin"
- "github.com/gorilla/websocket"
- "github.com/stretchr/testify/require"
- "github.com/tidwall/gjson"
- )
- func TestOpenAIGatewayService_Forward_WSv2_SuccessAndBindSticky(t *testing.T) {
- gin.SetMode(gin.TestMode)
- type receivedPayload struct {
- Type string
- PreviousResponseID string
- StreamExists bool
- Stream bool
- }
- receivedCh := make(chan receivedPayload, 1)
- upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
- wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- conn, err := upgrader.Upgrade(w, r, nil)
- if err != nil {
- t.Errorf("upgrade websocket failed: %v", err)
- return
- }
- defer func() {
- _ = conn.Close()
- }()
- var request map[string]any
- if err := conn.ReadJSON(&request); err != nil {
- t.Errorf("read ws request failed: %v", err)
- return
- }
- requestJSON := requestToJSONString(request)
- receivedCh <- receivedPayload{
- Type: strings.TrimSpace(gjson.Get(requestJSON, "type").String()),
- PreviousResponseID: strings.TrimSpace(gjson.Get(requestJSON, "previous_response_id").String()),
- StreamExists: gjson.Get(requestJSON, "stream").Exists(),
- Stream: gjson.Get(requestJSON, "stream").Bool(),
- }
- if err := conn.WriteJSON(map[string]any{
- "type": "response.created",
- "response": map[string]any{
- "id": "resp_new_1",
- "model": "gpt-5.1",
- },
- }); err != nil {
- t.Errorf("write response.created failed: %v", err)
- return
- }
- if err := conn.WriteJSON(map[string]any{
- "type": "response.completed",
- "response": map[string]any{
- "id": "resp_new_1",
- "model": "gpt-5.1",
- "usage": map[string]any{
- "input_tokens": 12,
- "output_tokens": 7,
- "input_tokens_details": map[string]any{
- "cached_tokens": 3,
- },
- },
- },
- }); err != nil {
- t.Errorf("write response.completed failed: %v", err)
- return
- }
- }))
- defer wsServer.Close()
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
- groupID := int64(1001)
- c.Set("api_key", &APIKey{GroupID: &groupID})
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
- cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
- cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
- cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 30
- cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 10
- cfg.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = 3600
- upstream := &httpUpstreamRecorder{
- resp: &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"application/json"}},
- Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":1}}`)),
- },
- }
- cache := &stubGatewayCache{}
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: upstream,
- cache: cache,
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- }
- account := &Account{
- ID: 9,
- Name: "openai-ws",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 2,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "base_url": wsServer.URL,
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_prev_1","input":[{"type":"input_text","text":"hello"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, 12, result.Usage.InputTokens)
- require.Equal(t, 7, result.Usage.OutputTokens)
- require.Equal(t, 3, result.Usage.CacheReadInputTokens)
- require.Equal(t, "resp_new_1", result.RequestID)
- require.True(t, result.OpenAIWSMode)
- require.False(t, gjson.GetBytes(upstream.lastBody, "model").Exists(), "WSv2 成功时不应回落 HTTP 上游")
- received := <-receivedCh
- require.Equal(t, "response.create", received.Type)
- require.Equal(t, "resp_prev_1", received.PreviousResponseID)
- require.True(t, received.StreamExists, "WS 请求应携带 stream 字段")
- require.False(t, received.Stream, "应保持客户端 stream=false 的原始语义")
- store := svc.getOpenAIWSStateStore()
- mappedAccountID, getErr := store.GetResponseAccount(context.Background(), groupID, "resp_new_1")
- require.NoError(t, getErr)
- require.Equal(t, account.ID, mappedAccountID)
- connID, ok := store.GetResponseConn("resp_new_1")
- require.True(t, ok)
- require.NotEmpty(t, connID)
- responseBody := rec.Body.Bytes()
- require.Equal(t, "resp_new_1", gjson.GetBytes(responseBody, "id").String())
- }
- func TestOpenAIGatewayService_Forward_WSv2_UsesPatchedBodyAfterValidationDecode(t *testing.T) {
- gin.SetMode(gin.TestMode)
- type receivedPayload struct {
- MaxCompletionTokensExists bool
- }
- receivedCh := make(chan receivedPayload, 1)
- upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
- wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- conn, err := upgrader.Upgrade(w, r, nil)
- if err != nil {
- t.Errorf("upgrade websocket failed: %v", err)
- return
- }
- defer func() { _ = conn.Close() }()
- var request map[string]any
- if err := conn.ReadJSON(&request); err != nil {
- t.Errorf("read ws request failed: %v", err)
- return
- }
- requestJSON := requestToJSONString(request)
- receivedCh <- receivedPayload{MaxCompletionTokensExists: gjson.Get(requestJSON, "max_completion_tokens").Exists()}
- if err := conn.WriteJSON(map[string]any{
- "type": "response.completed",
- "response": map[string]any{
- "id": "resp_patched_ws_1",
- "model": "gpt-5.3-codex-spark",
- "usage": map[string]any{"input_tokens": 1, "output_tokens": 1},
- },
- }); err != nil {
- t.Errorf("write response.completed failed: %v", err)
- return
- }
- }))
- defer wsServer.Close()
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
- cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 30
- cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 10
- svc := &OpenAIGatewayService{
- cfg: cfg,
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- }
- account := &Account{
- ID: 10,
- Name: "openai-ws",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "base_url": wsServer.URL,
- },
- Extra: map[string]any{"responses_websockets_v2_enabled": true},
- }
- body := []byte(`{"model":"gpt-5.4","stream":false,"max_completion_tokens":12,"tools":[{"type":"image_generation"}],"input":[{"type":"input_text","text":"hello"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.True(t, result.OpenAIWSMode)
- received := <-receivedCh
- require.False(t, received.MaxCompletionTokensExists)
- }
- func TestOpenAIGatewayService_Forward_WSv2_ImageGenerationCountsOutputs(t *testing.T) {
- gin.SetMode(gin.TestMode)
- upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
- wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- conn, err := upgrader.Upgrade(w, r, nil)
- if err != nil {
- t.Errorf("upgrade websocket failed: %v", err)
- return
- }
- defer func() {
- _ = conn.Close()
- }()
- var request map[string]any
- if err := conn.ReadJSON(&request); err != nil {
- t.Errorf("read ws request failed: %v", err)
- return
- }
- if err := conn.WriteJSON(map[string]any{
- "type": "response.output_item.done",
- "item": map[string]any{
- "id": "ig_ws_1",
- "type": "image_generation_call",
- "status": "generating",
- "result": "final-image",
- },
- }); err != nil {
- t.Errorf("write response.output_item.done failed: %v", err)
- return
- }
- if err := conn.WriteJSON(map[string]any{
- "type": "response.completed",
- "response": map[string]any{
- "id": "resp_ws_image_1",
- "model": "gpt-5.4",
- "output": []any{
- map[string]any{
- "id": "ig_ws_1",
- "type": "image_generation_call",
- "status": "in_progress",
- "result": "final-image",
- },
- },
- "usage": map[string]any{
- "input_tokens": 9,
- "output_tokens": 4,
- },
- },
- }); err != nil {
- t.Errorf("write response.completed failed: %v", err)
- return
- }
- }))
- defer wsServer.Close()
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- groupID := int64(1010)
- c.Set("api_key", &APIKey{
- GroupID: &groupID,
- Group: &Group{
- ID: groupID,
- AllowImageGeneration: true,
- },
- })
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
- cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
- cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5
- cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- }
- account := &Account{
- ID: 10,
- Name: "openai-ws-image",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "base_url": wsServer.URL,
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- 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"}}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "resp_ws_image_1", result.RequestID)
- require.Equal(t, 1, result.ImageCount)
- require.Equal(t, "1K", result.ImageSize)
- require.Equal(t, "gpt-image-2", result.BillingModel)
- require.Equal(t, 9, result.Usage.InputTokens)
- require.Equal(t, 4, result.Usage.OutputTokens)
- require.True(t, result.OpenAIWSMode)
- require.Equal(t, "resp_ws_image_1", gjson.GetBytes(rec.Body.Bytes(), "id").String())
- require.Equal(t, "completed", gjson.GetBytes(rec.Body.Bytes(), "output.0.status").String())
- }
- func requestToJSONString(payload map[string]any) string {
- if len(payload) == 0 {
- return "{}"
- }
- b, err := json.Marshal(payload)
- if err != nil {
- return "{}"
- }
- return string(b)
- }
- func TestOpenAIGatewayService_BuildOpenAIWSHeadersPreservesCodexIdentity(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
- c.Request.Header.Set("X-Codex-Window-ID", "window-ws")
- c.Request.Header.Set("X-Codex-Installation-ID", "installation-ws")
- c.Request.Header.Set("session-id", "session-ws")
- c.Request.Header.Set("thread-id", "thread-ws")
- c.Request.Header.Set("x-client-request-id", "client-request-ws")
- c.Request.Header.Set("X-Test", "blocked")
- svc := &OpenAIGatewayService{}
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- headers, _, err := svc.buildOpenAIWSHeaders(
- context.Background(),
- c,
- account,
- "token",
- OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2},
- true,
- "",
- "",
- "",
- "",
- "",
- )
- require.NoError(t, err)
- require.Equal(t, "window-ws", headers.Get("X-Codex-Window-ID"))
- require.Equal(t, "installation-ws", headers.Get("X-Codex-Installation-ID"))
- require.Equal(t, "session-ws", headers.Get("session-id"))
- require.Equal(t, "thread-ws", headers.Get("thread-id"))
- require.Equal(t, "client-request-ws", headers.Get("x-client-request-id"))
- require.Empty(t, headers.Get("X-Test"))
- }
- func TestOpenAIGatewayService_BuildOpenAIWSHeadersDeviceModePreservesNamespacedClientSessionIdentity(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
- c.Request.Header.Set("X-Codex-Installation-ID", "client-installation")
- c.Request.Header.Set("X-Codex-Window-ID", "client-window")
- c.Request.Header.Set("session-id", "client-session")
- c.Request.Header.Set("thread-id", "client-thread")
- c.Request.Header.Set("x-client-request-id", "client-request")
- account := newTestOAuthAccount(1300, map[string]any{codexFingerprintModeExtraKey: "device"})
- ids := resolveCodexFingerprintIDsFromRequest(account, c.Request.Header)
- require.NotNil(t, ids)
- stageCodexFingerprintIDs(c, ids)
- svc := &OpenAIGatewayService{}
- headers, _, err := svc.buildOpenAIWSHeaders(
- context.Background(),
- c,
- account,
- "token",
- OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2},
- true,
- "",
- "",
- "",
- "",
- "",
- )
- require.NoError(t, err)
- require.Equal(t, ids.installationID, headers.Get("x-codex-installation-id"))
- require.NotEqual(t, "client-installation", headers.Get("x-codex-installation-id"))
- require.Equal(t, scopeCodexAccountIdentityValue(account, 0, "window", "client-window"), headers.Get("x-codex-window-id"))
- require.Equal(t, scopeCodexAccountIdentityValue(account, 0, "session", "client-session"), headers.Get("session-id"))
- require.Equal(t, scopeCodexAccountIdentityValue(account, 0, "thread", "client-thread"), headers.Get("thread-id"))
- require.Equal(t, scopeCodexAccountIdentityValue(account, 0, "request", "client-request"), headers.Get("x-client-request-id"))
- }
- func TestLogOpenAIWSBindResponseAccountWarn(t *testing.T) {
- require.NotPanics(t, func() {
- logOpenAIWSBindResponseAccountWarn(1, 2, "resp_ok", nil)
- })
- require.NotPanics(t, func() {
- logOpenAIWSBindResponseAccountWarn(1, 2, "resp_err", errors.New("bind failed"))
- })
- }
- func TestOpenAIGatewayService_Forward_WSv2_RewriteModelAndToolCallsOnCompletedEvent(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
- groupID := int64(3001)
- c.Set("api_key", &APIKey{GroupID: &groupID})
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
- cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
- cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5
- cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
- captureConn := &openAIWSCaptureConn{
- events: [][]byte{
- []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\"}"}}]}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 1301,
- Name: "openai-rewrite",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "model_mapping": map[string]any{
- "custom-original-model": "gpt-5.1",
- },
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"custom-original-model","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "resp_model_tool_1", result.RequestID)
- require.Equal(t, "response.completed", result.UpstreamTerminalEvent)
- require.True(t, result.SucceededForScheduling())
- require.Equal(t, "custom-original-model", gjson.GetBytes(rec.Body.Bytes(), "model").String(), "响应模型应回写为原始请求模型")
- require.Equal(t, "edit", gjson.GetBytes(rec.Body.Bytes(), "tool_calls.0.function.name").String(), "工具名称应被修正为 OpenCode 规范")
- }
- func TestOpenAIGatewayService_Forward_WSv2_ResponseFailedIsNotSchedulingSuccess(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
- cfg := newOpenAIWSV2TestConfig()
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- captureConn := &openAIWSCaptureConn{events: [][]byte{
- []byte(`{"type":"response.failed","response":{"id":"resp_failed_1","model":"gpt-5.5","error":{"code":"server_error","message":"Internal error"}}}`),
- }}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(&openAIWSCaptureDialer{conn: captureConn})
- svc := &OpenAIGatewayService{
- cfg: cfg,
- rateLimitService: NewRateLimitService(transientCooldownAccountRepo{}, nil, cfg, nil, nil),
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 1302,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{"api_key": "sk-test"},
- Extra: map[string]any{"responses_websockets_v2_enabled": true},
- }
- svc.recordOpenAIAccountModelTransientFailure(account, "gpt-5.5", time.Now())
- result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.5","stream":false,"input":"hello"}`))
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "response.failed", result.UpstreamTerminalEvent)
- require.False(t, result.SucceededForScheduling())
- require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5"))
- }
- func TestOpenAIWSPayloadString_OnlyAcceptsStringValues(t *testing.T) {
- payload := map[string]any{
- "type": nil,
- "model": 123,
- "prompt_cache_key": " cache-key ",
- "previous_response_id": []byte(" resp_1 "),
- }
- require.Equal(t, "", openAIWSPayloadString(payload, "type"))
- require.Equal(t, "", openAIWSPayloadString(payload, "model"))
- require.Equal(t, "cache-key", openAIWSPayloadString(payload, "prompt_cache_key"))
- require.Equal(t, "resp_1", openAIWSPayloadString(payload, "previous_response_id"))
- }
- func TestOpenAIGatewayService_Forward_WSv2_PoolReuseNotOneToOne(t *testing.T) {
- gin.SetMode(gin.TestMode)
- var upgradeCount atomic.Int64
- var sequence atomic.Int64
- upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
- wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- upgradeCount.Add(1)
- conn, err := upgrader.Upgrade(w, r, nil)
- if err != nil {
- t.Errorf("upgrade websocket failed: %v", err)
- return
- }
- defer func() {
- _ = conn.Close()
- }()
- for {
- var request map[string]any
- if err := conn.ReadJSON(&request); err != nil {
- return
- }
- idx := sequence.Add(1)
- responseID := "resp_reuse_" + strconv.FormatInt(idx, 10)
- if err := conn.WriteJSON(map[string]any{
- "type": "response.created",
- "response": map[string]any{
- "id": responseID,
- "model": "gpt-5.1",
- },
- }); err != nil {
- return
- }
- if err := conn.WriteJSON(map[string]any{
- "type": "response.completed",
- "response": map[string]any{
- "id": responseID,
- "model": "gpt-5.1",
- "usage": map[string]any{
- "input_tokens": 2,
- "output_tokens": 1,
- },
- },
- }); err != nil {
- return
- }
- }
- }))
- defer wsServer.Close()
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
- cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
- cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 30
- cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 10
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- }
- account := &Account{
- ID: 19,
- Name: "openai-ws",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 2,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "base_url": wsServer.URL,
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- for i := 0; i < 2; i++ {
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
- groupID := int64(2001)
- c.Set("api_key", &APIKey{GroupID: &groupID})
- body := []byte(`{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_prev_reuse","input":[{"type":"input_text","text":"hello"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.True(t, strings.HasPrefix(result.RequestID, "resp_reuse_"))
- }
- // 条件式 MarkBroken:正常终端事件退出后连接归还复用,不再无条件销毁。
- require.Equal(t, int64(1), upgradeCount.Load(), "正常完成后连接应归还复用,不应每次新建")
- metrics := svc.SnapshotOpenAIWSPoolMetrics()
- require.GreaterOrEqual(t, metrics.AcquireReuseTotal, int64(1))
- require.GreaterOrEqual(t, metrics.ConnPickTotal, int64(1))
- }
- func TestOpenAIGatewayService_Forward_WSv2_OAuthStoreFalseByDefault(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
- c.Request.Header.Set("session_id", "sess-oauth-1")
- c.Request.Header.Set("conversation_id", "conv-oauth-1")
- c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2")
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.AllowStoreRecovery = false
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- captureConn := &openAIWSCaptureConn{
- events: [][]byte{
- []byte(`{"type":"response.completed","response":{"id":"resp_oauth_1","model":"gpt-5.1","usage":{"input_tokens":3,"output_tokens":2}}}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 29,
- Name: "openai-oauth",
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "access_token": "oauth-token-1",
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"store":true,"input":[{"type":"input_text","text":"hello","namespace":"native-wsv2"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "resp_oauth_1", result.RequestID)
- require.NotNil(t, captureConn.lastWrite)
- requestJSON := requestToJSONString(captureConn.lastWrite)
- require.True(t, gjson.Get(requestJSON, "store").Exists(), "OAuth WSv2 应显式写入 store 字段")
- require.False(t, gjson.Get(requestJSON, "store").Bool(), "默认策略应将 OAuth store 置为 false")
- require.True(t, gjson.Get(requestJSON, "stream").Exists(), "WSv2 payload 应保留 stream 字段")
- require.True(t, gjson.Get(requestJSON, "stream").Bool(), "OAuth Codex 规范化后应强制 stream=true")
- require.Equal(t, "native-wsv2", gjson.Get(requestJSON, "input.0.namespace").String(), "OAuth WSv2 应保留原生 namespace")
- require.Equal(t, openAIWSBetaV2Value, captureDialer.lastHeaders.Get("OpenAI-Beta"))
- require.Equal(t, "remote_compaction_v2", captureDialer.lastHeaders.Get("x-codex-beta-features"))
- // OAuth 账号的 session_id/conversation_id 应同时按 API key 和上游账号隔离,
- // 测试中未设置 api_key 到 context,apiKeyID=0。
- require.Equal(t, isolateOpenAIUpstreamSessionID(0, account, "sess-oauth-1"), captureDialer.lastHeaders.Get("session_id"))
- require.Equal(t, isolateOpenAIUpstreamSessionID(0, account, "conv-oauth-1"), captureDialer.lastHeaders.Get("conversation_id"))
- }
- func TestOpenAIGatewayService_Forward_WSv2_OAuthSanitizesInvalidNativeToolItemID(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
- groupID := int64(5662)
- c.Set("api_key", &APIKey{GroupID: &groupID})
- cfg := newOpenAIWSV2TestConfig()
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
- cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
- cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
- cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
- captureConn := &openAIWSCaptureConn{events: [][]byte{
- []byte(`{"type":"response.completed","response":{"id":"resp_oauth_tool_history","model":"gpt-5.6-sol","usage":{"input_tokens":1,"output_tokens":1}}}`),
- }}
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 5662,
- Name: "openai-oauth-ws-tool-history",
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{"access_token": "test-oauth-token"},
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- 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"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- captureConn.mu.Lock()
- requestPayload := cloneMapStringAny(captureConn.lastWrite)
- captureConn.mu.Unlock()
- requestJSON := requestToJSONString(requestPayload)
- require.Equal(t, "response.create", gjson.Get(requestJSON, "type").String())
- require.False(t, gjson.Get(requestJSON, "input.0.id").Exists(), "stale fc_* ID must not be replayed as a native custom_tool_call ID")
- require.Equal(t, "ctc_hotfix", gjson.Get(requestJSON, "input.0.call_id").String())
- require.Equal(t, "custom_tool_call_output", gjson.Get(requestJSON, "input.1.type").String())
- require.Equal(t, "ctc_hotfix", gjson.Get(requestJSON, "input.1.call_id").String())
- }
- func TestOpenAIGatewayService_Forward_WSv2_OAuthOriginatorCompatibility(t *testing.T) {
- gin.SetMode(gin.TestMode)
- // WS 握手头与 HTTP 出站共用身份收口:同样强制统一为网关规范身份,
- // 客户端自报的 originator / user-agent 不参与构造(issue #3901 的配对不变式自然满足)。
- tests := []struct {
- name string
- userAgent string
- originator string
- }{
- {name: "official desktop ua", userAgent: "Codex Desktop/1.2.3"},
- {
- name: "mismatched originator",
- userAgent: "codex_vscode/0.140.2 (Mac OS X 14.0; arm64) vscode (codex_vscode; 0.140.2)",
- originator: "codex_cli_rs",
- },
- {
- name: "tui identity",
- userAgent: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
- originator: "codex-tui",
- },
- {name: "official originator without ua", originator: "codex_vscode"},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- if tt.userAgent != "" {
- c.Request.Header.Set("User-Agent", tt.userAgent)
- }
- if tt.originator != "" {
- c.Request.Header.Set("originator", tt.originator)
- }
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.AllowStoreRecovery = false
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- captureConn := &openAIWSCaptureConn{
- events: [][]byte{
- []byte(`{"type":"response.completed","response":{"id":"resp_oauth_originator","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 129,
- Name: "openai-oauth",
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "access_token": "oauth-token-1",
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, openai.CodexDefaultOriginator, captureDialer.lastHeaders.Get("originator"))
- require.Equal(t, codexCLIUserAgent, captureDialer.lastHeaders.Get("user-agent"))
- require.Equal(t, codexCLIVersion, captureDialer.lastHeaders.Get("version"))
- })
- }
- }
- // 账号级自定义 UA 是管理员的显式配置,WS 握手与 HTTP 出站必须一视同仁地生效——
- // 否则同一个账号在两种传输上以不同身份出站。指纹保留、版本段重建、originator 配套。
- func TestOpenAIGatewayService_Forward_WSv2_OAuthHonorsAccountUserAgent(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "luna/1.0.0")
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.AllowStoreRecovery = false
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- captureConn := &openAIWSCaptureConn{
- events: [][]byte{
- []byte(`{"type":"response.completed","response":{"id":"resp_oauth_account_ua","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 130,
- Name: "openai-oauth-custom-ua",
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "access_token": "oauth-token-1",
- // 填写于某个历史版本的账号级 UA:指纹要保留,版本段不能被逐字沿用。
- "user_agent": "codex-tui/0.125.0 (Mac OS X 15.1.0; arm64) iTerm.app",
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "codex-tui", captureDialer.lastHeaders.Get("originator"))
- require.Equal(t,
- "codex-tui/"+codexCLIVersion+" (Mac OS X 15.1.0; arm64) iTerm.app",
- captureDialer.lastHeaders.Get("user-agent"),
- )
- require.Equal(t, codexCLIVersion, captureDialer.lastHeaders.Get("version"))
- }
- func TestOpenAIGatewayService_Forward_WSv2_HeaderSessionFallbackFromPromptCacheKey(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- captureConn := &openAIWSCaptureConn{
- events: [][]byte{
- []byte(`{"type":"response.completed","response":{"id":"resp_prompt_cache_key","model":"gpt-5.1","usage":{"input_tokens":2,"output_tokens":1}}}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 31,
- Name: "openai-oauth",
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "access_token": "oauth-token-1",
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":true,"prompt_cache_key":"pcache_123","input":[{"type":"input_text","text":"hi"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "resp_prompt_cache_key", result.RequestID)
- // OAuth 账号的 session_id 应同时按 API key 和上游账号隔离(apiKeyID=0)。
- require.Equal(t, isolateOpenAIUpstreamSessionID(0, account, "pcache_123"), captureDialer.lastHeaders.Get("session_id"))
- require.Empty(t, captureDialer.lastHeaders.Get("conversation_id"))
- require.NotNil(t, captureConn.lastWrite)
- require.True(t, gjson.Get(requestToJSONString(captureConn.lastWrite), "stream").Exists())
- }
- func TestOpenAIGatewayService_Forward_WSv2_CodexFingerprintHandshakeBodyParityAndDefaultCacheKey(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
- c.Request.Header.Set("originator", "codex_cli_rs")
- c.Request.Header.Set("session-id", "header-session")
- 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"}`)
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- captureConn := &openAIWSCaptureConn{
- events: [][]byte{
- []byte(`{"type":"response.completed","response":{"id":"resp_ws_fingerprint","model":"gpt-5.2","usage":{"input_tokens":2,"output_tokens":1}}}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := newTestOAuthAccount(4405, map[string]any{
- codexFingerprintModeExtraKey: "session",
- "responses_websockets_v2_enabled": true,
- })
- account.Name = "oauth-ws-fingerprint"
- account.Status = StatusActive
- account.Schedulable = true
- account.Concurrency = 1
- account.Credentials = map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"}
- 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"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "resp_ws_fingerprint", result.RequestID)
- require.NotNil(t, captureConn.lastWrite)
- seed, ok := codexFingerprintSeed(account.Extra)
- require.True(t, ok)
- wantInstall := resolveConvergedInstallationID(account, seed)
- wantSession := resolveConvergedSessionID(seed)
- wantThread := resolveConvergedThreadID(seed, "header-session")
- payloadJSON := requestToJSONString(captureConn.lastWrite)
- require.Equal(t, wantInstall, captureDialer.lastHeaders.Get("x-codex-installation-id"))
- require.Equal(t, wantSession, captureDialer.lastHeaders.Get("session-id"))
- require.Equal(t, wantSession, captureDialer.lastHeaders.Get("session_id"))
- require.Equal(t, wantThread, captureDialer.lastHeaders.Get("thread-id"))
- require.Equal(t, wantThread, captureDialer.lastHeaders.Get("x-client-request-id"))
- require.Equal(t, wantThread+":0", captureDialer.lastHeaders.Get("x-codex-window-id"))
- require.Equal(t, wantSession, gjson.Get(payloadJSON, "prompt_cache_key").String())
- require.Equal(t, wantInstall, gjson.Get(payloadJSON, "client_metadata.x-codex-installation-id").String())
- require.Equal(t, wantSession, gjson.Get(payloadJSON, "client_metadata.session_id").String())
- require.Equal(t, wantThread, gjson.Get(payloadJSON, "client_metadata.thread_id").String())
- require.Equal(t, wantThread+":0", gjson.Get(payloadJSON, "client_metadata.x-codex-window-id").String())
- bodyTurnMetadata := gjson.Get(payloadJSON, "client_metadata.x-codex-turn-metadata").String()
- headerTurnMetadata := captureDialer.lastHeaders.Get("x-codex-turn-metadata")
- require.Equal(t, wantInstall, gjson.Get(bodyTurnMetadata, "installation_id").String())
- require.Equal(t, wantSession, gjson.Get(bodyTurnMetadata, "session_id").String())
- require.Equal(t, wantThread, gjson.Get(bodyTurnMetadata, "thread_id").String())
- require.Equal(t, wantSession, gjson.Get(headerTurnMetadata, "session_id").String())
- require.Equal(t, gjson.Get(bodyTurnMetadata, "turn_id").String(), gjson.Get(headerTurnMetadata, "turn_id").String())
- require.NotZero(t, gjson.Get(bodyTurnMetadata, "turn_started_at_unix_ms").Int())
- require.Equal(t,
- gjson.Get(bodyTurnMetadata, "turn_started_at_unix_ms").Int(),
- gjson.Get(headerTurnMetadata, "turn_started_at_unix_ms").Int(),
- )
- }
- func TestOpenAIGatewayService_Forward_WSv2_ResponseDoneUsageParsed(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "unit-test-agent/1.0")
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- captureConn := &openAIWSCaptureConn{
- events: [][]byte{
- []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}}}}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 32,
- Name: "openai-ws-done",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hi"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "resp_done_usage", result.RequestID)
- require.Equal(t, 13, result.Usage.InputTokens)
- require.Equal(t, 8, result.Usage.OutputTokens)
- require.Equal(t, 5, result.Usage.CacheReadInputTokens)
- require.Equal(t, 2, result.Usage.CacheCreationInputTokens)
- require.Equal(t, 4, result.Usage.ImageOutputTokens)
- }
- func TestOpenAIGatewayService_Forward_WSv1_Unsupported(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsockets = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = false
- upstream := &httpUpstreamRecorder{
- resp: &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"application/json"}},
- Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":1}}`)),
- },
- }
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: upstream,
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- }
- account := &Account{
- ID: 39,
- Name: "openai-ws-v1",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "base_url": "https://api.openai.com/v1/responses",
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"previous_response_id":"resp_prev_v1","input":[{"type":"input_text","text":"hello"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.Error(t, err)
- require.Nil(t, result)
- require.Contains(t, err.Error(), "ws v1")
- require.Equal(t, http.StatusBadRequest, rec.Code)
- require.Contains(t, rec.Body.String(), "WSv1")
- require.Nil(t, upstream.lastReq, "WSv1 不支持时不应触发 HTTP 上游请求")
- }
- func TestOpenAIGatewayService_Forward_WSv2_TurnStateAndMetadataReplayOnReconnect(t *testing.T) {
- gin.SetMode(gin.TestMode)
- var connIndex atomic.Int64
- headersCh := make(chan http.Header, 4)
- upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
- wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- idx := connIndex.Add(1)
- headersCh <- cloneHeader(r.Header)
- respHeader := http.Header{}
- if idx == 1 {
- respHeader.Set("x-codex-turn-state", "turn_state_first")
- }
- conn, err := upgrader.Upgrade(w, r, respHeader)
- if err != nil {
- t.Errorf("upgrade websocket failed: %v", err)
- return
- }
- defer func() {
- _ = conn.Close()
- }()
- var request map[string]any
- if err := conn.ReadJSON(&request); err != nil {
- t.Errorf("read ws request failed: %v", err)
- return
- }
- responseID := "resp_turn_" + strconv.FormatInt(idx, 10)
- if err := conn.WriteJSON(map[string]any{
- "type": "response.completed",
- "response": map[string]any{
- "id": responseID,
- "model": "gpt-5.1",
- "usage": map[string]any{
- "input_tokens": 2,
- "output_tokens": 1,
- },
- },
- }); err != nil {
- t.Errorf("write response.completed failed: %v", err)
- return
- }
- }))
- defer wsServer.Close()
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 0
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- }
- account := &Account{
- ID: 49,
- Name: "openai-turn-state",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "base_url": wsServer.URL,
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- reqBody := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
- rec1 := httptest.NewRecorder()
- c1, _ := gin.CreateTestContext(rec1)
- c1.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c1.Request.Header.Set("session_id", "session_turn_state")
- c1.Request.Header.Set("x-codex-turn-metadata", "turn_meta_1")
- result1, err := svc.Forward(context.Background(), c1, account, reqBody)
- require.NoError(t, err)
- require.NotNil(t, result1)
- sessionHash := svc.GenerateSessionHash(c1, reqBody)
- store := svc.getOpenAIWSStateStore()
- turnState, ok := store.GetSessionTurnState(0, sessionHash)
- require.True(t, ok)
- require.Equal(t, "turn_state_first", turnState)
- // 主动淘汰连接,模拟下一次请求发生重连。
- connID, hasConn := store.GetResponseConn(result1.RequestID)
- require.True(t, hasConn)
- svc.getOpenAIWSConnPool().evictConn(account.ID, connID)
- rec2 := httptest.NewRecorder()
- c2, _ := gin.CreateTestContext(rec2)
- c2.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c2.Request.Header.Set("session_id", "session_turn_state")
- c2.Request.Header.Set("x-codex-turn-metadata", "turn_meta_2")
- result2, err := svc.Forward(context.Background(), c2, account, reqBody)
- require.NoError(t, err)
- require.NotNil(t, result2)
- firstHandshakeHeaders := <-headersCh
- secondHandshakeHeaders := <-headersCh
- require.Equal(t, "turn_meta_1", firstHandshakeHeaders.Get("X-Codex-Turn-Metadata"))
- require.Equal(t, "turn_meta_2", secondHandshakeHeaders.Get("X-Codex-Turn-Metadata"))
- require.Equal(t, "turn_state_first", secondHandshakeHeaders.Get("X-Codex-Turn-State"))
- }
- func TestOpenAIGatewayService_Forward_WSv2_GeneratePrewarm(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("session_id", "session-prewarm")
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.PrewarmGenerateEnabled = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- captureConn := &openAIWSCaptureConn{
- events: [][]byte{
- []byte(`{"type":"response.completed","response":{"id":"resp_prewarm_1","model":"gpt-5.1","usage":{"input_tokens":0,"output_tokens":0}}}`),
- []byte(`{"type":"response.completed","response":{"id":"resp_main_1","model":"gpt-5.1","usage":{"input_tokens":4,"output_tokens":2}}}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 59,
- Name: "openai-prewarm",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "resp_main_1", result.RequestID)
- require.Len(t, captureConn.writes, 2, "开启 generate=false 预热后应发送两次 WS 请求")
- firstWrite := requestToJSONString(captureConn.writes[0])
- secondWrite := requestToJSONString(captureConn.writes[1])
- require.True(t, gjson.Get(firstWrite, "generate").Exists())
- require.False(t, gjson.Get(firstWrite, "generate").Bool())
- require.False(t, gjson.Get(secondWrite, "generate").Exists())
- }
- func TestOpenAIGatewayService_PrewarmReadHonorsParentContext(t *testing.T) {
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.PrewarmGenerateEnabled = true
- cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5
- cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
- svc := &OpenAIGatewayService{
- cfg: cfg,
- toolCorrector: NewCodexToolCorrector(),
- }
- account := &Account{
- ID: 601,
- Name: "openai-prewarm-timeout",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- }
- conn := newOpenAIWSConn("prewarm_ctx_conn", account.ID, &openAIWSBlockingConn{
- readDelay: 200 * time.Millisecond,
- }, nil)
- lease := &openAIWSConnLease{
- accountID: account.ID,
- conn: conn,
- }
- payload := map[string]any{
- "type": "response.create",
- "model": "gpt-5.1",
- }
- ctx, cancel := context.WithTimeout(context.Background(), 40*time.Millisecond)
- defer cancel()
- start := time.Now()
- err := svc.performOpenAIWSGeneratePrewarm(
- ctx,
- lease,
- OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2},
- payload,
- "",
- map[string]any{"model": "gpt-5.1"},
- account,
- nil,
- 0,
- )
- elapsed := time.Since(start)
- require.Error(t, err)
- require.Contains(t, err.Error(), "prewarm_read_event")
- require.Less(t, elapsed, 180*time.Millisecond, "预热读取应受父 context 取消控制,不应阻塞到 read_timeout")
- }
- func TestOpenAIGatewayService_Forward_WSv2_TurnMetadataInPayloadOnConnReuse(t *testing.T) {
- gin.SetMode(gin.TestMode)
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- captureConn := &openAIWSCaptureConn{
- events: [][]byte{
- []byte(`{"type":"response.completed","response":{"id":"resp_meta_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
- []byte(`{"type":"response.completed","response":{"id":"resp_meta_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 69,
- Name: "openai-turn-metadata",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
- rec1 := httptest.NewRecorder()
- c1, _ := gin.CreateTestContext(rec1)
- c1.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c1.Request.Header.Set("session_id", "session-metadata-reuse")
- c1.Request.Header.Set("x-codex-turn-metadata", "turn_meta_payload_1")
- result1, err := svc.Forward(context.Background(), c1, account, body)
- require.NoError(t, err)
- require.NotNil(t, result1)
- require.Equal(t, "resp_meta_1", result1.RequestID)
- require.Len(t, captureConn.writes, 1)
- firstWrite := requestToJSONString(captureConn.writes[0])
- require.Equal(t, "turn_meta_payload_1", gjson.Get(firstWrite, "client_metadata.x-codex-turn-metadata").String())
- rec2 := httptest.NewRecorder()
- c2, _ := gin.CreateTestContext(rec2)
- c2.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c2.Request.Header.Set("session_id", "session-metadata-reuse")
- c2.Request.Header.Set("x-codex-turn-metadata", "turn_meta_payload_2")
- result2, err := svc.Forward(context.Background(), c2, account, body)
- require.NoError(t, err)
- require.NotNil(t, result2)
- require.Equal(t, "resp_meta_2", result2.RequestID)
- require.Equal(t, 1, captureDialer.DialCount(), "同一账号两轮请求应复用同一 WS 连接")
- require.Len(t, captureConn.writes, 2)
- firstWrite = requestToJSONString(captureConn.writes[0])
- secondWrite := requestToJSONString(captureConn.writes[1])
- require.Equal(t, "turn_meta_payload_1", gjson.Get(firstWrite, "client_metadata.x-codex-turn-metadata").String())
- require.Equal(t, "turn_meta_payload_2", gjson.Get(secondWrite, "client_metadata.x-codex-turn-metadata").String())
- }
- func TestOpenAIGatewayService_Forward_WSv2StoreFalseSessionConnIsolation(t *testing.T) {
- gin.SetMode(gin.TestMode)
- var upgradeCount atomic.Int64
- var sequence atomic.Int64
- upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
- wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- upgradeCount.Add(1)
- conn, err := upgrader.Upgrade(w, r, nil)
- if err != nil {
- t.Errorf("upgrade websocket failed: %v", err)
- return
- }
- defer func() {
- _ = conn.Close()
- }()
- for {
- var request map[string]any
- if err := conn.ReadJSON(&request); err != nil {
- return
- }
- responseID := "resp_store_false_" + strconv.FormatInt(sequence.Add(1), 10)
- if err := conn.WriteJSON(map[string]any{
- "type": "response.completed",
- "response": map[string]any{
- "id": responseID,
- "model": "gpt-5.1",
- "usage": map[string]any{
- "input_tokens": 1,
- "output_tokens": 1,
- },
- },
- }); err != nil {
- return
- }
- }
- }))
- defer wsServer.Close()
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 4
- cfg.Gateway.OpenAIWS.StoreDisabledForceNewConn = true
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- }
- account := &Account{
- ID: 79,
- Name: "openai-store-false",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 2,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "base_url": wsServer.URL,
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
- rec1 := httptest.NewRecorder()
- c1, _ := gin.CreateTestContext(rec1)
- c1.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c1.Request.Header.Set("session_id", "session_store_false_a")
- result1, err := svc.Forward(context.Background(), c1, account, body)
- require.NoError(t, err)
- require.NotNil(t, result1)
- require.Equal(t, int64(1), upgradeCount.Load())
- rec2 := httptest.NewRecorder()
- c2, _ := gin.CreateTestContext(rec2)
- c2.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c2.Request.Header.Set("session_id", "session_store_false_a")
- result2, err := svc.Forward(context.Background(), c2, account, body)
- require.NoError(t, err)
- require.NotNil(t, result2)
- require.Equal(t, int64(1), upgradeCount.Load(), "同一 session(store=false) 应复用同一 WS 连接")
- rec3 := httptest.NewRecorder()
- c3, _ := gin.CreateTestContext(rec3)
- c3.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c3.Request.Header.Set("session_id", "session_store_false_b")
- result3, err := svc.Forward(context.Background(), c3, account, body)
- require.NoError(t, err)
- require.NotNil(t, result3)
- require.Equal(t, int64(2), upgradeCount.Load(), "不同 session(store=false) 应隔离连接,避免续链状态互相覆盖")
- }
- func TestOpenAIGatewayService_Forward_WSv2StoreFalseDisableForceNewConnAllowsReuse(t *testing.T) {
- gin.SetMode(gin.TestMode)
- var upgradeCount atomic.Int64
- var sequence atomic.Int64
- upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
- wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- upgradeCount.Add(1)
- conn, err := upgrader.Upgrade(w, r, nil)
- if err != nil {
- t.Errorf("upgrade websocket failed: %v", err)
- return
- }
- defer func() {
- _ = conn.Close()
- }()
- for {
- var request map[string]any
- if err := conn.ReadJSON(&request); err != nil {
- return
- }
- responseID := "resp_store_false_reuse_" + strconv.FormatInt(sequence.Add(1), 10)
- if err := conn.WriteJSON(map[string]any{
- "type": "response.completed",
- "response": map[string]any{
- "id": responseID,
- "model": "gpt-5.1",
- "usage": map[string]any{
- "input_tokens": 1,
- "output_tokens": 1,
- },
- },
- }); err != nil {
- return
- }
- }
- }))
- defer wsServer.Close()
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- cfg.Gateway.OpenAIWS.StoreDisabledForceNewConn = false
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: &httpUpstreamRecorder{},
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- }
- account := &Account{
- ID: 80,
- Name: "openai-store-false-reuse",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 2,
- Credentials: map[string]any{
- "api_key": "sk-test",
- "base_url": wsServer.URL,
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"store":false,"input":[{"type":"input_text","text":"hello"}]}`)
- rec1 := httptest.NewRecorder()
- c1, _ := gin.CreateTestContext(rec1)
- c1.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c1.Request.Header.Set("session_id", "session_store_false_reuse_a")
- result1, err := svc.Forward(context.Background(), c1, account, body)
- require.NoError(t, err)
- require.NotNil(t, result1)
- require.Equal(t, int64(1), upgradeCount.Load())
- rec2 := httptest.NewRecorder()
- c2, _ := gin.CreateTestContext(rec2)
- c2.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c2.Request.Header.Set("session_id", "session_store_false_reuse_b")
- result2, err := svc.Forward(context.Background(), c2, account, body)
- require.NoError(t, err)
- require.NotNil(t, result2)
- require.Equal(t, int64(1), upgradeCount.Load(), "关闭强制新连后,不同 session(store=false) 可复用连接")
- }
- func TestOpenAIGatewayService_Forward_WSv2ReadTimeoutAppliesPerRead(t *testing.T) {
- gin.SetMode(gin.TestMode)
- rec := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(rec)
- c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
- c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0")
- cfg := &config.Config{}
- cfg.Security.URLAllowlist.Enabled = false
- cfg.Security.URLAllowlist.AllowInsecureHTTP = true
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
- cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
- cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
- cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
- cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
- cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 1
- cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
- captureConn := &openAIWSCaptureConn{
- readDelays: []time.Duration{
- 700 * time.Millisecond,
- 700 * time.Millisecond,
- },
- events: [][]byte{
- []byte(`{"type":"response.created","response":{"id":"resp_timeout_ok","model":"gpt-5.1"}}`),
- []byte(`{"type":"response.completed","response":{"id":"resp_timeout_ok","model":"gpt-5.1","usage":{"input_tokens":2,"output_tokens":1}}}`),
- },
- }
- captureDialer := &openAIWSCaptureDialer{conn: captureConn}
- pool := newOpenAIWSConnPool(cfg)
- pool.setClientDialerForTest(captureDialer)
- upstream := &httpUpstreamRecorder{
- resp: &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"application/json"}},
- Body: io.NopCloser(strings.NewReader(`{"id":"resp_http_fallback","usage":{"input_tokens":1,"output_tokens":1}}`)),
- },
- }
- svc := &OpenAIGatewayService{
- cfg: cfg,
- httpUpstream: upstream,
- cache: &stubGatewayCache{},
- openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
- toolCorrector: NewCodexToolCorrector(),
- openaiWSPool: pool,
- }
- account := &Account{
- ID: 81,
- Name: "openai-read-timeout",
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "api_key": "sk-test",
- },
- Extra: map[string]any{
- "responses_websockets_v2_enabled": true,
- },
- }
- body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`)
- result, err := svc.Forward(context.Background(), c, account, body)
- require.NoError(t, err)
- require.NotNil(t, result)
- require.Equal(t, "resp_timeout_ok", result.RequestID)
- require.Nil(t, upstream.lastReq, "每次 Read 都应独立应用超时;总时长超过 read_timeout 不应误回退 HTTP")
- }
- type openAIWSCaptureDialer struct {
- mu sync.Mutex
- conn *openAIWSCaptureConn
- lastHeaders http.Header
- handshake http.Header
- dialCount int
- }
- func (d *openAIWSCaptureDialer) Dial(
- ctx context.Context,
- wsURL string,
- headers http.Header,
- proxyURL string,
- ) (openAIWSClientConn, int, http.Header, error) {
- _ = ctx
- _ = wsURL
- _ = proxyURL
- d.mu.Lock()
- d.lastHeaders = cloneHeader(headers)
- d.dialCount++
- respHeaders := cloneHeader(d.handshake)
- d.mu.Unlock()
- return d.conn, 0, respHeaders, nil
- }
- func (d *openAIWSCaptureDialer) DialCount() int {
- d.mu.Lock()
- defer d.mu.Unlock()
- return d.dialCount
- }
- type openAIWSCaptureConn struct {
- mu sync.Mutex
- readDelays []time.Duration
- events [][]byte
- lastWrite map[string]any
- writes []map[string]any
- closed bool
- }
- func (c *openAIWSCaptureConn) WriteJSON(ctx context.Context, value any) error {
- _ = ctx
- c.mu.Lock()
- defer c.mu.Unlock()
- if c.closed {
- return errOpenAIWSConnClosed
- }
- switch payload := value.(type) {
- case map[string]any:
- c.lastWrite = cloneMapStringAny(payload)
- c.writes = append(c.writes, cloneMapStringAny(payload))
- case json.RawMessage:
- var parsed map[string]any
- if err := decodeOpenAIJSONUseNumber(payload, &parsed); err == nil {
- c.lastWrite = cloneMapStringAny(parsed)
- c.writes = append(c.writes, cloneMapStringAny(parsed))
- }
- case []byte:
- var parsed map[string]any
- if err := decodeOpenAIJSONUseNumber(payload, &parsed); err == nil {
- c.lastWrite = cloneMapStringAny(parsed)
- c.writes = append(c.writes, cloneMapStringAny(parsed))
- }
- }
- return nil
- }
- func (c *openAIWSCaptureConn) ReadMessage(ctx context.Context) ([]byte, error) {
- if ctx == nil {
- ctx = context.Background()
- }
- c.mu.Lock()
- if c.closed {
- c.mu.Unlock()
- return nil, errOpenAIWSConnClosed
- }
- if len(c.events) == 0 {
- c.mu.Unlock()
- return nil, io.EOF
- }
- delay := time.Duration(0)
- if len(c.readDelays) > 0 {
- delay = c.readDelays[0]
- c.readDelays = c.readDelays[1:]
- }
- event := c.events[0]
- c.events = c.events[1:]
- c.mu.Unlock()
- if delay > 0 {
- timer := time.NewTimer(delay)
- defer timer.Stop()
- select {
- case <-ctx.Done():
- return nil, ctx.Err()
- case <-timer.C:
- }
- }
- return event, nil
- }
- func (c *openAIWSCaptureConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) {
- payload, err := c.ReadMessage(ctx)
- if err != nil {
- return coderws.MessageText, nil, err
- }
- return coderws.MessageText, payload, nil
- }
- func (c *openAIWSCaptureConn) WriteFrame(ctx context.Context, _ coderws.MessageType, payload []byte) error {
- return c.WriteJSON(ctx, json.RawMessage(payload))
- }
- func (c *openAIWSCaptureConn) Ping(ctx context.Context) error {
- _ = ctx
- return nil
- }
- func (c *openAIWSCaptureConn) Close() error {
- c.mu.Lock()
- defer c.mu.Unlock()
- c.closed = true
- return nil
- }
- func cloneMapStringAny(src map[string]any) map[string]any {
- if src == nil {
- return nil
- }
- dst := make(map[string]any, len(src))
- for k, v := range src {
- dst[k] = v
- }
- return dst
- }
|