openai_ws_http_bridge_test.go 92 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180
  1. package service
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "io"
  7. "net/http"
  8. "net/http/httptest"
  9. "strings"
  10. "testing"
  11. "time"
  12. "github.com/Wei-Shaw/sub2api/internal/config"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  14. coderws "github.com/coder/websocket"
  15. "github.com/gin-gonic/gin"
  16. "github.com/stretchr/testify/require"
  17. "github.com/tidwall/gjson"
  18. )
  19. func TestResolveOpenAIWSClientFirstMessageTimeout(t *testing.T) {
  20. defaultTimeout := time.Duration(config.DefaultOpenAIWSClientFirstMessageTimeoutSeconds) * time.Second
  21. require.Equal(t, defaultTimeout, ResolveOpenAIWSClientFirstMessageTimeout(nil))
  22. cfg := &config.Config{}
  23. require.Equal(t, defaultTimeout, ResolveOpenAIWSClientFirstMessageTimeout(cfg))
  24. cfg.Gateway.OpenAIWS.ClientFirstMessageTimeoutSeconds = 120
  25. require.Equal(t, 120*time.Second, ResolveOpenAIWSClientFirstMessageTimeout(cfg))
  26. }
  27. func TestPrepareOpenAIWSHTTPBridgeBodyStripsWSFields(t *testing.T) {
  28. body, err := prepareOpenAIWSHTTPBridgeBody(nil, []byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":false,"previous_response_id":"resp_prev","input":"hi","sequence":900719925474099312345}`))
  29. require.NoError(t, err)
  30. require.False(t, gjson.GetBytes(body, "type").Exists())
  31. require.False(t, gjson.GetBytes(body, "generate").Exists())
  32. require.False(t, gjson.GetBytes(body, "previous_response_id").Exists())
  33. require.Equal(t, "gpt-5", gjson.GetBytes(body, "model").String())
  34. require.True(t, gjson.GetBytes(body, "stream").Bool())
  35. require.Equal(t, "hi", gjson.GetBytes(body, "input").String())
  36. require.Equal(t, "900719925474099312345", gjson.GetBytes(body, "sequence").Raw)
  37. _, err = prepareOpenAIWSHTTPBridgeBody(nil, []byte(`{"type":"response.create"}{"trailing":true}`))
  38. require.Error(t, err)
  39. }
  40. func TestPrepareOpenAIWSHTTPBridgeBodyStripsNoneReasoningForCompatibleEndpoint(t *testing.T) {
  41. payload := []byte(`{"type":"response.create","model":"company-coding-model","reasoning":{"effort":"none"},"input":"hi"}`)
  42. compatible := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{
  43. "base_url": "https://compat.example/v1",
  44. }}
  45. body, err := prepareOpenAIWSHTTPBridgeBody(compatible, payload)
  46. require.NoError(t, err)
  47. require.False(t, gjson.GetBytes(body, "reasoning.effort").Exists())
  48. require.False(t, gjson.GetBytes(body, "reasoning").Exists())
  49. officialBody, err := prepareOpenAIWSHTTPBridgeBody(&Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, payload)
  50. require.NoError(t, err)
  51. require.Equal(t, "none", gjson.GetBytes(officialBody, "reasoning.effort").String())
  52. }
  53. func TestProxyOpenAIWSHTTPBridgeTurn_KeepsOutboundAndObservedServiceTiersSeparate(t *testing.T) {
  54. gin.SetMode(gin.TestMode)
  55. // proxyOpenAIWSHTTPBridgeTurn 是 client WS→HTTP bridge,本身不 canonicalize
  56. // fast→priority;生产入口的归一化在 openai_ws_forwarder_ingress.go 的 fast
  57. // policy. This test covers the local observer while preserving the canonical
  58. // outbound priority tier independently.
  59. sse := strings.Join([]string{
  60. `data: {"type":"response.completed","response":{"id":"resp_tier","model":"gpt-5.5","status":"completed","service_tier":"default","usage":{"input_tokens":1,"output_tokens":1}}}`,
  61. ``,
  62. }, "\n")
  63. upstream := &httpUpstreamRecorder{resp: &http.Response{
  64. StatusCode: http.StatusOK,
  65. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  66. Body: io.NopCloser(strings.NewReader(sse)),
  67. }}
  68. svc := &OpenAIGatewayService{
  69. cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
  70. httpUpstream: upstream,
  71. }
  72. account := &Account{ID: 5881, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  73. payload := []byte(`{"type":"response.create","model":"gpt-5.5","stream":true,"service_tier":"priority","input":"hi"}`)
  74. recorder := httptest.NewRecorder()
  75. c, _ := gin.CreateTestContext(recorder)
  76. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  77. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  78. context.Background(), c, account, "test-token", payload, len(payload),
  79. "gpt-5.5", "", "", "", "", 1,
  80. func([]byte) error { return nil },
  81. )
  82. require.NoError(t, err)
  83. require.NotNil(t, result)
  84. require.Equal(t, "priority", gjson.GetBytes(upstream.lastBody, "service_tier").String())
  85. require.NotNil(t, result.ServiceTier)
  86. require.Equal(t, "priority", *result.ServiceTier)
  87. require.Equal(t, "default", result.UpstreamResponseServiceTier)
  88. }
  89. func TestProxyOpenAIWSHTTPBridgeTurn_NormalizesFastWithoutLosingObservedDefault(t *testing.T) {
  90. gin.SetMode(gin.TestMode)
  91. // The client alias fast is canonicalized to priority while the local observer
  92. // independently captures the upstream default declaration.
  93. sse := strings.Join([]string{
  94. `data: {"type":"response.completed","response":{"id":"resp_tier2","model":"gpt-5.5","status":"completed","service_tier":"default","usage":{"input_tokens":1,"output_tokens":1}}}`,
  95. ``,
  96. }, "\n")
  97. upstream := &httpUpstreamRecorder{resp: &http.Response{
  98. StatusCode: http.StatusOK,
  99. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  100. Body: io.NopCloser(strings.NewReader(sse)),
  101. }}
  102. svc := &OpenAIGatewayService{
  103. cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
  104. httpUpstream: upstream,
  105. }
  106. account := &Account{ID: 5882, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  107. payload := []byte(`{"type":"response.create","model":"gpt-5.5","stream":true,"service_tier":"fast","input":"hi"}`)
  108. recorder := httptest.NewRecorder()
  109. c, _ := gin.CreateTestContext(recorder)
  110. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  111. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  112. context.Background(), c, account, "test-token", payload, len(payload),
  113. "gpt-5.5", "", "", "", "", 1,
  114. func([]byte) error { return nil },
  115. )
  116. require.NoError(t, err)
  117. require.NotNil(t, result)
  118. require.NotNil(t, result.ServiceTier)
  119. require.Equal(t, "priority", *result.ServiceTier)
  120. require.Equal(t, "default", result.UpstreamResponseServiceTier)
  121. }
  122. func TestProxyOpenAIWSHTTPBridgeTurnAPIKeyAdaptsClientTools(t *testing.T) {
  123. gin.SetMode(gin.TestMode)
  124. sse := strings.Join([]string{
  125. `data: {"type":"response.output_item.added","sequence_number":0,"output_index":0,"item":{"type":"function_call","id":"item_exec","call_id":"call_exec","name":"exec","status":"in_progress"}}`,
  126. ``,
  127. `data: {"type":"response.function_call_arguments.done","sequence_number":1,"output_index":0,"item_id":"item_exec","call_id":"call_exec","name":"exec","arguments":"{\"input\":\"pwd\"}"}`,
  128. ``,
  129. `data: {"type":"response.output_item.done","sequence_number":2,"output_index":0,"item":{"type":"function_call","id":"item_exec","call_id":"call_exec","name":"exec","arguments":"{\"input\":\"pwd\"}","status":"completed"}}`,
  130. ``,
  131. `data: {"type":"response.completed","sequence_number":3,"response":{"id":"resp_tools","status":"completed","output":[{"type":"function_call","id":"item_exec","call_id":"call_exec","name":"exec","arguments":"{\"input\":\"pwd\"}","status":"completed"}],"usage":{"input_tokens":1,"output_tokens":1}}}`,
  132. ``,
  133. }, "\n")
  134. upstream := &httpUpstreamRecorder{resp: &http.Response{
  135. StatusCode: http.StatusOK,
  136. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  137. Body: io.NopCloser(strings.NewReader(sse)),
  138. }}
  139. svc := &OpenAIGatewayService{
  140. cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
  141. httpUpstream: upstream,
  142. }
  143. account := &Account{ID: 5659, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  144. payload := []byte(`{
  145. "type":"response.create","model":"gpt-5","stream":true,
  146. "tools":[{"type":"custom","name":"exec","description":"Run a command"}],
  147. "input":[
  148. {"type":"custom_tool_call","id":"previous_item","call_id":"previous_call","name":"exec","input":"echo ready"},
  149. {"type":"custom_tool_call_output","call_id":"previous_call","output":"ready"}
  150. ]
  151. }`)
  152. recorder := httptest.NewRecorder()
  153. c, _ := gin.CreateTestContext(recorder)
  154. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  155. var events [][]byte
  156. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  157. context.Background(), c, account, "test-token", payload, len(payload),
  158. "gpt-5", "", "", "", "", 2,
  159. func(message []byte) error {
  160. events = append(events, append([]byte(nil), message...))
  161. return nil
  162. },
  163. )
  164. require.NoError(t, err)
  165. require.NotNil(t, result)
  166. require.Equal(t, "function", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
  167. require.Equal(t, "function_call", gjson.GetBytes(upstream.lastBody, "input.0.type").String())
  168. require.JSONEq(t, `{"input":"echo ready"}`, gjson.GetBytes(upstream.lastBody, "input.0.arguments").String())
  169. require.False(t, gjson.GetBytes(upstream.lastBody, "input.0.input").Exists())
  170. require.Equal(t, "function_call_output", gjson.GetBytes(upstream.lastBody, "input.1.type").String())
  171. var outputDone, completed []byte
  172. for _, event := range events {
  173. switch gjson.GetBytes(event, "type").String() {
  174. case "response.output_item.done":
  175. outputDone = event
  176. case "response.completed":
  177. completed = event
  178. }
  179. }
  180. require.NotEmpty(t, outputDone)
  181. require.Equal(t, "custom_tool_call", gjson.GetBytes(outputDone, "item.type").String())
  182. require.Equal(t, "pwd", gjson.GetBytes(outputDone, "item.input").String())
  183. require.False(t, gjson.GetBytes(outputDone, "item.arguments").Exists())
  184. require.NotEmpty(t, completed)
  185. require.Equal(t, "custom_tool_call", gjson.GetBytes(completed, "response.output.0.type").String())
  186. require.Equal(t, "pwd", gjson.GetBytes(completed, "response.output.0.input").String())
  187. require.True(t, result.wsReplayInputExists)
  188. require.Len(t, result.wsReplayInput, 1)
  189. require.Equal(t, "custom_tool_call", gjson.GetBytes(result.wsReplayInput[0], "type").String())
  190. require.Equal(t, "pwd", gjson.GetBytes(result.wsReplayInput[0], "input").String())
  191. }
  192. func TestProxyOpenAIWSHTTPBridgeTurnAPIKeyRestoresClientToolsInResponseDone(t *testing.T) {
  193. gin.SetMode(gin.TestMode)
  194. sse := strings.Join([]string{
  195. `data: {"type":"response.output_item.added","sequence_number":0,"output_index":0,"item":{"type":"function_call","id":"item_exec","call_id":"call_exec","name":"exec","status":"in_progress"}}`,
  196. ``,
  197. `data: {"type":"response.function_call_arguments.done","sequence_number":1,"output_index":0,"item_id":"item_exec","call_id":"call_exec","name":"exec","arguments":"{\"input\":\"pwd\"}"}`,
  198. ``,
  199. `data: {"type":"response.done","sequence_number":2,"response":{"id":"resp_tools","status":"completed","output":[{"type":"function_call","id":"item_exec","call_id":"call_exec","name":"exec","arguments":"{\"input\":\"pwd\"}","status":"completed"}],"usage":{"input_tokens":1,"output_tokens":1}}}`,
  200. ``,
  201. }, "\n")
  202. upstream := &httpUpstreamRecorder{resp: &http.Response{
  203. StatusCode: http.StatusOK,
  204. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  205. Body: io.NopCloser(strings.NewReader(sse)),
  206. }}
  207. svc := &OpenAIGatewayService{
  208. cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
  209. httpUpstream: upstream,
  210. }
  211. account := &Account{ID: 5764, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  212. payload := []byte(`{
  213. "type":"response.create","model":"gpt-5","stream":true,
  214. "tools":[{"type":"custom","name":"exec","description":"Run a command"}],
  215. "input":"run pwd"
  216. }`)
  217. recorder := httptest.NewRecorder()
  218. c, _ := gin.CreateTestContext(recorder)
  219. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  220. var events [][]byte
  221. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  222. context.Background(), c, account, "test-token", payload, len(payload),
  223. "gpt-5", "", "", "", "", 1,
  224. func(message []byte) error {
  225. events = append(events, append([]byte(nil), message...))
  226. return nil
  227. },
  228. )
  229. require.NoError(t, err)
  230. require.NotNil(t, result)
  231. require.Len(t, events, 4)
  232. terminal := events[len(events)-1]
  233. require.Equal(t, "response.done", gjson.GetBytes(terminal, "type").String())
  234. require.Equal(t, int64(3), gjson.GetBytes(terminal, "sequence_number").Int())
  235. require.Equal(t, "custom_tool_call", gjson.GetBytes(terminal, "response.output.0.type").String())
  236. require.Equal(t, "pwd", gjson.GetBytes(terminal, "response.output.0.input").String())
  237. require.False(t, gjson.GetBytes(terminal, "response.output.0.arguments").Exists())
  238. require.True(t, result.wsReplayInputExists)
  239. require.Len(t, result.wsReplayInput, 1)
  240. require.Equal(t, "custom_tool_call", gjson.GetBytes(result.wsReplayInput[0], "type").String())
  241. }
  242. func TestProxyOpenAIWSHTTPBridgeTurnGrokPromotesDiscoveryAndRestoresNamespaceSSE(t *testing.T) {
  243. gin.SetMode(gin.TestMode)
  244. sse := strings.Join([]string{
  245. `data: {"type":"response.output_item.added","sequence_number":0,"output_index":0,"item":{"type":"function_call","id":"item_spawn","call_id":"call_spawn","name":"multi_agent_v1__spawn_agent","status":"in_progress"}}`,
  246. "",
  247. `data: {"type":"response.function_call_arguments.done","sequence_number":1,"output_index":0,"item_id":"item_spawn","call_id":"call_spawn","name":"multi_agent_v1__spawn_agent","arguments":"{\"message\":\"work\"}"}`,
  248. "",
  249. `data: {"type":"response.output_item.done","sequence_number":2,"output_index":0,"item":{"type":"function_call","id":"item_spawn","call_id":"call_spawn","name":"multi_agent_v1__spawn_agent","arguments":"{\"message\":\"work\"}","status":"completed"}}`,
  250. "",
  251. `data: {"type":"response.completed","sequence_number":3,"response":{"id":"resp_spawn","status":"completed","output":[{"type":"function_call","id":"item_spawn","call_id":"call_spawn","name":"multi_agent_v1__spawn_agent","arguments":"{\"message\":\"work\"}","status":"completed"}],"usage":{"input_tokens":1,"output_tokens":1}}}`,
  252. "",
  253. }, "\n")
  254. upstream := &httpUpstreamRecorder{resp: &http.Response{
  255. StatusCode: http.StatusOK,
  256. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  257. Body: io.NopCloser(strings.NewReader(sse)),
  258. }}
  259. svc := &OpenAIGatewayService{
  260. cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
  261. httpUpstream: upstream,
  262. }
  263. account := &Account{
  264. ID: 5765, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1,
  265. Credentials: map[string]any{"base_url": xai.DefaultCLIBaseURL},
  266. }
  267. payload := []byte(`{
  268. "type":"response.create","model":"grok-4.5","stream":true,
  269. "tools":[{"type":"tool_search"}],
  270. "input":[
  271. {"type":"tool_search_call","call_id":"call_search","arguments":{"query":"subagent"},"status":"completed"},
  272. {"type":"tool_search_output","call_id":"call_search","execution":"client","status":"completed","tools":[
  273. {"type":"namespace","name":"multi_agent_v1","tools":[{"type":"function","name":"spawn_agent","parameters":{"type":"object"}}]}
  274. ]}
  275. ]
  276. }`)
  277. recorder := httptest.NewRecorder()
  278. c, _ := gin.CreateTestContext(recorder)
  279. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  280. var events [][]byte
  281. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  282. context.Background(), c, account, "access-token", payload, len(payload),
  283. "grok-4.5", "", "", "", "", 1,
  284. func(message []byte) error {
  285. events = append(events, append([]byte(nil), message...))
  286. return nil
  287. },
  288. )
  289. require.NoError(t, err)
  290. require.NotNil(t, result)
  291. require.Equal(t, "multi_agent_v1__spawn_agent", gjson.GetBytes(upstream.lastBody, "tools.1.name").String())
  292. require.Equal(t, "function_call_output", gjson.GetBytes(upstream.lastBody, "input.1.type").String())
  293. require.False(t, gjson.GetBytes(upstream.lastBody, "input.1.tools").Exists())
  294. require.False(t, gjson.GetBytes(upstream.lastBody, "input.1.status").Exists())
  295. require.False(t, gjson.GetBytes(upstream.lastBody, "input.1.execution").Exists())
  296. state, ok := openAIWSHTTPBridgeToolStateFromContext(c)
  297. require.True(t, ok)
  298. require.Equal(t, "multi_agent_v1", state.ClientMapping.NamespaceTools["multi_agent_v1__spawn_agent"].Namespace)
  299. require.Len(t, events, 4)
  300. require.Equal(t, "spawn_agent", gjson.GetBytes(events[0], "item.name").String())
  301. require.Equal(t, "multi_agent_v1", gjson.GetBytes(events[0], "item.namespace").String())
  302. require.Equal(t, "spawn_agent", gjson.GetBytes(events[2], "item.name").String())
  303. require.Equal(t, "multi_agent_v1", gjson.GetBytes(events[2], "item.namespace").String())
  304. require.Equal(t, "spawn_agent", gjson.GetBytes(events[3], "response.output.0.name").String())
  305. require.Equal(t, "multi_agent_v1", gjson.GetBytes(events[3], "response.output.0.namespace").String())
  306. }
  307. func TestProxyOpenAIWSHTTPBridgeTurnGrokInheritsToolSearchAndPromotesFollowupDiscovery(t *testing.T) {
  308. gin.SetMode(gin.TestMode)
  309. firstSSE := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_first\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n"
  310. secondSSE := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_second\",\"output\":[{\"type\":\"function_call\",\"id\":\"item_spawn\",\"call_id\":\"call_spawn\",\"name\":\"multi_agent_v1__spawn_agent\",\"arguments\":\"{}\"}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n"
  311. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  312. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(firstSSE))},
  313. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(secondSSE))},
  314. }}
  315. svc := &OpenAIGatewayService{
  316. cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
  317. httpUpstream: upstream,
  318. }
  319. account := &Account{
  320. ID: 5766, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1,
  321. Credentials: map[string]any{"base_url": xai.DefaultCLIBaseURL, "subscription_tier": "free"},
  322. }
  323. recorder := httptest.NewRecorder()
  324. c, _ := gin.CreateTestContext(recorder)
  325. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  326. first := []byte(`{"type":"response.create","model":"grok-4.5","stream":true,"tools":[{"type":"tool_search"}],"input":"discover tools"}`)
  327. _, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  328. context.Background(), c, account, "access-token", first, len(first),
  329. "grok-4.5", "", "", "", "grok-ws-cache", 1, func([]byte) error { return nil },
  330. )
  331. require.NoError(t, err)
  332. second := []byte(`{
  333. "type":"response.create","model":"grok-4.5","stream":true,
  334. "input":[{"type":"tool_search_output","call_id":"call_search","status":"completed","tools":[
  335. {"type":"namespace","name":"multi_agent_v1","tools":[{"type":"function","name":"spawn_agent","parameters":{"type":"object"}}]}
  336. ]}]
  337. }`)
  338. var events [][]byte
  339. _, err = svc.proxyOpenAIWSHTTPBridgeTurn(
  340. context.Background(), c, account, "access-token", second, len(second),
  341. "grok-4.5", "", "", "", "grok-ws-cache", 2,
  342. func(message []byte) error {
  343. events = append(events, append([]byte(nil), message...))
  344. return nil
  345. },
  346. )
  347. require.NoError(t, err)
  348. require.Len(t, upstream.bodies, 2)
  349. require.Equal(t, "tool_search", gjson.GetBytes(upstream.bodies[1], "tools.0.name").String())
  350. require.Equal(t, "multi_agent_v1__spawn_agent", gjson.GetBytes(upstream.bodies[1], "tools.1.name").String())
  351. require.NotEqual(t, grokFreeCacheDisabledToolChoice, gjson.GetBytes(upstream.bodies[1], "tool_choice").String())
  352. require.Equal(t, "grok-ws-cache", gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String())
  353. require.Equal(t, "function_call_output", gjson.GetBytes(upstream.bodies[1], "input.0.type").String())
  354. require.Len(t, events, 1)
  355. require.Equal(t, "spawn_agent", gjson.GetBytes(events[0], "response.output.0.name").String())
  356. require.Equal(t, "multi_agent_v1", gjson.GetBytes(events[0], "response.output.0.namespace").String())
  357. }
  358. func TestOpenAIWSHTTPBridgeAPIKeyReusesClientToolMappingWhenFollowupOmitsTools(t *testing.T) {
  359. gin.SetMode(gin.TestMode)
  360. firstSSEBody := strings.Join([]string{
  361. `data: {"type":"response.completed","response":{"id":"resp_custom_first","model":"gpt-5.6-sol","output":[{"type":"function_call","id":"fc_custom_1","call_id":"call_custom_1","name":"exec","arguments":"{\"input\":\"pwd\"}"}],"usage":{"input_tokens":9,"output_tokens":1}}}`,
  362. "",
  363. }, "\n")
  364. secondSSEBody := strings.Join([]string{
  365. `data: {"type":"response.completed","response":{"id":"resp_custom_second","model":"gpt-5.6-sol","output":[],"usage":{"input_tokens":1,"output_tokens":1}}}`,
  366. "",
  367. }, "\n")
  368. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  369. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(firstSSEBody))},
  370. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(secondSSEBody))},
  371. }}
  372. cfg := &config.Config{}
  373. cfg.Security.URLAllowlist.Enabled = false
  374. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  375. cfg.Gateway.OpenAIWS.Enabled = true
  376. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  377. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  378. cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true
  379. cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = 1
  380. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  381. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  382. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  383. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  384. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  385. svc := &OpenAIGatewayService{
  386. cfg: cfg, httpUpstream: upstream, cache: &stubGatewayCache{},
  387. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), toolCorrector: NewCodexToolCorrector(),
  388. }
  389. account := &Account{
  390. ID: 9001, Name: "api-key-custom-followup", Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
  391. Credentials: map[string]any{"api_key": "sk-upstream"}, Extra: map[string]any{"responses_websockets_v2_enabled": true},
  392. Concurrency: 1, Status: StatusActive, Schedulable: true,
  393. }
  394. errCh := make(chan error, 1)
  395. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  396. conn, err := coderws.Accept(w, r, nil)
  397. if err != nil {
  398. errCh <- err
  399. return
  400. }
  401. defer func() { _ = conn.CloseNow() }()
  402. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  403. _, firstMessage, err := conn.Read(readCtx)
  404. cancelRead()
  405. if err != nil {
  406. errCh <- err
  407. return
  408. }
  409. rec := httptest.NewRecorder()
  410. ginCtx, _ := gin.CreateTestContext(rec)
  411. ginCtx.Request = r.Clone(r.Context())
  412. errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  413. }))
  414. defer wsServer.Close()
  415. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  416. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  417. cancelDial()
  418. require.NoError(t, err)
  419. defer func() { _ = clientConn.CloseNow() }()
  420. writeMessage := func(payload string) {
  421. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  422. defer cancelWrite()
  423. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  424. }
  425. readMessage := func() []byte {
  426. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  427. defer cancelRead()
  428. messageType, event, readErr := clientConn.Read(readCtx)
  429. require.NoError(t, readErr)
  430. require.Equal(t, coderws.MessageText, messageType)
  431. return event
  432. }
  433. writeMessage(`{"type":"response.create","model":"gpt-5.6-sol","stream":true,"tools":[{"type":"custom","name":"exec"}],"input":"run pwd"}`)
  434. firstEvent := readMessage()
  435. require.Equal(t, "response.completed", gjson.GetBytes(firstEvent, "type").String())
  436. require.Equal(t, "custom_tool_call", gjson.GetBytes(firstEvent, "response.output.0.type").String())
  437. require.Equal(t, "pwd", gjson.GetBytes(firstEvent, "response.output.0.input").String())
  438. writeMessage(`{"type":"response.create","model":"gpt-5.6-sol","stream":true,"previous_response_id":"resp_custom_first","input":[{"type":"custom_tool_call_output","id":"ctco_client_output_1","call_id":"call_custom_1","output":"ok"}]}`)
  439. secondEvent := readMessage()
  440. require.Equal(t, "response.completed", gjson.GetBytes(secondEvent, "type").String())
  441. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  442. select {
  443. case proxyErr := <-errCh:
  444. require.NoError(t, proxyErr)
  445. case <-time.After(5 * time.Second):
  446. t.Fatal("timed out waiting for websocket bridge proxy to finish")
  447. }
  448. require.Len(t, upstream.bodies, 2)
  449. firstTools := gjson.GetBytes(upstream.bodies[0], "tools").Array()
  450. require.Len(t, firstTools, 1)
  451. require.Equal(t, "function", firstTools[0].Get("type").String())
  452. secondTools := gjson.GetBytes(upstream.bodies[1], "tools").Array()
  453. require.Len(t, secondTools, 1)
  454. require.Equal(t, "function", secondTools[0].Get("type").String())
  455. require.Equal(t, "exec", secondTools[0].Get("name").String())
  456. secondInput := gjson.GetBytes(upstream.bodies[1], "input").Array()
  457. require.Len(t, secondInput, 3)
  458. require.Equal(t, "run pwd", secondInput[0].String())
  459. require.Equal(t, "function_call", secondInput[1].Get("type").String())
  460. require.Equal(t, "fc_custom_1", secondInput[1].Get("id").String())
  461. require.JSONEq(t, `{"input":"pwd"}`, secondInput[1].Get("arguments").String())
  462. require.False(t, secondInput[1].Get("input").Exists())
  463. require.Equal(t, "function_call_output", secondInput[2].Get("type").String())
  464. require.False(t, secondInput[2].Get("id").Exists())
  465. }
  466. func TestOpenAIWSHTTPBridgeFullCustomToolHistoryWithoutPreviousResponseIDDoesNotReplay(t *testing.T) {
  467. gin.SetMode(gin.TestMode)
  468. completed := func(responseID string, output string) string {
  469. return "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"" + responseID + "\",\"model\":\"gpt-5.1\",\"output\":" + output + ",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n"
  470. }
  471. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  472. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(completed("resp_1", `[{"type":"custom_tool_call","id":"item_1","call_id":"call_1","name":"exec","input":"pwd"}]`)))},
  473. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(completed("resp_2", `[]`)))},
  474. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(completed("resp_3", `[]`)))},
  475. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(completed("resp_4", `[]`)))},
  476. }}
  477. cfg := &config.Config{}
  478. cfg.Security.URLAllowlist.Enabled = false
  479. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  480. cfg.Gateway.OpenAIWS.Enabled = true
  481. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  482. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  483. cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true
  484. cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = 1
  485. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  486. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  487. svc := &OpenAIGatewayService{
  488. cfg: cfg, httpUpstream: upstream, cache: &stubGatewayCache{},
  489. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), toolCorrector: NewCodexToolCorrector(),
  490. }
  491. account := &Account{
  492. ID: 9002, Name: "oauth-full-context", Platform: PlatformOpenAI, Type: AccountTypeOAuth,
  493. Credentials: map[string]any{"access_token": "test-token"}, Extra: map[string]any{"responses_websockets_v2_enabled": true},
  494. Concurrency: 1, Status: StatusActive, Schedulable: true,
  495. }
  496. errCh := make(chan error, 1)
  497. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  498. conn, err := coderws.Accept(w, r, nil)
  499. if err != nil {
  500. errCh <- err
  501. return
  502. }
  503. defer func() { _ = conn.CloseNow() }()
  504. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  505. _, firstMessage, err := conn.Read(readCtx)
  506. cancelRead()
  507. if err != nil {
  508. errCh <- err
  509. return
  510. }
  511. rec := httptest.NewRecorder()
  512. ginCtx, _ := gin.CreateTestContext(rec)
  513. ginCtx.Request = r.Clone(r.Context())
  514. errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "test-token", firstMessage, nil)
  515. }))
  516. defer wsServer.Close()
  517. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  518. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  519. cancelDial()
  520. require.NoError(t, err)
  521. defer func() { _ = clientConn.CloseNow() }()
  522. writeAndRead := func(payload string) {
  523. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  524. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  525. cancelWrite()
  526. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  527. _, event, readErr := clientConn.Read(readCtx)
  528. cancelRead()
  529. require.NoError(t, readErr)
  530. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  531. }
  532. writeAndRead(`{"type":"response.create","model":"gpt-5.1","input":"run pwd"}`)
  533. writeAndRead(`{"type":"response.create","model":"gpt-5.1","previous_response_id":"resp_1","input":[{"role":"user","content":"continue without tool output"}]}`)
  534. fullContext := `{"type":"response.create","model":"gpt-5.1","input":[{"type":"custom_tool_call","id":"item_1","call_id":"call_1","name":"exec","input":"pwd"},{"type":"custom_tool_call_output","call_id":"call_1","output":"/tmp"},{"role":"user","content":"continue"}]}`
  535. writeAndRead(fullContext)
  536. writeAndRead(fullContext)
  537. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  538. select {
  539. case proxyErr := <-errCh:
  540. require.NoError(t, proxyErr)
  541. case <-time.After(5 * time.Second):
  542. t.Fatal("timed out waiting for websocket bridge proxy to finish")
  543. }
  544. require.Len(t, upstream.bodies, 4)
  545. orphanInput := gjson.GetBytes(upstream.bodies[1], "input").Array()
  546. require.Len(t, orphanInput, 2)
  547. require.Equal(t, "run pwd", orphanInput[0].String())
  548. require.Equal(t, "user", orphanInput[1].Get("role").String())
  549. for _, body := range upstream.bodies[2:] {
  550. input := gjson.GetBytes(body, "input").Array()
  551. require.Len(t, input, 3)
  552. require.Equal(t, "custom_tool_call", input[0].Get("type").String())
  553. require.Equal(t, "call_1", input[0].Get("call_id").String())
  554. require.Equal(t, "custom_tool_call_output", input[1].Get("type").String())
  555. require.Equal(t, "call_1", input[1].Get("call_id").String())
  556. }
  557. }
  558. func TestOpenAIWSHTTPBridgeObjectToolOutputWithoutPreviousResponseIDReplaysMatchingCall(t *testing.T) {
  559. gin.SetMode(gin.TestMode)
  560. completed := func(responseID string, output string) string {
  561. return "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"" + responseID + "\",\"model\":\"gpt-5.1\",\"output\":" + output + ",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n"
  562. }
  563. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  564. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(completed("resp_1", `[{"type":"custom_tool_call","id":"item_1","call_id":"call_1","name":"exec","input":"pwd"}]`)))},
  565. {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(completed("resp_2", `[]`)))},
  566. }}
  567. cfg := &config.Config{}
  568. cfg.Security.URLAllowlist.Enabled = false
  569. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  570. cfg.Gateway.OpenAIWS.Enabled = true
  571. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  572. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  573. cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true
  574. cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = 1
  575. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  576. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  577. svc := &OpenAIGatewayService{
  578. cfg: cfg, httpUpstream: upstream, cache: &stubGatewayCache{},
  579. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), toolCorrector: NewCodexToolCorrector(),
  580. }
  581. account := &Account{
  582. ID: 9003, Name: "oauth-output-only", Platform: PlatformOpenAI, Type: AccountTypeOAuth,
  583. Credentials: map[string]any{"access_token": "test-token"}, Extra: map[string]any{"responses_websockets_v2_enabled": true},
  584. Concurrency: 1, Status: StatusActive, Schedulable: true,
  585. }
  586. errCh := make(chan error, 1)
  587. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  588. conn, err := coderws.Accept(w, r, nil)
  589. if err != nil {
  590. errCh <- err
  591. return
  592. }
  593. defer func() { _ = conn.CloseNow() }()
  594. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  595. _, firstMessage, err := conn.Read(readCtx)
  596. cancelRead()
  597. if err != nil {
  598. errCh <- err
  599. return
  600. }
  601. rec := httptest.NewRecorder()
  602. ginCtx, _ := gin.CreateTestContext(rec)
  603. ginCtx.Request = r.Clone(r.Context())
  604. errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "test-token", firstMessage, nil)
  605. }))
  606. defer wsServer.Close()
  607. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  608. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  609. cancelDial()
  610. require.NoError(t, err)
  611. defer func() { _ = clientConn.CloseNow() }()
  612. writeAndRead := func(payload string) {
  613. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  614. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  615. cancelWrite()
  616. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  617. _, event, readErr := clientConn.Read(readCtx)
  618. cancelRead()
  619. require.NoError(t, readErr)
  620. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  621. }
  622. writeAndRead(`{"type":"response.create","model":"gpt-5.1","input":"run pwd"}`)
  623. writeAndRead(`{"type":"response.create","model":"gpt-5.1","input":{"type":"custom_tool_call_output","call_id":"call_1","output":"/tmp"}}`)
  624. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  625. select {
  626. case proxyErr := <-errCh:
  627. require.NoError(t, proxyErr)
  628. case <-time.After(5 * time.Second):
  629. t.Fatal("timed out waiting for websocket bridge proxy to finish")
  630. }
  631. require.Len(t, upstream.bodies, 2)
  632. secondInput := gjson.GetBytes(upstream.bodies[1], "input").Array()
  633. require.Len(t, secondInput, 3)
  634. require.Equal(t, "custom_tool_call", secondInput[1].Get("type").String())
  635. require.Equal(t, "call_1", secondInput[1].Get("call_id").String())
  636. require.Equal(t, "custom_tool_call_output", secondInput[2].Get("type").String())
  637. require.Equal(t, "call_1", secondInput[2].Get("call_id").String())
  638. }
  639. func TestOpenAIWSHTTPBridgeDecisionKeepsSmallFramesOnWS(t *testing.T) {
  640. svc := &OpenAIGatewayService{
  641. cfg: &config.Config{
  642. Gateway: config.GatewayConfig{
  643. OpenAIWS: config.GatewayOpenAIWSConfig{
  644. HTTPBridgeEnabled: true,
  645. HTTPBridgeThresholdBytes: 100,
  646. },
  647. },
  648. },
  649. }
  650. require.False(t, svc.shouldBridgeOpenAIWSHTTP(nil, 99, ""))
  651. require.True(t, svc.shouldBridgeOpenAIWSHTTP(nil, 100, ""))
  652. require.False(t, svc.shouldBridgeOpenAIWSHTTP(nil, 1000, "resp_existing"))
  653. svc.cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = false
  654. require.False(t, svc.shouldBridgeOpenAIWSHTTP(nil, 1000, ""))
  655. require.True(t, svc.shouldBridgeOpenAIWSHTTP(&Account{Platform: PlatformGrok}, 1, "resp_existing"))
  656. }
  657. func TestOpenAIWSPassthroughFirstMessageBridgeDecision(t *testing.T) {
  658. const threshold = 100
  659. svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{
  660. OpenAIWS: config.GatewayOpenAIWSConfig{
  661. HTTPBridgeEnabled: true,
  662. HTTPBridgeThresholdBytes: threshold,
  663. },
  664. }}}
  665. exactThresholdPayload := `{"type":"response.create","input":"` +
  666. strings.Repeat("x", threshold-len(`{"type":"response.create","input":""}`)) + `"}`
  667. tests := []struct {
  668. name string
  669. payload string
  670. want bool
  671. }{
  672. {name: "oversized without previous response id bridges", payload: `{"type":"response.create","input":"` + strings.Repeat("x", 100) + `"}`, want: true},
  673. {name: "omitted type bridges as response create", payload: `{"input":"` + strings.Repeat("x", 100) + `"}`, want: true},
  674. {name: "blank type bridges as response create", payload: `{"type":" ","padding":"` + strings.Repeat("x", 100) + `"}`, want: true},
  675. {name: "exact threshold bridges", payload: exactThresholdPayload, want: true},
  676. {name: "small stays passthrough", payload: `{"type":"response.create","input":"x"}`},
  677. {name: "previous response id stays passthrough", payload: `{"type":"response.create","previous_response_id":"resp_previous","input":"` + strings.Repeat("x", 100) + `"}`},
  678. {name: "duplicate type stays passthrough", payload: `{"type":"response.create","type":"response.create","input":"` + strings.Repeat("x", 100) + `"}`},
  679. {name: "duplicate previous response id stays passthrough", payload: `{"type":"response.create","previous_response_id":null,"previous_response_id":null,"input":"` + strings.Repeat("x", 100) + `"}`},
  680. {name: "null type bridges as response create", payload: `{"type":null,"padding":"` + strings.Repeat("x", 100) + `"}`, want: true},
  681. {name: "non-string type stays passthrough", payload: `{"type":123,"padding":"` + strings.Repeat("x", 100) + `"}`},
  682. {name: "oversized response cancel stays passthrough", payload: `{"type":"response.cancel","padding":"` + strings.Repeat("x", 100) + `"}`},
  683. {name: "oversized other event stays passthrough", payload: `{"type":"session.update","padding":"` + strings.Repeat("x", 100) + `"}`},
  684. {name: "oversized malformed JSON stays passthrough", payload: `{"type":"response.create","padding":"` + strings.Repeat("x", 100)},
  685. }
  686. for _, tt := range tests {
  687. t.Run(tt.name, func(t *testing.T) {
  688. require.Equal(t, tt.want, svc.shouldBridgeOpenAIWSPassthroughFirstMessage(nil, []byte(tt.payload)))
  689. })
  690. }
  691. }
  692. func TestProxyOpenAIWSHTTPBridgeTurnTransportErrorFailoverSafety(t *testing.T) {
  693. gin.SetMode(gin.TestMode)
  694. tests := []struct {
  695. name string
  696. turn int
  697. wantFailover bool
  698. wantWrites int
  699. }{
  700. {name: "first_turn_fails_over_before_downstream_event", turn: 1, wantFailover: true},
  701. {name: "later_turn_does_not_replay_completed_turns", turn: 2, wantWrites: 1},
  702. }
  703. for _, tt := range tests {
  704. t.Run(tt.name, func(t *testing.T) {
  705. upstream := &httpUpstreamRecorder{err: io.EOF}
  706. svc := &OpenAIGatewayService{
  707. cfg: &config.Config{},
  708. httpUpstream: upstream,
  709. }
  710. account := &Account{
  711. ID: 8,
  712. Name: "api-key",
  713. Platform: PlatformOpenAI,
  714. Type: AccountTypeAPIKey,
  715. Concurrency: 1,
  716. }
  717. recorder := httptest.NewRecorder()
  718. c, _ := gin.CreateTestContext(recorder)
  719. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  720. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  721. var writes [][]byte
  722. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  723. context.Background(), c, account, "sk-test", payload, len(payload),
  724. "gpt-5", "", "", "", "", tt.turn,
  725. func(message []byte) error {
  726. writes = append(writes, append([]byte(nil), message...))
  727. return nil
  728. },
  729. )
  730. require.Nil(t, result)
  731. var failoverErr *UpstreamFailoverError
  732. if tt.wantFailover {
  733. require.ErrorAs(t, err, &failoverErr)
  734. require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
  735. require.JSONEq(t, string(openAITransportFailoverBody), string(failoverErr.ResponseBody))
  736. } else {
  737. require.Error(t, err)
  738. require.False(t, errors.As(err, &failoverErr))
  739. }
  740. require.Len(t, writes, tt.wantWrites)
  741. if tt.wantWrites > 0 {
  742. require.Equal(t, "error", gjson.GetBytes(writes[0], "type").String())
  743. require.Equal(t, int64(http.StatusBadGateway), gjson.GetBytes(writes[0], "status").Int())
  744. }
  745. })
  746. }
  747. }
  748. func TestProxyOpenAIWSHTTPBridgeTurnHTTPStatusFailoverSafety(t *testing.T) {
  749. gin.SetMode(gin.TestMode)
  750. tests := []struct {
  751. name string
  752. turn int
  753. status int
  754. wantFailover bool
  755. wantWrites int
  756. }{
  757. {name: "first_turn_401", turn: 1, status: http.StatusUnauthorized, wantFailover: true},
  758. {name: "first_turn_429", turn: 1, status: http.StatusTooManyRequests, wantFailover: true},
  759. {name: "first_turn_500", turn: 1, status: http.StatusInternalServerError, wantFailover: true},
  760. {name: "later_turn_500_does_not_replay", turn: 2, status: http.StatusInternalServerError, wantWrites: 1},
  761. }
  762. for _, tt := range tests {
  763. t.Run(tt.name, func(t *testing.T) {
  764. upstream := &httpUpstreamRecorder{resp: &http.Response{
  765. StatusCode: tt.status,
  766. Header: make(http.Header),
  767. Body: io.NopCloser(strings.NewReader(`{"error":{"type":"server_error","message":"temporary upstream failure"}}`)),
  768. }}
  769. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  770. account := &Account{ID: 9, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  771. recorder := httptest.NewRecorder()
  772. c, _ := gin.CreateTestContext(recorder)
  773. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  774. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  775. var writes [][]byte
  776. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  777. context.Background(), c, account, "sk-test", payload, len(payload),
  778. "gpt-5", "", "", "", "", tt.turn,
  779. func(message []byte) error {
  780. writes = append(writes, append([]byte(nil), message...))
  781. return nil
  782. },
  783. )
  784. require.Nil(t, result)
  785. var failoverErr *UpstreamFailoverError
  786. if tt.wantFailover {
  787. require.ErrorAs(t, err, &failoverErr)
  788. require.Equal(t, tt.status, failoverErr.StatusCode)
  789. } else {
  790. require.Error(t, err)
  791. require.False(t, errors.As(err, &failoverErr))
  792. }
  793. require.Len(t, writes, tt.wantWrites)
  794. })
  795. }
  796. }
  797. func TestProxyOpenAIWSHTTPBridgeTurnRetriesRejectedFieldBeforeClientOutput(t *testing.T) {
  798. gin.SetMode(gin.TestMode)
  799. upstream := &httpUpstreamSequenceRecorder{responses: []*http.Response{
  800. {
  801. StatusCode: http.StatusBadRequest,
  802. Header: make(http.Header),
  803. Body: io.NopCloser(strings.NewReader(
  804. `{"error":{"type":"invalid_request_error","code":"invalid_parameter","param":"truncation","message":"Unsupported parameter: truncation"}}`,
  805. )),
  806. },
  807. {
  808. StatusCode: http.StatusOK,
  809. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  810. Body: io.NopCloser(strings.NewReader(
  811. "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_retry\",\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n",
  812. )),
  813. },
  814. }}
  815. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  816. account := &Account{ID: 91, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  817. recorder := httptest.NewRecorder()
  818. c, _ := gin.CreateTestContext(recorder)
  819. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  820. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi","truncation":"auto"}`)
  821. var writes [][]byte
  822. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  823. context.Background(), c, account, "sk-test", payload, len(payload),
  824. "gpt-5", "", "", "", "", 1,
  825. func(message []byte) error {
  826. writes = append(writes, append([]byte(nil), message...))
  827. return nil
  828. },
  829. )
  830. require.NoError(t, err)
  831. require.NotNil(t, result)
  832. require.Equal(t, 2, upstream.callCount)
  833. require.Len(t, upstream.bodies, 2)
  834. require.True(t, gjson.GetBytes(upstream.bodies[0], "truncation").Exists())
  835. require.False(t, gjson.GetBytes(upstream.bodies[1], "truncation").Exists())
  836. require.Len(t, writes, 1)
  837. require.Equal(t, "response.completed", gjson.GetBytes(writes[0], "type").String())
  838. }
  839. func TestProxyOpenAIWSHTTPBridgeTurnSSEErrorFailoverSafety(t *testing.T) {
  840. gin.SetMode(gin.TestMode)
  841. for _, turn := range []int{1, 2} {
  842. t.Run(fmt.Sprintf("turn_%d", turn), func(t *testing.T) {
  843. upstream := &httpUpstreamRecorder{resp: &http.Response{
  844. StatusCode: http.StatusOK,
  845. Header: make(http.Header),
  846. Body: io.NopCloser(strings.NewReader(
  847. "data: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"code\":\"rate_limit_exceeded\",\"message\":\"limited\"}}\n\n",
  848. )),
  849. }}
  850. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  851. account := &Account{ID: 10, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  852. recorder := httptest.NewRecorder()
  853. c, _ := gin.CreateTestContext(recorder)
  854. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  855. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  856. var writes [][]byte
  857. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  858. context.Background(), c, account, "sk-test", payload, len(payload),
  859. "gpt-5", "", "", "", "", turn,
  860. func(message []byte) error {
  861. writes = append(writes, append([]byte(nil), message...))
  862. return nil
  863. },
  864. )
  865. var failoverErr *UpstreamFailoverError
  866. require.Nil(t, result)
  867. require.ErrorAs(t, err, &failoverErr)
  868. require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
  869. require.Empty(t, writes)
  870. })
  871. }
  872. }
  873. // 桥接转发 error / response.failed 给 WS 客户端前必须把容量降载码改写为可重试
  874. // 的 server_error:Codex 对 server_is_overloaded/slow_down 判致命并终止会话。
  875. // 账号状态判定使用改写前的原始事件,不受影响。
  876. func TestProxyOpenAIWSHTTPBridgeTurnRewritesCapacityShedCodeForClient(t *testing.T) {
  877. gin.SetMode(gin.TestMode)
  878. tests := []struct {
  879. name string
  880. turn int
  881. body string
  882. wantErr bool
  883. }{
  884. {
  885. name: "turn2_error_frame",
  886. turn: 2,
  887. body: "data: {\"type\":\"error\",\"error\":{\"type\":\"service_unavailable_error\",\"code\":\"server_is_overloaded\",\"message\":\"Our servers are currently overloaded. Please try again later.\"}}\n\n",
  888. wantErr: true,
  889. },
  890. {
  891. // 后续 turn 不允许 replay,容量错误必须改写后交给客户端重试。
  892. name: "turn2_bare_response_failed",
  893. turn: 2,
  894. body: "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_shed\",\"status\":\"failed\",\"error\":{\"code\":\"server_is_overloaded\",\"message\":\"Our servers are currently overloaded. Please try again later.\"}}}\n\n",
  895. },
  896. }
  897. for _, tt := range tests {
  898. t.Run(tt.name, func(t *testing.T) {
  899. upstream := &httpUpstreamRecorder{resp: &http.Response{
  900. StatusCode: http.StatusOK,
  901. Header: make(http.Header),
  902. Body: io.NopCloser(strings.NewReader(tt.body)),
  903. }}
  904. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  905. account := &Account{ID: 11, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1}
  906. recorder := httptest.NewRecorder()
  907. c, _ := gin.CreateTestContext(recorder)
  908. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  909. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  910. var writes [][]byte
  911. _, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  912. context.Background(), c, account, "sk-test", payload, len(payload),
  913. "gpt-5", "", "", "", "", tt.turn,
  914. func(message []byte) error {
  915. writes = append(writes, append([]byte(nil), message...))
  916. return nil
  917. },
  918. )
  919. if tt.wantErr {
  920. require.Error(t, err)
  921. } else {
  922. require.NoError(t, err)
  923. }
  924. require.Len(t, writes, 1)
  925. require.Contains(t, string(writes[0]), `"code":"server_error"`)
  926. require.NotContains(t, string(writes[0]), "server_is_overloaded")
  927. require.Contains(t, string(writes[0]), "Our servers are currently overloaded")
  928. })
  929. }
  930. }
  931. func TestProxyOpenAIWSHTTPBridgeTurnBareErrorUsesAuthoritativeFailed(t *testing.T) {
  932. gin.SetMode(gin.TestMode)
  933. body := strings.Join([]string{
  934. `data: {"type":"response.output_text.delta","response_id":"resp_failed","delta":"partial"}`,
  935. ``,
  936. `data: {"type":"error","error":{"status_code":403,"code":"workspace_suspended","message":"workspace is suspended"}}`,
  937. ``,
  938. `data: {"type":"response.failed","response":{"id":"resp_failed","status":"failed","usage":{"input_tokens":9,"output_tokens":2},"error":{"status_code":403,"code":"workspace_suspended","message":"workspace is suspended"}}}`,
  939. ``,
  940. }, "\n")
  941. upstream := &httpUpstreamRecorder{resp: &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}}
  942. repo := &openAIStream403AccountRepo{}
  943. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream, rateLimitService: &RateLimitService{accountRepo: repo}}
  944. account := &Account{ID: 111, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  945. recorder := httptest.NewRecorder()
  946. c, _ := gin.CreateTestContext(recorder)
  947. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  948. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  949. var writes [][]byte
  950. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(context.Background(), c, account, "sk-test", payload, len(payload), "gpt-5", "", "", "", "", 2, func(message []byte) error {
  951. writes = append(writes, append([]byte(nil), message...))
  952. return nil
  953. })
  954. require.NoError(t, err)
  955. require.NotNil(t, result)
  956. require.Equal(t, 9, result.Usage.InputTokens)
  957. require.Equal(t, 2, result.Usage.OutputTokens)
  958. require.Equal(t, 1, repo.setErrorCalls)
  959. require.Len(t, writes, 2)
  960. require.Equal(t, "response.output_text.delta", gjson.GetBytes(writes[0], "type").String())
  961. require.Equal(t, "response.failed", gjson.GetBytes(writes[1], "type").String())
  962. }
  963. func TestProxyOpenAIWSHTTPBridgeTurnMarksCyberPolicyForFailureShapes(t *testing.T) {
  964. gin.SetMode(gin.TestMode)
  965. tests := []struct {
  966. name string
  967. statusCode int
  968. body string
  969. wantStatus int
  970. wantInput int
  971. wantOutput int
  972. wantResult bool
  973. wantError bool
  974. }{
  975. {
  976. name: "http_error_body",
  977. statusCode: http.StatusBadRequest,
  978. body: `{"error":{"code":"cyber_policy","message":"blocked by HTTP cyber policy"}}`,
  979. wantStatus: http.StatusBadRequest,
  980. wantError: true,
  981. },
  982. {
  983. name: "sse_error_event",
  984. statusCode: http.StatusOK,
  985. body: "data: {\"type\":\"error\",\"error\":{\"code\":\"cyber_policy\",\"message\":\"blocked by error event\"},\"usage\":{\"input_tokens\":5,\"output_tokens\":1}}\n\n",
  986. wantStatus: http.StatusOK,
  987. wantInput: 5,
  988. wantOutput: 1,
  989. wantResult: true,
  990. wantError: true,
  991. },
  992. {
  993. name: "sse_response_failed_event",
  994. statusCode: http.StatusOK,
  995. body: "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_cyber\",\"error\":{\"code\":\"cyber_policy\",\"message\":\"blocked by failed event\"},\"usage\":{\"input_tokens\":9,\"output_tokens\":2}}}\n\n",
  996. wantStatus: http.StatusOK,
  997. wantInput: 9,
  998. wantOutput: 2,
  999. wantResult: true,
  1000. },
  1001. }
  1002. for _, tt := range tests {
  1003. t.Run(tt.name, func(t *testing.T) {
  1004. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1005. StatusCode: tt.statusCode,
  1006. Header: make(http.Header),
  1007. Body: io.NopCloser(strings.NewReader(tt.body)),
  1008. }}
  1009. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  1010. account := &Account{ID: 112, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  1011. recorder := httptest.NewRecorder()
  1012. c, _ := gin.CreateTestContext(recorder)
  1013. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  1014. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  1015. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  1016. context.Background(), c, account, "sk-test", payload, len(payload),
  1017. "gpt-5", "", "", "", "", 1, func([]byte) error { return nil },
  1018. )
  1019. require.Equal(t, tt.wantResult, result != nil)
  1020. require.Equal(t, tt.wantError, err != nil)
  1021. mark := GetOpsCyberPolicy(c)
  1022. require.NotNil(t, mark)
  1023. require.Equal(t, "cyber_policy", mark.Code)
  1024. require.Equal(t, tt.wantStatus, mark.UpstreamStatus)
  1025. require.Equal(t, tt.wantInput, mark.UpstreamInTok)
  1026. require.Equal(t, tt.wantOutput, mark.UpstreamOutTok)
  1027. require.Contains(t, mark.Body, `"code":"cyber_policy"`)
  1028. })
  1029. }
  1030. }
  1031. func TestProxyOpenAIWSHTTPBridgeTurnBareErrorEOFSynthesizesFailed(t *testing.T) {
  1032. gin.SetMode(gin.TestMode)
  1033. body := "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_eof\",\"status\":\"in_progress\"}}\n\n" +
  1034. "data: {\"type\":\"error\",\"error\":{\"code\":\"invalid_request\",\"message\":\"bad request\"}}\n\n"
  1035. upstream := &httpUpstreamRecorder{resp: &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}}
  1036. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  1037. account := &Account{ID: 112, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  1038. recorder := httptest.NewRecorder()
  1039. c, _ := gin.CreateTestContext(recorder)
  1040. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  1041. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  1042. var writes [][]byte
  1043. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(context.Background(), c, account, "sk-test", payload, len(payload), "gpt-5", "", "", "", "", 2, func(message []byte) error {
  1044. writes = append(writes, append([]byte(nil), message...))
  1045. return nil
  1046. })
  1047. require.EqualError(t, err, "bad request")
  1048. require.NotNil(t, result)
  1049. require.Equal(t, "response.failed", result.UpstreamTerminalEvent)
  1050. require.Len(t, writes, 2)
  1051. require.Equal(t, "response.failed", gjson.GetBytes(writes[1], "type").String())
  1052. require.Equal(t, "failed", gjson.GetBytes(writes[1], "response.status").String())
  1053. require.Equal(t, "resp_eof", gjson.GetBytes(writes[1], "response.id").String())
  1054. }
  1055. func TestProxyOpenAIWSHTTPBridgeTurnBareErrorFollowedByCompletedUsesCompleted(t *testing.T) {
  1056. gin.SetMode(gin.TestMode)
  1057. body := strings.Join([]string{
  1058. `data: {"type":"response.created","response":{"id":"resp_recovered","status":"in_progress"}}`,
  1059. ``,
  1060. `data: {"type":"error","error":{"code":"transient","message":"retrying"}}`,
  1061. ``,
  1062. `data: {"type":"response.completed","response":{"id":"resp_recovered","status":"completed","output":[],"usage":{"input_tokens":8,"output_tokens":4}}}`,
  1063. ``,
  1064. }, "\n")
  1065. upstream := &httpUpstreamRecorder{resp: &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}}
  1066. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  1067. account := &Account{ID: 113, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1}
  1068. recorder := httptest.NewRecorder()
  1069. c, _ := gin.CreateTestContext(recorder)
  1070. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  1071. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  1072. var writes [][]byte
  1073. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(context.Background(), c, account, "sk-test", payload, len(payload), "gpt-5", "", "", "", "", 2, func(message []byte) error {
  1074. writes = append(writes, append([]byte(nil), message...))
  1075. return nil
  1076. })
  1077. require.NoError(t, err)
  1078. require.NotNil(t, result)
  1079. require.Equal(t, "response.completed", result.UpstreamTerminalEvent)
  1080. require.Equal(t, 8, result.Usage.InputTokens)
  1081. require.Equal(t, 4, result.Usage.OutputTokens)
  1082. require.Len(t, writes, 2)
  1083. require.Equal(t, "response.created", gjson.GetBytes(writes[0], "type").String())
  1084. require.Equal(t, "response.completed", gjson.GetBytes(writes[1], "type").String())
  1085. }
  1086. func TestProxyOpenAIWSHTTPBridgeTurnStagesMetadataBeforeCapacityFailover(t *testing.T) {
  1087. gin.SetMode(gin.TestMode)
  1088. body := strings.Join([]string{
  1089. `data: {"type":"response.created","response":{"id":"resp_shed"}}`,
  1090. "",
  1091. `data: {"type":"response.in_progress","response":{"id":"resp_shed"}}`,
  1092. "",
  1093. `data: {"type":"response.failed","response":{"id":"resp_shed","status":"failed","error":{"message":"Our servers are currently overloaded. Please try again later."}}}`,
  1094. "",
  1095. }, "\n")
  1096. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1097. StatusCode: http.StatusOK,
  1098. Header: http.Header{"X-Request-Id": []string{"rid-ws-bridge-capacity"}},
  1099. Body: io.NopCloser(strings.NewReader(body)),
  1100. }}
  1101. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  1102. account := &Account{ID: 12, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1}
  1103. recorder := httptest.NewRecorder()
  1104. c, _ := gin.CreateTestContext(recorder)
  1105. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  1106. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  1107. var writes [][]byte
  1108. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  1109. context.Background(), c, account, "sk-test", payload, len(payload),
  1110. "gpt-5", "", "", "", "", 1,
  1111. func(message []byte) error {
  1112. writes = append(writes, append([]byte(nil), message...))
  1113. return nil
  1114. },
  1115. )
  1116. require.Nil(t, result)
  1117. var failoverErr *UpstreamFailoverError
  1118. require.ErrorAs(t, err, &failoverErr)
  1119. require.True(t, failoverErr.RetryableOnSameAccount)
  1120. require.True(t, failoverErr.RequestScopedTransient)
  1121. require.Empty(t, writes)
  1122. }
  1123. func TestProxyOpenAIWSHTTPBridgeTurnDoesNotReplayCapacityAfterSemanticOutput(t *testing.T) {
  1124. gin.SetMode(gin.TestMode)
  1125. logSink, restore := captureStructuredLog(t)
  1126. defer restore()
  1127. body := strings.Join([]string{
  1128. `data: {"type":"response.created","response":{"id":"resp_partial"}}`,
  1129. "",
  1130. `data: {"type":"response.output_text.delta","delta":"partial"}`,
  1131. "",
  1132. `data: {"type":"response.failed","response":{"id":"resp_partial","status":"failed","error":{"code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}}}`,
  1133. "",
  1134. }, "\n")
  1135. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1136. StatusCode: http.StatusOK,
  1137. Header: http.Header{"X-Request-Id": []string{"rid-ws-bridge-post-output"}},
  1138. Body: io.NopCloser(strings.NewReader(body)),
  1139. }}
  1140. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  1141. account := &Account{ID: 13, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1}
  1142. recorder := httptest.NewRecorder()
  1143. c, _ := gin.CreateTestContext(recorder)
  1144. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  1145. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  1146. var writes [][]byte
  1147. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  1148. context.Background(), c, account, "sk-test", payload, len(payload),
  1149. "gpt-5", "", "", "", "", 1,
  1150. func(message []byte) error {
  1151. writes = append(writes, append([]byte(nil), message...))
  1152. return nil
  1153. },
  1154. )
  1155. require.NotNil(t, result)
  1156. require.NoError(t, err)
  1157. require.Len(t, writes, 3)
  1158. var failoverErr *UpstreamFailoverError
  1159. require.False(t, errors.As(err, &failoverErr))
  1160. require.Contains(t, string(writes[2]), `"code":"server_error"`)
  1161. require.NotContains(t, string(writes[2]), "server_is_overloaded")
  1162. require.True(t, logSink.ContainsMessage("gateway.failover_suppressed_after_semantic_output"))
  1163. require.True(t, logSink.ContainsFieldValue("path", "ws_http_bridge"))
  1164. }
  1165. func TestProxyOpenAIWSHTTPBridgeTurnRequiresTerminalEvent(t *testing.T) {
  1166. gin.SetMode(gin.TestMode)
  1167. tests := []struct {
  1168. name string
  1169. body string
  1170. wantFailover bool
  1171. wantWrites int
  1172. }{
  1173. {name: "done_without_events_fails_over", body: "data: [DONE]\n\n", wantFailover: true},
  1174. {
  1175. name: "created_then_done_fails_over_before_semantic_output",
  1176. body: "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_truncated\"}}\n\n" +
  1177. "data: [DONE]\n\n",
  1178. wantFailover: true,
  1179. },
  1180. }
  1181. for _, tt := range tests {
  1182. t.Run(tt.name, func(t *testing.T) {
  1183. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1184. StatusCode: http.StatusOK,
  1185. Header: make(http.Header),
  1186. Body: io.NopCloser(strings.NewReader(tt.body)),
  1187. }}
  1188. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  1189. account := &Account{ID: 11, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  1190. recorder := httptest.NewRecorder()
  1191. c, _ := gin.CreateTestContext(recorder)
  1192. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  1193. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  1194. var writes [][]byte
  1195. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  1196. context.Background(), c, account, "sk-test", payload, len(payload),
  1197. "gpt-5", "", "", "", "", 1,
  1198. func(message []byte) error {
  1199. writes = append(writes, append([]byte(nil), message...))
  1200. return nil
  1201. },
  1202. )
  1203. var failoverErr *UpstreamFailoverError
  1204. if tt.wantFailover {
  1205. require.Nil(t, result)
  1206. require.ErrorAs(t, err, &failoverErr)
  1207. } else {
  1208. require.NotNil(t, result)
  1209. require.Error(t, err)
  1210. require.False(t, errors.As(err, &failoverErr))
  1211. }
  1212. require.Len(t, writes, tt.wantWrites)
  1213. })
  1214. }
  1215. }
  1216. func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) {
  1217. gin.SetMode(gin.TestMode)
  1218. sseBody := strings.Join([]string{
  1219. `data: {"type":"response.created","response":{"id":"resp_bridge","model":"gpt-5"}}`,
  1220. "",
  1221. `data: {"type":"response.output_text.delta","response":{"id":"resp_bridge"},"delta":"ok"}`,
  1222. "",
  1223. `data: {"type":"response.completed","response":{"id":"resp_bridge","model":"gpt-5","usage":{"input_tokens":3,"output_tokens":2}}}`,
  1224. "",
  1225. }, "\n")
  1226. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1227. StatusCode: http.StatusOK,
  1228. Header: http.Header{
  1229. "Content-Type": []string{"text/event-stream"},
  1230. "x-request-id": []string{"rid_bridge"},
  1231. },
  1232. Body: io.NopCloser(strings.NewReader(sseBody)),
  1233. }}
  1234. svc := &OpenAIGatewayService{
  1235. cfg: &config.Config{
  1236. Gateway: config.GatewayConfig{
  1237. MaxLineSize: defaultMaxLineSize,
  1238. OpenAIWS: config.GatewayOpenAIWSConfig{
  1239. HTTPBridgeEnabled: true,
  1240. HTTPBridgeThresholdBytes: 1,
  1241. },
  1242. },
  1243. },
  1244. httpUpstream: upstream,
  1245. toolCorrector: NewCodexToolCorrector(),
  1246. }
  1247. account := &Account{
  1248. ID: 7,
  1249. Name: "api-key",
  1250. Platform: PlatformOpenAI,
  1251. Type: AccountTypeAPIKey,
  1252. Concurrency: 1,
  1253. Status: StatusActive,
  1254. }
  1255. payload := []byte(`{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"},"input":"hi","parallel_tool_calls":true}`)
  1256. type bridgeResult struct {
  1257. result *OpenAIForwardResult
  1258. err error
  1259. }
  1260. resultCh := make(chan bridgeResult, 1)
  1261. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1262. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  1263. if err != nil {
  1264. resultCh <- bridgeResult{err: err}
  1265. return
  1266. }
  1267. defer func() { _ = conn.CloseNow() }()
  1268. rec := httptest.NewRecorder()
  1269. ginCtx, _ := gin.CreateTestContext(rec)
  1270. req := r.Clone(r.Context())
  1271. req.Header = req.Header.Clone()
  1272. ginCtx.Request = req
  1273. writeClient := func(message []byte) error {
  1274. writeCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
  1275. defer cancel()
  1276. return conn.Write(writeCtx, coderws.MessageText, message)
  1277. }
  1278. result, bridgeErr := svc.proxyOpenAIWSHTTPBridgeTurn(
  1279. r.Context(),
  1280. ginCtx,
  1281. account,
  1282. "sk-test",
  1283. payload,
  1284. len(payload),
  1285. "gpt-5",
  1286. "",
  1287. "",
  1288. "",
  1289. "",
  1290. 1,
  1291. writeClient,
  1292. )
  1293. resultCh <- bridgeResult{result: result, err: bridgeErr}
  1294. }))
  1295. defer wsServer.Close()
  1296. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1297. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1298. cancelDial()
  1299. require.NoError(t, err)
  1300. defer func() { _ = clientConn.CloseNow() }()
  1301. readEvent := func() []byte {
  1302. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  1303. msgType, event, readErr := clientConn.Read(readCtx)
  1304. cancelRead()
  1305. require.NoError(t, readErr)
  1306. require.Equal(t, coderws.MessageText, msgType)
  1307. return event
  1308. }
  1309. created := readEvent()
  1310. delta := readEvent()
  1311. completed := readEvent()
  1312. require.Equal(t, "response.created", gjson.GetBytes(created, "type").String())
  1313. require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String())
  1314. require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String())
  1315. select {
  1316. case bridge := <-resultCh:
  1317. require.NoError(t, bridge.err)
  1318. require.NotNil(t, bridge.result)
  1319. require.Equal(t, "resp_bridge", bridge.result.RequestID)
  1320. require.Equal(t, 3, bridge.result.Usage.InputTokens)
  1321. require.Equal(t, 2, bridge.result.Usage.OutputTokens)
  1322. require.True(t, bridge.result.OpenAIWSMode)
  1323. case <-time.After(3 * time.Second):
  1324. t.Fatal("timed out waiting for bridge result")
  1325. }
  1326. require.NotNil(t, upstream.lastReq)
  1327. require.Equal(t, http.MethodPost, upstream.lastReq.Method)
  1328. require.Equal(t, "true", upstream.lastReq.Header.Get(responsesLiteHeader))
  1329. require.False(t, gjson.GetBytes(upstream.lastBody, "type").Exists())
  1330. require.False(t, gjson.GetBytes(upstream.lastBody, "generate").Exists())
  1331. require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
  1332. require.True(t, gjson.GetBytes(upstream.lastBody, "parallel_tool_calls").Exists())
  1333. require.False(t, gjson.GetBytes(upstream.lastBody, "parallel_tool_calls").Bool())
  1334. }
  1335. func TestProxyOpenAIWSHTTPBridgeTurnForGrokDefaultsEmptyModelTo45(t *testing.T) {
  1336. gin.SetMode(gin.TestMode)
  1337. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1338. StatusCode: http.StatusOK,
  1339. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  1340. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1341. `data: {"type":"response.created","response":{"id":"resp_grok_default","model":"grok-4.5"}}`,
  1342. "",
  1343. `data: {"type":"response.completed","response":{"id":"resp_grok_default","model":"grok-4.5","usage":{"input_tokens":1,"output_tokens":1}}}`,
  1344. "",
  1345. }, "\n"))),
  1346. }}
  1347. svc := &OpenAIGatewayService{
  1348. cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
  1349. httpUpstream: upstream,
  1350. }
  1351. account := &Account{
  1352. ID: 72,
  1353. Platform: PlatformGrok,
  1354. Type: AccountTypeOAuth,
  1355. Concurrency: 1,
  1356. Credentials: map[string]any{"base_url": xai.DefaultCLIBaseURL},
  1357. }
  1358. payload := []byte(`{"type":"response.create","generate":true,"stream":true,"input":"hi"}`)
  1359. recorder := httptest.NewRecorder()
  1360. c, _ := gin.CreateTestContext(recorder)
  1361. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  1362. var events [][]byte
  1363. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  1364. context.Background(), c, account, "access-token", payload, len(payload),
  1365. "", "", "", "", "", 1,
  1366. func(message []byte) error {
  1367. events = append(events, append([]byte(nil), message...))
  1368. return nil
  1369. },
  1370. )
  1371. require.NoError(t, err)
  1372. require.NotNil(t, result)
  1373. require.Equal(t, grokDefaultResponsesModel, gjson.GetBytes(upstream.lastBody, "model").String())
  1374. require.Len(t, events, 2)
  1375. }
  1376. func TestProxyOpenAIWSHTTPBridgeTurnPromotesCodexAdditionalToolsForMixedCache(t *testing.T) {
  1377. gin.SetMode(gin.TestMode)
  1378. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1379. StatusCode: http.StatusOK,
  1380. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  1381. Body: io.NopCloser(strings.NewReader(strings.Join([]string{
  1382. `data: {"type":"response.created","response":{"id":"resp_grok_codex_lite","model":"grok-4.5"}}`,
  1383. "",
  1384. `data: {"type":"response.completed","response":{"id":"resp_grok_codex_lite","model":"grok-4.5","usage":{"input_tokens":4,"output_tokens":1}}}`,
  1385. "",
  1386. }, "\n"))),
  1387. }}
  1388. svc := &OpenAIGatewayService{
  1389. cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
  1390. httpUpstream: upstream,
  1391. }
  1392. account := &Account{
  1393. ID: 73,
  1394. Platform: PlatformGrok,
  1395. Type: AccountTypeOAuth,
  1396. Concurrency: 1,
  1397. Credentials: map[string]any{
  1398. "base_url": xai.DefaultCLIBaseURL,
  1399. "subscription_tier": "free",
  1400. },
  1401. }
  1402. payload := []byte(`{
  1403. "type":"response.create","generate":true,"model":"grok","stream":true,
  1404. "input":[
  1405. {"type":"additional_tools","role":"developer","tools":[
  1406. {"type":"function","name":"lookup","parameters":{"type":"object"}},
  1407. {"type":"function","name":"web_search","parameters":{"type":"object"}},
  1408. {"type":"custom","name":"apply_patch"},
  1409. {"type":"namespace","name":"collaboration"}
  1410. ]},
  1411. {"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}
  1412. ]
  1413. }`)
  1414. recorder := httptest.NewRecorder()
  1415. c, _ := gin.CreateTestContext(recorder)
  1416. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  1417. c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache")
  1418. var events [][]byte
  1419. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  1420. context.Background(), c, account, "access-token", payload, len(payload),
  1421. "grok", "", "", "", "isolated-ws-cache-id", 1,
  1422. func(message []byte) error {
  1423. events = append(events, append([]byte(nil), message...))
  1424. return nil
  1425. },
  1426. )
  1427. require.NoError(t, err)
  1428. require.NotNil(t, result)
  1429. require.Len(t, events, 2)
  1430. require.False(t, gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools")`).Exists())
  1431. tools := gjson.GetBytes(upstream.lastBody, "tools").Array()
  1432. require.Len(t, tools, 4)
  1433. require.Equal(t, "function", tools[0].Get("type").String())
  1434. require.Equal(t, "lookup", tools[0].Get("name").String())
  1435. require.Equal(t, "web_search", tools[1].Get("type").String())
  1436. require.Equal(t, "function", tools[2].Get("type").String())
  1437. require.Equal(t, "apply_patch", tools[2].Get("name").String())
  1438. require.Equal(t, "x_search", tools[3].Get("type").String())
  1439. require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
  1440. require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="custom")`).Exists())
  1441. require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="namespace")`).Exists())
  1442. require.Equal(t, "isolated-ws-cache-id", gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
  1443. require.Equal(t, "isolated-ws-cache-id", upstream.lastReq.Header.Get(grokConversationIDHeader))
  1444. require.Empty(t, upstream.lastReq.Header.Get(grokClientToolCacheOptInHeader))
  1445. }
  1446. func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridgeAndPreservesMappedModels(t *testing.T) {
  1447. gin.SetMode(gin.TestMode)
  1448. bridgeResponse := func(responseID, requestID string, cachedTokens int) *http.Response {
  1449. sseBody := strings.Join([]string{
  1450. `data: {"type":"response.created","response":{"id":"` + responseID + `","model":"grok-4.3"}}`,
  1451. "",
  1452. `data: {"type":"response.output_text.delta","response":{"id":"` + responseID + `"},"delta":"ok"}`,
  1453. "",
  1454. `data: {"type":"response.completed","response":{"id":"` + responseID + `","model":"grok-4.3","usage":{"input_tokens":4,"output_tokens":2,"input_tokens_details":{"cached_tokens":` + fmt.Sprintf("%d", cachedTokens) + `}}}}`,
  1455. "",
  1456. }, "\n")
  1457. return &http.Response{
  1458. StatusCode: http.StatusOK,
  1459. Header: http.Header{
  1460. "Content-Type": []string{"text/event-stream"},
  1461. "Xai-Request-Id": []string{requestID},
  1462. },
  1463. Body: io.NopCloser(strings.NewReader(sseBody)),
  1464. }
  1465. }
  1466. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  1467. bridgeResponse("resp_grok_ws_1", "xai-ws-req-1", 0),
  1468. bridgeResponse("resp_grok_ws_2", "xai-ws-req-2", 3),
  1469. bridgeResponse("resp_grok_ws_3", "xai-ws-req-3", 0),
  1470. }}
  1471. svc := &OpenAIGatewayService{
  1472. cfg: &config.Config{
  1473. Gateway: config.GatewayConfig{
  1474. MaxLineSize: defaultMaxLineSize,
  1475. },
  1476. },
  1477. httpUpstream: upstream,
  1478. }
  1479. account := &Account{
  1480. ID: 71,
  1481. Name: "grok",
  1482. Platform: PlatformGrok,
  1483. Type: AccountTypeOAuth,
  1484. Concurrency: 1,
  1485. Status: StatusActive,
  1486. Credentials: map[string]any{
  1487. "base_url": xai.DefaultCLIBaseURL,
  1488. },
  1489. }
  1490. errCh := make(chan error, 1)
  1491. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1492. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  1493. if err != nil {
  1494. errCh <- err
  1495. return
  1496. }
  1497. defer func() { _ = conn.CloseNow() }()
  1498. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  1499. msgType, firstMessage, err := conn.Read(readCtx)
  1500. cancelRead()
  1501. if err != nil {
  1502. errCh <- err
  1503. return
  1504. }
  1505. if msgType != coderws.MessageText {
  1506. errCh <- errors.New("first message was not text")
  1507. return
  1508. }
  1509. rec := httptest.NewRecorder()
  1510. ginCtx, _ := gin.CreateTestContext(rec)
  1511. req := r.Clone(r.Context())
  1512. req.Header = req.Header.Clone()
  1513. ginCtx.Request = req
  1514. ginCtx.Set("api_key", &APIKey{ID: 7101})
  1515. errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "access-token", firstMessage, &OpenAIWSIngressHooks{
  1516. MapRequestModel: func(_ int, originalModel string) (string, error) {
  1517. if originalModel == "channel-alias" {
  1518. return "grok-4.3", nil
  1519. }
  1520. return originalModel, nil
  1521. },
  1522. })
  1523. }))
  1524. defer wsServer.Close()
  1525. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1526. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1527. cancelDial()
  1528. require.NoError(t, err)
  1529. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  1530. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","generate":true,"model":"grok","stream":true,"input":"hi","prompt_cache_retention":"24h"}`))
  1531. cancelWrite()
  1532. require.NoError(t, err)
  1533. readEvent := func() []byte {
  1534. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  1535. msgType, event, readErr := clientConn.Read(readCtx)
  1536. cancelRead()
  1537. require.NoError(t, readErr)
  1538. require.Equal(t, coderws.MessageText, msgType)
  1539. return event
  1540. }
  1541. created := readEvent()
  1542. delta := readEvent()
  1543. completed := readEvent()
  1544. require.Equal(t, "response.created", gjson.GetBytes(created, "type").String())
  1545. require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String())
  1546. require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String())
  1547. require.Equal(t, "resp_grok_ws_1", gjson.GetBytes(completed, "response.id").String())
  1548. writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
  1549. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","generate":true,"model":"channel-alias","stream":true,"previous_response_id":"resp_grok_ws_1","input":"second turn"}`))
  1550. cancelWrite()
  1551. require.NoError(t, err)
  1552. created = readEvent()
  1553. delta = readEvent()
  1554. completed = readEvent()
  1555. require.Equal(t, "response.created", gjson.GetBytes(created, "type").String())
  1556. require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String())
  1557. require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String())
  1558. require.Equal(t, "resp_grok_ws_2", gjson.GetBytes(completed, "response.id").String())
  1559. require.Equal(t, "channel-alias", gjson.GetBytes(completed, "response.model").String())
  1560. writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
  1561. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","generate":true,"model":"grok-4.3","stream":true,"previous_response_id":"resp_grok_ws_2","input":"third turn with a different model"}`))
  1562. cancelWrite()
  1563. require.NoError(t, err)
  1564. created = readEvent()
  1565. delta = readEvent()
  1566. completed = readEvent()
  1567. require.Equal(t, "response.created", gjson.GetBytes(created, "type").String())
  1568. require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String())
  1569. require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String())
  1570. require.Equal(t, "resp_grok_ws_3", gjson.GetBytes(completed, "response.id").String())
  1571. _ = clientConn.Close(coderws.StatusNormalClosure, "done")
  1572. select {
  1573. case proxyErr := <-errCh:
  1574. require.NoError(t, proxyErr)
  1575. case <-time.After(3 * time.Second):
  1576. require.Fail(t, "proxy did not finish after client close")
  1577. }
  1578. require.Len(t, upstream.requests, 3)
  1579. require.Len(t, upstream.bodies, 3)
  1580. require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
  1581. require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
  1582. require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), upstream.lastReq.Header.Get("User-Agent"))
  1583. require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
  1584. require.Equal(t, "grok-4.6", gjson.GetBytes(upstream.bodies[0], "model").String())
  1585. require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.bodies[1], "model").String())
  1586. require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.bodies[2], "model").String())
  1587. require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
  1588. require.Equal(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String(), upstream.lastReq.Header.Get(grokConversationIDHeader))
  1589. require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
  1590. require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String())
  1591. require.Equal(t, "none", gjson.GetBytes(upstream.lastBody, "tool_choice").String())
  1592. firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
  1593. secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
  1594. thirdIdentity := gjson.GetBytes(upstream.bodies[2], "prompt_cache_key").String()
  1595. require.NotEmpty(t, firstIdentity)
  1596. require.NotEmpty(t, secondIdentity)
  1597. require.NotEqual(t, firstIdentity, secondIdentity)
  1598. require.NotEmpty(t, thirdIdentity)
  1599. require.Equal(t, secondIdentity, thirdIdentity)
  1600. require.Equal(t, firstIdentity, upstream.requests[0].Header.Get(grokConversationIDHeader))
  1601. require.Equal(t, secondIdentity, upstream.requests[1].Header.Get(grokConversationIDHeader))
  1602. require.Equal(t, thirdIdentity, upstream.requests[2].Header.Get(grokConversationIDHeader))
  1603. require.False(t, gjson.GetBytes(upstream.lastBody, "type").Exists())
  1604. require.False(t, gjson.GetBytes(upstream.lastBody, "generate").Exists())
  1605. require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_retention").Exists())
  1606. }
  1607. func TestOpenAIWSHTTPBridgeAcceptsFirstFrameAboveLegacy16MiB(t *testing.T) {
  1608. gin.SetMode(gin.TestMode)
  1609. sseBody := strings.Join([]string{
  1610. `data: {"type":"response.created","response":{"id":"resp_large_bridge","model":"gpt-5"}}`,
  1611. "",
  1612. `data: {"type":"response.completed","response":{"id":"resp_large_bridge","model":"gpt-5","usage":{"input_tokens":9,"output_tokens":1}}}`,
  1613. "",
  1614. }, "\n")
  1615. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1616. StatusCode: http.StatusOK,
  1617. Header: http.Header{
  1618. "Content-Type": []string{"text/event-stream"},
  1619. "x-request-id": []string{"rid_large_bridge"},
  1620. },
  1621. Body: io.NopCloser(strings.NewReader(sseBody)),
  1622. }}
  1623. cfg := &config.Config{
  1624. Gateway: config.GatewayConfig{
  1625. MaxLineSize: defaultMaxLineSize,
  1626. OpenAIWS: config.GatewayOpenAIWSConfig{
  1627. Enabled: true,
  1628. APIKeyEnabled: true,
  1629. ResponsesWebsocketsV2: true,
  1630. ModeRouterV2Enabled: true,
  1631. IngressModeDefault: OpenAIWSIngressModeCtxPool,
  1632. ClientReadLimitBytes: 64 * 1024 * 1024,
  1633. HTTPBridgeEnabled: true,
  1634. HTTPBridgeThresholdBytes: 17*1024*1024 + 512,
  1635. },
  1636. },
  1637. }
  1638. svc := &OpenAIGatewayService{
  1639. cfg: cfg,
  1640. httpUpstream: upstream,
  1641. toolCorrector: NewCodexToolCorrector(),
  1642. }
  1643. account := &Account{
  1644. ID: 9,
  1645. Name: "api-key",
  1646. Platform: PlatformOpenAI,
  1647. Type: AccountTypeAPIKey,
  1648. Credentials: map[string]any{
  1649. "api_key": "sk-upstream",
  1650. "base_url": "https://env-openai.example/v1",
  1651. "model_provider": "env-openai",
  1652. },
  1653. Extra: map[string]any{
  1654. "openai_apikey_responses_websockets_v2_enabled": true,
  1655. "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough,
  1656. },
  1657. Concurrency: 1,
  1658. Status: StatusActive,
  1659. }
  1660. payload := []byte(strings.Repeat(" ", 1024) + `{"type":"response.create","generate":true,"model":"gpt-5","stream":true,"input":"` + strings.Repeat("x", 17*1024*1024) + `"}`)
  1661. require.Greater(t, len(payload), 16*1024*1024)
  1662. require.GreaterOrEqual(t, int64(len(payload)), cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes)
  1663. require.Less(t, int64(len(payload)), ResolveOpenAIWSClientReadLimitBytes(cfg))
  1664. type turnOutcome struct {
  1665. turn int
  1666. result *OpenAIForwardResult
  1667. err error
  1668. }
  1669. turnOutcomeCh := make(chan turnOutcome, 1)
  1670. hooks := &OpenAIWSIngressHooks{
  1671. AfterTurn: func(turn int, result *OpenAIForwardResult, turnErr error) {
  1672. turnOutcomeCh <- turnOutcome{turn: turn, result: result, err: turnErr}
  1673. },
  1674. }
  1675. errCh := make(chan error, 1)
  1676. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1677. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  1678. if err != nil {
  1679. errCh <- err
  1680. return
  1681. }
  1682. defer func() { _ = conn.CloseNow() }()
  1683. conn.SetReadLimit(ResolveOpenAIWSClientReadLimitBytes(cfg))
  1684. readCtx, cancelRead := context.WithTimeout(r.Context(), 10*time.Second)
  1685. msgType, firstMessage, err := conn.Read(readCtx)
  1686. cancelRead()
  1687. if err != nil {
  1688. errCh <- err
  1689. return
  1690. }
  1691. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  1692. errCh <- NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "unexpected client websocket message type", nil)
  1693. return
  1694. }
  1695. rec := httptest.NewRecorder()
  1696. ginCtx, _ := gin.CreateTestContext(rec)
  1697. req := r.Clone(r.Context())
  1698. req.Header = req.Header.Clone()
  1699. req.Header.Set("User-Agent", "codex_cli_rs/0.135.0")
  1700. ginCtx.Request = req
  1701. proxyCtx, cancelProxy := context.WithTimeout(r.Context(), 20*time.Second)
  1702. defer cancelProxy()
  1703. errCh <- svc.ProxyResponsesWebSocketFromClient(proxyCtx, ginCtx, conn, account, "sk-test", firstMessage, hooks)
  1704. }))
  1705. defer wsServer.Close()
  1706. dialCtx, cancelDial := context.WithTimeout(context.Background(), 5*time.Second)
  1707. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1708. cancelDial()
  1709. require.NoError(t, err)
  1710. defer func() { _ = clientConn.CloseNow() }()
  1711. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 20*time.Second)
  1712. err = clientConn.Write(writeCtx, coderws.MessageText, payload)
  1713. cancelWrite()
  1714. require.NoError(t, err)
  1715. var eventTypes []string
  1716. for {
  1717. readCtx, cancelRead := context.WithTimeout(context.Background(), 10*time.Second)
  1718. msgType, event, readErr := clientConn.Read(readCtx)
  1719. cancelRead()
  1720. require.NoError(t, readErr)
  1721. require.Equal(t, coderws.MessageText, msgType)
  1722. eventType := gjson.GetBytes(event, "type").String()
  1723. eventTypes = append(eventTypes, eventType)
  1724. if eventType == "response.completed" {
  1725. break
  1726. }
  1727. }
  1728. require.Contains(t, eventTypes, "response.created")
  1729. require.Contains(t, eventTypes, "response.completed")
  1730. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  1731. select {
  1732. case proxyErr := <-errCh:
  1733. require.NoError(t, proxyErr)
  1734. case <-time.After(10 * time.Second):
  1735. t.Fatal("timed out waiting for websocket bridge proxy to finish")
  1736. }
  1737. select {
  1738. case outcome := <-turnOutcomeCh:
  1739. require.NoError(t, outcome.err)
  1740. require.Equal(t, 1, outcome.turn)
  1741. require.NotNil(t, outcome.result)
  1742. require.Equal(t, 9, outcome.result.Usage.InputTokens)
  1743. require.Equal(t, 1, outcome.result.Usage.OutputTokens)
  1744. require.Equal(t, "response.completed", outcome.result.UpstreamTerminalEvent)
  1745. default:
  1746. t.Fatal("AfterTurn was not called for websocket HTTP bridge turn")
  1747. }
  1748. require.NotNil(t, upstream.lastReq)
  1749. require.Equal(t, http.MethodPost, upstream.lastReq.Method)
  1750. require.Equal(t, "https://env-openai.example/v1/responses", upstream.lastReq.URL.String())
  1751. require.Equal(t, "env-openai", account.GetCredential("model_provider"))
  1752. require.Less(t, int64(len(upstream.lastBody)), cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes)
  1753. require.Greater(t, len(upstream.lastBody), 16*1024*1024)
  1754. require.False(t, gjson.GetBytes(upstream.lastBody, "type").Exists())
  1755. require.False(t, gjson.GetBytes(upstream.lastBody, "generate").Exists())
  1756. require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
  1757. require.Equal(t, "gpt-5", gjson.GetBytes(upstream.lastBody, "model").String())
  1758. }
  1759. func TestOpenAIWSHTTPBridgeKeepsContinuationFramesOnHTTPWithoutPreviousResponseID(t *testing.T) {
  1760. gin.SetMode(gin.TestMode)
  1761. firstSSEBody := strings.Join([]string{
  1762. `data: {"type":"response.completed","response":{"id":"resp_bridge_first","model":"gpt-5.1","output":[{"type":"function_call","id":"fc_bridge_1","call_id":"call_bridge_1","name":"shell","arguments":"{}"}],"usage":{"input_tokens":9,"output_tokens":1}}}`,
  1763. "",
  1764. }, "\n")
  1765. secondSSEBody := strings.Join([]string{
  1766. `data: {"type":"response.completed","response":{"id":"resp_bridge_second","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`,
  1767. "",
  1768. }, "\n")
  1769. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  1770. {
  1771. StatusCode: http.StatusOK,
  1772. Header: http.Header{
  1773. "Content-Type": []string{"text/event-stream"},
  1774. },
  1775. Body: io.NopCloser(strings.NewReader(firstSSEBody)),
  1776. },
  1777. {
  1778. StatusCode: http.StatusOK,
  1779. Header: http.Header{
  1780. "Content-Type": []string{"text/event-stream"},
  1781. },
  1782. Body: io.NopCloser(strings.NewReader(secondSSEBody)),
  1783. },
  1784. }}
  1785. cfg := &config.Config{}
  1786. cfg.Security.URLAllowlist.Enabled = false
  1787. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1788. cfg.Gateway.OpenAIWS.Enabled = true
  1789. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1790. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1791. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1792. cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true
  1793. cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = 1
  1794. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1795. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1796. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1797. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  1798. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  1799. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
  1800. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  1801. captureConn := &openAIWSCaptureConn{}
  1802. captureDialer := &openAIWSCaptureDialer{conn: captureConn}
  1803. pool := newOpenAIWSConnPool(cfg)
  1804. pool.setClientDialerForTest(captureDialer)
  1805. svc := &OpenAIGatewayService{
  1806. cfg: cfg,
  1807. httpUpstream: upstream,
  1808. cache: &stubGatewayCache{},
  1809. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1810. toolCorrector: NewCodexToolCorrector(),
  1811. openaiWSPool: pool,
  1812. }
  1813. account := &Account{
  1814. ID: 19,
  1815. Name: "api-key-bridge-handoff",
  1816. Platform: PlatformOpenAI,
  1817. Type: AccountTypeAPIKey,
  1818. Credentials: map[string]any{"api_key": "sk-upstream"},
  1819. Extra: map[string]any{
  1820. "responses_websockets_v2_enabled": true,
  1821. },
  1822. Concurrency: 1,
  1823. Status: StatusActive,
  1824. Schedulable: true,
  1825. }
  1826. errCh := make(chan error, 1)
  1827. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1828. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  1829. if err != nil {
  1830. errCh <- err
  1831. return
  1832. }
  1833. defer func() { _ = conn.CloseNow() }()
  1834. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  1835. msgType, firstMessage, err := conn.Read(readCtx)
  1836. cancelRead()
  1837. if err != nil {
  1838. errCh <- err
  1839. return
  1840. }
  1841. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  1842. errCh <- NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "unexpected client websocket message type", nil)
  1843. return
  1844. }
  1845. rec := httptest.NewRecorder()
  1846. ginCtx, _ := gin.CreateTestContext(rec)
  1847. req := r.Clone(r.Context())
  1848. req.Header = req.Header.Clone()
  1849. req.Header.Set("User-Agent", "codex_cli_rs/0.135.0")
  1850. ginCtx.Request = req
  1851. errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  1852. }))
  1853. defer wsServer.Close()
  1854. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1855. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1856. cancelDial()
  1857. require.NoError(t, err)
  1858. defer func() { _ = clientConn.CloseNow() }()
  1859. writeMessage := func(payload string) {
  1860. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  1861. defer cancelWrite()
  1862. require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload)))
  1863. }
  1864. readMessage := func() []byte {
  1865. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  1866. defer cancelRead()
  1867. msgType, event, readErr := clientConn.Read(readCtx)
  1868. require.NoError(t, readErr)
  1869. require.Equal(t, coderws.MessageText, msgType)
  1870. return event
  1871. }
  1872. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":true,"input":"first"}`)
  1873. firstTurnEvent := readMessage()
  1874. require.Equal(t, "response.completed", gjson.GetBytes(firstTurnEvent, "type").String())
  1875. require.Equal(t, "resp_bridge_first", gjson.GetBytes(firstTurnEvent, "response.id").String())
  1876. writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"previous_response_id":"resp_bridge_first","input":[{"type":"function_call_output","call_id":"call_bridge_1","output":"ok"}]}`)
  1877. secondTurnEvent := readMessage()
  1878. require.Equal(t, "response.completed", gjson.GetBytes(secondTurnEvent, "type").String())
  1879. require.Equal(t, "resp_bridge_second", gjson.GetBytes(secondTurnEvent, "response.id").String())
  1880. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  1881. select {
  1882. case proxyErr := <-errCh:
  1883. require.NoError(t, proxyErr)
  1884. case <-time.After(5 * time.Second):
  1885. t.Fatal("timed out waiting for websocket bridge proxy to finish")
  1886. }
  1887. require.Len(t, upstream.bodies, 2, "进入 HTTP bridge 后同一客户端 WS 连接内应保持 HTTP/SSE bridge")
  1888. require.False(t, gjson.GetBytes(upstream.bodies[0], "previous_response_id").Exists())
  1889. require.False(t, gjson.GetBytes(upstream.bodies[1], "previous_response_id").Exists())
  1890. secondInput := gjson.GetBytes(upstream.bodies[1], "input").Array()
  1891. require.Len(t, secondInput, 3)
  1892. require.Equal(t, "first", secondInput[0].String())
  1893. require.Equal(t, "function_call", secondInput[1].Get("type").String())
  1894. require.Equal(t, "call_bridge_1", secondInput[1].Get("call_id").String())
  1895. require.Equal(t, "function_call_output", secondInput[2].Get("type").String())
  1896. require.Equal(t, "call_bridge_1", secondInput[2].Get("call_id").String())
  1897. require.Equal(t, 0, captureDialer.DialCount())
  1898. require.Empty(t, captureConn.writes)
  1899. }
  1900. func TestOpenAIWSHTTPBridge_IdleTimeoutClosesClientSession(t *testing.T) {
  1901. gin.SetMode(gin.TestMode)
  1902. sseBody := strings.Join([]string{
  1903. `data: {"type":"response.completed","response":{"id":"resp_bridge_idle","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`,
  1904. "",
  1905. }, "\n")
  1906. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1907. StatusCode: http.StatusOK,
  1908. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  1909. Body: io.NopCloser(strings.NewReader(sseBody)),
  1910. }}
  1911. cfg := &config.Config{}
  1912. cfg.Security.URLAllowlist.Enabled = false
  1913. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  1914. cfg.Gateway.OpenAIWS.Enabled = true
  1915. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1916. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1917. cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true
  1918. cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = 1
  1919. cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 1
  1920. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1921. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1922. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  1923. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
  1924. svc := &OpenAIGatewayService{
  1925. cfg: cfg,
  1926. httpUpstream: upstream,
  1927. cache: &stubGatewayCache{},
  1928. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  1929. toolCorrector: NewCodexToolCorrector(),
  1930. }
  1931. account := &Account{
  1932. ID: 20,
  1933. Name: "api-key-bridge-idle-timeout",
  1934. Platform: PlatformOpenAI,
  1935. Type: AccountTypeAPIKey,
  1936. Credentials: map[string]any{"api_key": "sk-upstream"},
  1937. Extra: map[string]any{"responses_websockets_v2_enabled": true},
  1938. Concurrency: 1,
  1939. Status: StatusActive,
  1940. Schedulable: true,
  1941. }
  1942. errCh := make(chan error, 1)
  1943. wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1944. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  1945. if err != nil {
  1946. errCh <- err
  1947. return
  1948. }
  1949. defer func() { _ = conn.CloseNow() }()
  1950. readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
  1951. _, firstMessage, err := conn.Read(readCtx)
  1952. cancelRead()
  1953. if err != nil {
  1954. errCh <- err
  1955. return
  1956. }
  1957. rec := httptest.NewRecorder()
  1958. ginCtx, _ := gin.CreateTestContext(rec)
  1959. ginCtx.Request = r.Clone(r.Context())
  1960. errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
  1961. }))
  1962. defer wsServer.Close()
  1963. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  1964. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
  1965. cancelDial()
  1966. require.NoError(t, err)
  1967. defer func() { _ = clientConn.CloseNow() }()
  1968. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  1969. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false,"input":"hello"}`))
  1970. cancelWrite()
  1971. require.NoError(t, err)
  1972. readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
  1973. _, event, err := clientConn.Read(readCtx)
  1974. cancelRead()
  1975. require.NoError(t, err)
  1976. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  1977. closeReadCtx, cancelCloseRead := context.WithTimeout(context.Background(), 3*time.Second)
  1978. _, _, err = clientConn.Read(closeReadCtx)
  1979. cancelCloseRead()
  1980. var clientClose coderws.CloseError
  1981. require.ErrorAs(t, err, &clientClose)
  1982. require.Equal(t, coderws.StatusNormalClosure, clientClose.Code)
  1983. require.Equal(t, "websocket idle timeout", clientClose.Reason)
  1984. select {
  1985. case proxyErr := <-errCh:
  1986. var closeErr *OpenAIWSClientCloseError
  1987. require.ErrorAs(t, proxyErr, &closeErr)
  1988. require.Equal(t, coderws.StatusNormalClosure, closeErr.StatusCode())
  1989. require.Equal(t, "websocket idle timeout", closeErr.Reason())
  1990. case <-time.After(4 * time.Second):
  1991. t.Fatal("timed out waiting for idle HTTP bridge session to close")
  1992. }
  1993. require.Len(t, upstream.bodies, 1, "an idle client must not leave a continuation request running")
  1994. }