openai_ws_v2_passthrough_lifecycle_test.go 37 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904
  1. package service
  2. import (
  3. "context"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "net/http"
  8. "net/http/httptest"
  9. "strings"
  10. "sync"
  11. "sync/atomic"
  12. "testing"
  13. "time"
  14. "unicode/utf8"
  15. "github.com/Wei-Shaw/sub2api/internal/config"
  16. coderws "github.com/coder/websocket"
  17. "github.com/gin-gonic/gin"
  18. "github.com/stretchr/testify/require"
  19. "github.com/tidwall/gjson"
  20. )
  21. type stagedPassthroughFrame struct {
  22. messageType coderws.MessageType
  23. payload []byte
  24. err error
  25. }
  26. type stagedPassthroughConn struct {
  27. frames chan stagedPassthroughFrame
  28. writes chan []byte
  29. closed chan struct{}
  30. closeOnce sync.Once
  31. }
  32. func newStagedPassthroughConn() *stagedPassthroughConn {
  33. return &stagedPassthroughConn{
  34. frames: make(chan stagedPassthroughFrame, 4),
  35. writes: make(chan []byte, 4),
  36. closed: make(chan struct{}),
  37. }
  38. }
  39. func (c *stagedPassthroughConn) Send(payload string) {
  40. c.frames <- stagedPassthroughFrame{messageType: coderws.MessageText, payload: []byte(payload)}
  41. }
  42. func (c *stagedPassthroughConn) Fail(err error) {
  43. c.frames <- stagedPassthroughFrame{err: err}
  44. }
  45. func (c *stagedPassthroughConn) WriteJSON(context.Context, any) error { return nil }
  46. func (c *stagedPassthroughConn) ReadMessage(ctx context.Context) ([]byte, error) {
  47. _, payload, err := c.ReadFrame(ctx)
  48. return payload, err
  49. }
  50. func (c *stagedPassthroughConn) Ping(context.Context) error { return nil }
  51. func (c *stagedPassthroughConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) {
  52. if ctx == nil {
  53. ctx = context.Background()
  54. }
  55. select {
  56. case <-ctx.Done():
  57. return coderws.MessageText, nil, ctx.Err()
  58. case <-c.closed:
  59. return coderws.MessageText, nil, errOpenAIWSConnClosed
  60. case frame := <-c.frames:
  61. return frame.messageType, append([]byte(nil), frame.payload...), frame.err
  62. }
  63. }
  64. func (c *stagedPassthroughConn) WriteFrame(ctx context.Context, _ coderws.MessageType, payload []byte) error {
  65. if ctx == nil {
  66. ctx = context.Background()
  67. }
  68. select {
  69. case <-ctx.Done():
  70. return ctx.Err()
  71. case <-c.closed:
  72. return errOpenAIWSConnClosed
  73. default:
  74. }
  75. var parsed any
  76. if err := json.Unmarshal(payload, &parsed); err != nil {
  77. return err
  78. }
  79. select {
  80. case c.writes <- append([]byte(nil), payload...):
  81. case <-ctx.Done():
  82. return ctx.Err()
  83. case <-c.closed:
  84. return errOpenAIWSConnClosed
  85. }
  86. return nil
  87. }
  88. func (c *stagedPassthroughConn) Close() error {
  89. c.closeOnce.Do(func() { close(c.closed) })
  90. return nil
  91. }
  92. type stagedPassthroughDialer struct {
  93. conn openAIWSClientConn
  94. }
  95. func (d *stagedPassthroughDialer) Dial(context.Context, string, http.Header, string) (openAIWSClientConn, int, http.Header, error) {
  96. return d.conn, http.StatusSwitchingProtocols, http.Header{}, nil
  97. }
  98. func newPassthroughLifecycleService(cfg *config.Config, upstream *stagedPassthroughConn) *OpenAIGatewayService {
  99. return &OpenAIGatewayService{
  100. cfg: cfg,
  101. httpUpstream: &httpUpstreamRecorder{},
  102. cache: &stubGatewayCache{},
  103. openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
  104. toolCorrector: NewCodexToolCorrector(),
  105. openaiWSPassthroughDialer: &stagedPassthroughDialer{conn: upstream},
  106. }
  107. }
  108. func passthroughLifecycleConfig() *config.Config {
  109. cfg := &config.Config{}
  110. cfg.Security.URLAllowlist.Enabled = false
  111. cfg.Security.URLAllowlist.AllowInsecureHTTP = true
  112. cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 1
  113. cfg.Gateway.OpenAIWS.Enabled = true
  114. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  115. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  116. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  117. cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool
  118. cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 1
  119. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
  120. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 1
  121. cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
  122. return cfg
  123. }
  124. func passthroughLifecycleAccount() *Account {
  125. return &Account{
  126. ID: 901,
  127. Name: "passthrough-lifecycle",
  128. Platform: PlatformOpenAI,
  129. Type: AccountTypeAPIKey,
  130. Status: StatusActive,
  131. Schedulable: true,
  132. Concurrency: 1,
  133. Credentials: map[string]any{"api_key": "sk-test"},
  134. Extra: map[string]any{
  135. "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough,
  136. },
  137. }
  138. }
  139. func startPassthroughLifecycleServer(
  140. t *testing.T,
  141. controlCtx context.Context,
  142. svc *OpenAIGatewayService,
  143. account *Account,
  144. ) (*httptest.Server, <-chan error) {
  145. return startPassthroughLifecycleServerWithHooks(t, controlCtx, svc, account, nil)
  146. }
  147. func startPassthroughLifecycleServerWithHooks(
  148. t *testing.T,
  149. controlCtx context.Context,
  150. svc *OpenAIGatewayService,
  151. account *Account,
  152. hooksFactory func(*gin.Context) *OpenAIWSIngressHooks,
  153. ) (*httptest.Server, <-chan error) {
  154. t.Helper()
  155. serverErr := make(chan error, 1)
  156. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  157. conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
  158. if err != nil {
  159. serverErr <- err
  160. return
  161. }
  162. defer func() { _ = conn.CloseNow() }()
  163. msgType, firstMessage, err := ReadOpenAIWSClientMessage(
  164. controlCtx,
  165. conn,
  166. 3*time.Second,
  167. coderws.StatusPolicyViolation,
  168. "missing first response.create message",
  169. )
  170. if err != nil {
  171. serverErr <- err
  172. return
  173. }
  174. if msgType != coderws.MessageText {
  175. serverErr <- errors.New("first message was not text")
  176. return
  177. }
  178. recorder := httptest.NewRecorder()
  179. ginCtx, _ := gin.CreateTestContext(recorder)
  180. req := r.Clone(controlCtx)
  181. req.Header = req.Header.Clone()
  182. ginCtx.Request = req
  183. var hooks *OpenAIWSIngressHooks
  184. if hooksFactory != nil {
  185. hooks = hooksFactory(ginCtx)
  186. }
  187. serverErr <- svc.ProxyResponsesWebSocketFromClient(controlCtx, ginCtx, conn, account, "sk-test", firstMessage, hooks)
  188. }))
  189. return server, serverErr
  190. }
  191. func TestPassthroughLifecycle_LaterTurnPreOutputRateLimitRequestsReconnect(t *testing.T) {
  192. gin.SetMode(gin.TestMode)
  193. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  194. defer cancelControl(context.Canceled)
  195. upstream := newStagedPassthroughConn()
  196. account := passthroughLifecycleAccount()
  197. repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{*account}}}
  198. svc := newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream)
  199. svc.accountRepo = repo
  200. svc.rateLimitService = &RateLimitService{accountRepo: repo}
  201. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, svc, account)
  202. defer server.Close()
  203. clientConn := dialPassthroughLifecycleClient(t, server)
  204. defer func() { _ = clientConn.CloseNow() }()
  205. firstRequest := requirePassthroughUpstreamWrite(t, upstream, time.Second)
  206. require.Equal(t, "response.create", gjson.GetBytes(firstRequest, "type").String())
  207. upstream.Send(`{"type":"response.completed","response":{"id":"resp_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`)
  208. completed, err := readPassthroughLifecycleFrame(t, clientConn, time.Second)
  209. require.NoError(t, err)
  210. require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String())
  211. writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second)
  212. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`))
  213. cancelWrite()
  214. require.NoError(t, err)
  215. secondRequest := requirePassthroughUpstreamWrite(t, upstream, time.Second)
  216. require.Equal(t, "response.create", gjson.GetBytes(secondRequest, "type").String())
  217. resetAt := time.Now().Add(90 * time.Minute).Unix()
  218. upstream.Send(fmt.Sprintf(`{"type":"error","error":{"code":"rate_limit_exceeded","type":"usage_limit_reached","message":"The usage limit has been reached","resets_at":%d}}`, resetAt))
  219. _, err = readPassthroughLifecycleFrame(t, clientConn, time.Second)
  220. var websocketCloseErr coderws.CloseError
  221. require.ErrorAs(t, err, &websocketCloseErr)
  222. require.Equal(t, coderws.StatusTryAgainLater, websocketCloseErr.Code)
  223. require.Equal(t, "upstream rate limit exceeded; please reconnect", websocketCloseErr.Reason)
  224. require.Len(t, repo.rateLimitCalls, 1)
  225. require.WithinDuration(t, time.Unix(resetAt, 0), repo.rateLimitCalls[0], 2*time.Second)
  226. select {
  227. case err := <-serverErr:
  228. var closeErr *OpenAIWSClientCloseError
  229. require.ErrorAs(t, err, &closeErr)
  230. require.Equal(t, coderws.StatusTryAgainLater, closeErr.StatusCode())
  231. case <-time.After(time.Second):
  232. t.Fatal("later-turn rate limit did not terminate passthrough")
  233. }
  234. select {
  235. case replay := <-upstream.writes:
  236. t.Fatalf("later-turn reconnect must not replay the retained first request: %s", replay)
  237. default:
  238. }
  239. }
  240. func TestPassthroughLifecycle_CyberTerminalEventsMarkBeforeAfterTurn(t *testing.T) {
  241. gin.SetMode(gin.TestMode)
  242. tests := []struct {
  243. name string
  244. events []string
  245. wantBody string
  246. wantMessage string
  247. wantInput int
  248. wantOutput int
  249. }{
  250. {
  251. name: "error",
  252. events: []string{
  253. `{"type":"error","error":{"code":"cyber_policy","message":"blocked by error event"},"usage":{"input_tokens":5,"output_tokens":1}}`,
  254. `{"type":"response.failed","response":{"id":"resp_error","error":{"code":"cyber_policy","message":"blocked by paired failed event"},"usage":{"input_tokens":9,"output_tokens":2}}}`,
  255. },
  256. wantBody: `"type":"error"`,
  257. wantMessage: "blocked by error event",
  258. wantInput: 5,
  259. wantOutput: 1,
  260. },
  261. {
  262. name: "response_failed",
  263. events: []string{
  264. `{"type":"response.failed","response":{"id":"resp_failed","error":{"code":"cyber_policy","message":"blocked by failed event"},"usage":{"input_tokens":9,"output_tokens":2}}}`,
  265. },
  266. wantBody: `"type":"response.failed"`,
  267. wantMessage: "blocked by failed event",
  268. wantInput: 9,
  269. wantOutput: 2,
  270. },
  271. }
  272. for _, tt := range tests {
  273. t.Run(tt.name, func(t *testing.T) {
  274. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  275. defer cancelControl(context.Canceled)
  276. upstream := newStagedPassthroughConn()
  277. for _, event := range tt.events {
  278. upstream.Send(event)
  279. }
  280. markSeen := make(chan CyberPolicyMark, 1)
  281. afterTurnCalls := atomic.Int32{}
  282. server, serverErr := startPassthroughLifecycleServerWithHooks(
  283. t,
  284. controlCtx,
  285. newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream),
  286. passthroughLifecycleAccount(),
  287. func(c *gin.Context) *OpenAIWSIngressHooks {
  288. return &OpenAIWSIngressHooks{AfterTurn: func(_ int, _ *OpenAIForwardResult, _ error) {
  289. afterTurnCalls.Add(1)
  290. if mark := GetOpsCyberPolicy(c); mark != nil {
  291. select {
  292. case markSeen <- *mark:
  293. default:
  294. }
  295. }
  296. }}
  297. },
  298. )
  299. defer server.Close()
  300. clientConn := dialPassthroughLifecycleClient(t, server)
  301. defer func() { _ = clientConn.CloseNow() }()
  302. for range tt.events {
  303. _, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  304. require.NoError(t, err)
  305. }
  306. select {
  307. case mark := <-markSeen:
  308. require.Equal(t, "cyber_policy", mark.Code)
  309. require.Equal(t, tt.wantMessage, mark.Message)
  310. require.Contains(t, mark.Body, tt.wantBody)
  311. require.Equal(t, http.StatusOK, mark.UpstreamStatus)
  312. require.Equal(t, tt.wantInput, mark.UpstreamInTok)
  313. require.Equal(t, tt.wantOutput, mark.UpstreamOutTok)
  314. case <-time.After(3 * time.Second):
  315. t.Fatal("cyber mark was not visible to AfterTurn")
  316. }
  317. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  318. select {
  319. case <-serverErr:
  320. case <-time.After(3 * time.Second):
  321. t.Fatal("cyber passthrough test did not exit")
  322. }
  323. require.Equal(t, int32(1), afterTurnCalls.Load(), "error/response.failed pair must complete and record once")
  324. })
  325. }
  326. }
  327. func TestPassthroughLifecycle_NonCyberFailureKeepsAccountSideEffects(t *testing.T) {
  328. gin.SetMode(gin.TestMode)
  329. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  330. defer cancelControl(context.Canceled)
  331. upstream := newStagedPassthroughConn()
  332. upstream.Send(`{"type":"response.failed","response":{"id":"resp_non_cyber","error":{"type":"authentication_error","code":"invalid_api_key","status_code":401,"message":"credential rejected"},"usage":{"input_tokens":3,"output_tokens":1}}}`)
  333. repo := &openAIStream403AccountRepo{}
  334. svc := newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream)
  335. svc.rateLimitService = NewRateLimitService(repo, nil, svc.cfg, nil, nil)
  336. account := passthroughLifecycleAccount()
  337. markSeen := make(chan *CyberPolicyMark, 1)
  338. server, serverErr := startPassthroughLifecycleServerWithHooks(
  339. t,
  340. controlCtx,
  341. svc,
  342. account,
  343. func(c *gin.Context) *OpenAIWSIngressHooks {
  344. return &OpenAIWSIngressHooks{AfterTurn: func(_ int, _ *OpenAIForwardResult, _ error) {
  345. markSeen <- GetOpsCyberPolicy(c)
  346. }}
  347. },
  348. )
  349. defer server.Close()
  350. clientConn := dialPassthroughLifecycleClient(t, server)
  351. defer func() { _ = clientConn.CloseNow() }()
  352. event, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  353. require.NoError(t, err)
  354. require.Equal(t, "response.failed", gjson.GetBytes(event, "type").String())
  355. select {
  356. case mark := <-markSeen:
  357. require.Nil(t, mark)
  358. case <-time.After(3 * time.Second):
  359. t.Fatal("non-cyber terminal event did not complete its turn")
  360. }
  361. require.Equal(t, 1, repo.setErrorCalls, "non-cyber credential failure must retain account failure side effects")
  362. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  363. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  364. select {
  365. case <-serverErr:
  366. case <-time.After(3 * time.Second):
  367. t.Fatal("non-cyber passthrough test did not exit")
  368. }
  369. }
  370. func TestPassthroughLifecycle_CyberSkipsFailureAccountSideEffects(t *testing.T) {
  371. gin.SetMode(gin.TestMode)
  372. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  373. defer cancelControl(context.Canceled)
  374. upstream := newStagedPassthroughConn()
  375. upstream.Send(`{"type":"response.failed","response":{"id":"resp_cyber_auth","error":{"type":"authentication_error","code":"cyber_policy","status_code":401,"message":"request blocked"}}}`)
  376. repo := &openAIStream403AccountRepo{}
  377. svc := newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream)
  378. svc.rateLimitService = NewRateLimitService(repo, nil, svc.cfg, nil, nil)
  379. account := passthroughLifecycleAccount()
  380. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, svc, account)
  381. defer server.Close()
  382. clientConn := dialPassthroughLifecycleClient(t, server)
  383. defer func() { _ = clientConn.CloseNow() }()
  384. event, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  385. require.NoError(t, err)
  386. require.Equal(t, "response.failed", gjson.GetBytes(event, "type").String())
  387. require.Zero(t, repo.setErrorCalls, "cyber_policy is request-scoped and must not cool down the account")
  388. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  389. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  390. select {
  391. case <-serverErr:
  392. case <-time.After(3 * time.Second):
  393. t.Fatal("cyber side-effect test did not exit")
  394. }
  395. }
  396. func TestPassthroughLifecycle_CloseReasonTruncationPreservesUTF8(t *testing.T) {
  397. gin.SetMode(gin.TestMode)
  398. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  399. defer cancelControl(context.Canceled)
  400. upstream := newStagedPassthroughConn()
  401. originalReason := strings.Repeat("a", 119) + "界"
  402. upstream.Fail(NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, originalReason, errors.New("policy rejected")))
  403. server, serverErr := startPassthroughLifecycleServer(
  404. t,
  405. controlCtx,
  406. newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream),
  407. passthroughLifecycleAccount(),
  408. )
  409. defer server.Close()
  410. clientConn := dialPassthroughLifecycleClient(t, server)
  411. defer func() { _ = clientConn.CloseNow() }()
  412. _, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  413. var closeErr coderws.CloseError
  414. require.ErrorAs(t, err, &closeErr)
  415. require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
  416. require.True(t, utf8.ValidString(closeErr.Reason))
  417. require.LessOrEqual(t, len(closeErr.Reason), 120)
  418. require.Equal(t, strings.Repeat("a", 119), closeErr.Reason)
  419. select {
  420. case <-serverErr:
  421. case <-time.After(3 * time.Second):
  422. t.Fatal("passthrough close reason test did not exit")
  423. }
  424. }
  425. func dialPassthroughLifecycleClient(t *testing.T, server *httptest.Server) *coderws.Conn {
  426. t.Helper()
  427. return dialPassthroughLifecycleClientWithPayload(t, server, `{"type":"response.create","model":"gpt-5.1","stream":false}`)
  428. }
  429. func dialPassthroughLifecycleClientWithPayload(t *testing.T, server *httptest.Server, payload string) *coderws.Conn {
  430. t.Helper()
  431. dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
  432. clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(server.URL, "http"), nil)
  433. cancelDial()
  434. require.NoError(t, err)
  435. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  436. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(payload))
  437. cancelWrite()
  438. require.NoError(t, err)
  439. return clientConn
  440. }
  441. func readPassthroughLifecycleFrame(t *testing.T, clientConn *coderws.Conn, timeout time.Duration) ([]byte, error) {
  442. t.Helper()
  443. readCtx, cancelRead := context.WithTimeout(context.Background(), timeout)
  444. _, payload, err := clientConn.Read(readCtx)
  445. cancelRead()
  446. return payload, err
  447. }
  448. func requirePassthroughUpstreamWrite(t *testing.T, upstream *stagedPassthroughConn, timeout time.Duration) []byte {
  449. t.Helper()
  450. select {
  451. case payload := <-upstream.writes:
  452. return payload
  453. case <-time.After(timeout):
  454. t.Fatal("passthrough request was not forwarded upstream")
  455. return nil
  456. }
  457. }
  458. func TestPassthroughLifecycle_ResponsesLiteFirstFramePinsParallelToolCalls(t *testing.T) {
  459. gin.SetMode(gin.TestMode)
  460. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  461. defer cancelControl(context.Canceled)
  462. upstream := newStagedPassthroughConn()
  463. upstream.Send(`{"type":"response.completed","response":{"id":"resp_lite","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`)
  464. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount())
  465. defer server.Close()
  466. clientConn := dialPassthroughLifecycleClientWithPayload(t, server, `{
  467. "type":"response.create","model":"gpt-5.1","stream":false,
  468. "parallel_tool_calls":true,
  469. "client_metadata":{"ws_request_header_x_openai_internal_codex_responses_lite":"true"}
  470. }`)
  471. defer func() { _ = clientConn.CloseNow() }()
  472. upstreamBody := requirePassthroughUpstreamWrite(t, upstream, 3*time.Second)
  473. require.Equal(t, gjson.False, gjson.GetBytes(upstreamBody, "parallel_tool_calls").Type, string(upstreamBody))
  474. event, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  475. require.NoError(t, err)
  476. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  477. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  478. select {
  479. case <-serverErr:
  480. case <-time.After(3 * time.Second):
  481. t.Fatal("Lite 首帧测试等待 passthrough 退出超时")
  482. }
  483. }
  484. func TestOpenAIWSPassthroughTurnLifecycle_SerializesTerminalCommitAndNextTurn(t *testing.T) {
  485. clientFrameConn := &openAIWSClientFrameConn{interTurnStarted: make(chan struct{}, 1)}
  486. clientFrameConn.markTurnCompleted()
  487. lifecycle := newOpenAIWSPassthroughTurnLifecycle(true)
  488. lifecycle.beginTerminalWrite()
  489. admitted := make(chan bool, 1)
  490. go func() {
  491. admitted <- lifecycle.beginResponseCreate(clientFrameConn.markTurnStarted)
  492. }()
  493. select {
  494. case <-admitted:
  495. t.Fatal("next response.create was admitted before terminal commit completed")
  496. case <-time.After(50 * time.Millisecond):
  497. }
  498. lifecycle.finishTerminalWrite(true, clientFrameConn.markTurnCompleted)
  499. select {
  500. case ok := <-admitted:
  501. require.True(t, ok)
  502. case <-time.After(time.Second):
  503. t.Fatal("next response.create remained blocked after terminal commit")
  504. }
  505. require.False(t, clientFrameConn.waitingForNextTurn.Load(), "accepted next turn must win over terminal idle state")
  506. lifecycle = newOpenAIWSPassthroughTurnLifecycle(true)
  507. lifecycle.beginTerminalWrite()
  508. admitted = make(chan bool, 1)
  509. go func() {
  510. admitted <- lifecycle.beginResponseCreate(nil)
  511. }()
  512. lifecycle.finishTerminalWrite(false, func() {
  513. t.Error("failed terminal write must not commit idle state")
  514. })
  515. require.False(t, <-admitted, "failed terminal write must keep the current turn in flight")
  516. }
  517. func TestPassthroughLifecycle_LeaseLossSendsRetryClose(t *testing.T) {
  518. gin.SetMode(gin.TestMode)
  519. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  520. upstream := newStagedPassthroughConn()
  521. upstream.Send(`{"type":"response.created","response":{"id":"resp_lease","model":"gpt-5.1"}}`)
  522. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount())
  523. defer server.Close()
  524. clientConn := dialPassthroughLifecycleClient(t, server)
  525. defer func() { _ = clientConn.CloseNow() }()
  526. event, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  527. require.NoError(t, err)
  528. require.Equal(t, "response.created", gjson.GetBytes(event, "type").String())
  529. cancelControl(ErrOpenAIWSIngressLeaseLost)
  530. _, err = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  531. var closeErr coderws.CloseError
  532. require.ErrorAs(t, err, &closeErr)
  533. require.Equal(t, coderws.StatusTryAgainLater, closeErr.Code)
  534. require.Equal(t, "websocket ingress capacity lease lost; please reconnect", closeErr.Reason)
  535. select {
  536. case <-serverErr:
  537. case <-time.After(3 * time.Second):
  538. t.Fatal("passthrough lease-loss reader did not exit")
  539. }
  540. }
  541. func TestPassthroughLifecycle_CompletedTurnStartsInterTurnIdle(t *testing.T) {
  542. gin.SetMode(gin.TestMode)
  543. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  544. defer cancelControl(context.Canceled)
  545. upstream := newStagedPassthroughConn()
  546. upstream.Send(`{"type":"response.completed","response":{"id":"resp_idle","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`)
  547. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount())
  548. defer server.Close()
  549. clientConn := dialPassthroughLifecycleClient(t, server)
  550. defer func() { _ = clientConn.CloseNow() }()
  551. event, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  552. require.NoError(t, err)
  553. require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
  554. _, err = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  555. var closeErr coderws.CloseError
  556. require.ErrorAs(t, err, &closeErr)
  557. require.Equal(t, coderws.StatusNormalClosure, closeErr.Code)
  558. require.Equal(t, "websocket idle timeout", closeErr.Reason)
  559. select {
  560. case <-serverErr:
  561. case <-time.After(3 * time.Second):
  562. t.Fatal("passthrough idle reader did not exit")
  563. }
  564. }
  565. func TestPassthroughLifecycle_ActiveTurnInactivityUsesReadTimeout(t *testing.T) {
  566. gin.SetMode(gin.TestMode)
  567. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  568. defer cancelControl(context.Canceled)
  569. upstream := newStagedPassthroughConn()
  570. upstream.Send(`{"type":"response.output_text.delta","response_id":"resp_active","delta":"hello"}`)
  571. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount())
  572. defer server.Close()
  573. clientConn := dialPassthroughLifecycleClient(t, server)
  574. defer func() { _ = clientConn.CloseNow() }()
  575. delta, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  576. require.NoError(t, err)
  577. require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String())
  578. _, err = readPassthroughLifecycleFrame(t, clientConn, 2500*time.Millisecond)
  579. var websocketCloseErr coderws.CloseError
  580. require.ErrorAs(t, err, &websocketCloseErr)
  581. require.Equal(t, coderws.StatusGoingAway, websocketCloseErr.Code)
  582. require.Equal(t, "upstream websocket read timeout; please reconnect", websocketCloseErr.Reason)
  583. select {
  584. case err := <-serverErr:
  585. var closeErr *OpenAIWSClientCloseError
  586. require.ErrorAs(t, err, &closeErr)
  587. require.Equal(t, coderws.StatusGoingAway, closeErr.StatusCode())
  588. require.Equal(t, "upstream websocket read timeout; please reconnect", closeErr.Reason())
  589. case <-time.After(2500 * time.Millisecond):
  590. t.Fatal("passthrough active turn remained unbounded after upstream activity stopped")
  591. }
  592. }
  593. func TestPassthroughLifecycle_PreambleAllowsPromptClientCancel(t *testing.T) {
  594. gin.SetMode(gin.TestMode)
  595. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  596. defer cancelControl(context.Canceled)
  597. cfg := passthroughLifecycleConfig()
  598. cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 3
  599. upstream := newStagedPassthroughConn()
  600. upstream.Send(`{"type":"response.created","response":{"id":"resp_cancel","model":"gpt-5.1"}}`)
  601. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(cfg, upstream), passthroughLifecycleAccount())
  602. defer server.Close()
  603. clientConn := dialPassthroughLifecycleClient(t, server)
  604. defer func() { _ = clientConn.CloseNow() }()
  605. require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, time.Second), "type").String())
  606. created, err := readPassthroughLifecycleFrame(t, clientConn, time.Second)
  607. require.NoError(t, err)
  608. require.Equal(t, "response.created", gjson.GetBytes(created, "type").String())
  609. writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second)
  610. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.cancel","response_id":"resp_cancel"}`))
  611. cancelWrite()
  612. require.NoError(t, err)
  613. cancelFrame := requirePassthroughUpstreamWrite(t, upstream, 500*time.Millisecond)
  614. require.Equal(t, "response.cancel", gjson.GetBytes(cancelFrame, "type").String())
  615. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  616. select {
  617. case <-serverErr:
  618. case <-time.After(3 * time.Second):
  619. t.Fatal("passthrough cancel test did not exit")
  620. }
  621. }
  622. func TestPassthroughLifecycle_RejectsOverlappingResponseCreate(t *testing.T) {
  623. gin.SetMode(gin.TestMode)
  624. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  625. defer cancelControl(context.Canceled)
  626. cfg := passthroughLifecycleConfig()
  627. cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 3
  628. upstream := newStagedPassthroughConn()
  629. upstream.Send(`{"type":"response.created","response":{"id":"resp_overlap_first","model":"gpt-5.1"}}`)
  630. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(cfg, upstream), passthroughLifecycleAccount())
  631. defer server.Close()
  632. clientConn := dialPassthroughLifecycleClient(t, server)
  633. defer func() { _ = clientConn.CloseNow() }()
  634. require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, time.Second), "type").String())
  635. created, err := readPassthroughLifecycleFrame(t, clientConn, time.Second)
  636. require.NoError(t, err)
  637. require.Equal(t, "response.created", gjson.GetBytes(created, "type").String())
  638. writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second)
  639. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1"}`))
  640. cancelWrite()
  641. require.NoError(t, err)
  642. _, err = readPassthroughLifecycleFrame(t, clientConn, time.Second)
  643. var websocketCloseErr coderws.CloseError
  644. require.ErrorAs(t, err, &websocketCloseErr)
  645. require.Equal(t, coderws.StatusPolicyViolation, websocketCloseErr.Code)
  646. require.Equal(t, "overlapping response.create is not supported", websocketCloseErr.Reason)
  647. select {
  648. case err := <-serverErr:
  649. var closeErr *OpenAIWSClientCloseError
  650. require.ErrorAs(t, err, &closeErr)
  651. require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode())
  652. require.Equal(t, "overlapping response.create is not supported", closeErr.Reason())
  653. case <-time.After(3 * time.Second):
  654. t.Fatal("overlapping response.create did not terminate passthrough")
  655. }
  656. }
  657. func TestPassthroughLifecycle_ActiveTurnActivityRefreshesReadTimeout(t *testing.T) {
  658. gin.SetMode(gin.TestMode)
  659. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  660. defer cancelControl(context.Canceled)
  661. upstream := newStagedPassthroughConn()
  662. upstream.Send(`{"type":"response.output_text.delta","response_id":"resp_active_refresh","delta":"one"}`)
  663. go func() {
  664. for _, event := range []string{
  665. `{"type":"response.output_text.delta","response_id":"resp_active_refresh","delta":"two"}`,
  666. `{"type":"response.output_text.delta","response_id":"resp_active_refresh","delta":"three"}`,
  667. `{"type":"response.completed","response":{"id":"resp_active_refresh","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":3}}}`,
  668. } {
  669. timer := time.NewTimer(600 * time.Millisecond)
  670. <-timer.C
  671. timer.Stop()
  672. upstream.Send(event)
  673. }
  674. }()
  675. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount())
  676. defer server.Close()
  677. clientConn := dialPassthroughLifecycleClient(t, server)
  678. defer func() { _ = clientConn.CloseNow() }()
  679. for _, wantType := range []string{
  680. "response.output_text.delta",
  681. "response.output_text.delta",
  682. "response.output_text.delta",
  683. "response.completed",
  684. } {
  685. frame, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  686. require.NoError(t, err)
  687. require.Equal(t, wantType, gjson.GetBytes(frame, "type").String())
  688. }
  689. require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
  690. select {
  691. case <-serverErr:
  692. case <-time.After(3 * time.Second):
  693. t.Fatal("passthrough active-turn refresh test did not exit")
  694. }
  695. }
  696. func TestPassthroughLifecycle_TerminalSwitchesToInterTurnIdleTimeout(t *testing.T) {
  697. gin.SetMode(gin.TestMode)
  698. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  699. defer cancelControl(context.Canceled)
  700. cfg := passthroughLifecycleConfig()
  701. cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 1
  702. cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 2
  703. upstream := newStagedPassthroughConn()
  704. upstream.Send(`{"type":"response.completed","response":{"id":"resp_idle_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`)
  705. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(cfg, upstream), passthroughLifecycleAccount())
  706. defer server.Close()
  707. clientConn := dialPassthroughLifecycleClient(t, server)
  708. defer func() { _ = clientConn.CloseNow() }()
  709. require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, 3*time.Second), "type").String())
  710. completed, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  711. require.NoError(t, err)
  712. require.Equal(t, "resp_idle_first", gjson.GetBytes(completed, "response.id").String())
  713. time.Sleep(1300 * time.Millisecond)
  714. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  715. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","previous_response_id":"resp_idle_first"}`))
  716. cancelWrite()
  717. require.NoError(t, err)
  718. require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, 3*time.Second), "type").String())
  719. upstream.Send(`{"type":"response.completed","response":{"id":"resp_idle_second","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`)
  720. completed, err = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  721. require.NoError(t, err)
  722. require.Equal(t, "resp_idle_second", gjson.GetBytes(completed, "response.id").String())
  723. _, err = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  724. var websocketCloseErr coderws.CloseError
  725. require.ErrorAs(t, err, &websocketCloseErr)
  726. require.Equal(t, coderws.StatusNormalClosure, websocketCloseErr.Code)
  727. require.Equal(t, "websocket idle timeout", websocketCloseErr.Reason)
  728. select {
  729. case err := <-serverErr:
  730. var closeErr *OpenAIWSClientCloseError
  731. require.ErrorAs(t, err, &closeErr)
  732. require.Equal(t, coderws.StatusNormalClosure, closeErr.StatusCode())
  733. require.Equal(t, "websocket idle timeout", closeErr.Reason())
  734. case <-time.After(3 * time.Second):
  735. t.Fatal("passthrough terminal turn did not use inter-turn idle timeout")
  736. }
  737. }
  738. func TestPassthroughLifecycle_FirstOutputTimeoutRemainsBounded(t *testing.T) {
  739. gin.SetMode(gin.TestMode)
  740. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  741. defer cancelControl(context.Canceled)
  742. upstream := newStagedPassthroughConn()
  743. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount())
  744. defer server.Close()
  745. clientConn := dialPassthroughLifecycleClient(t, server)
  746. defer func() { _ = clientConn.CloseNow() }()
  747. select {
  748. case err := <-serverErr:
  749. var failoverErr *UpstreamFailoverError
  750. require.ErrorAs(t, err, &failoverErr)
  751. require.Equal(t, http.StatusGatewayTimeout, failoverErr.StatusCode)
  752. require.Contains(t, string(failoverErr.ResponseBody), "first_output_timeout")
  753. case <-time.After(2500 * time.Millisecond):
  754. t.Fatal("passthrough first output was left unbounded")
  755. }
  756. }
  757. func TestPassthroughLifecycle_ResponseCreatedTimeoutClosesWithoutFailover(t *testing.T) {
  758. gin.SetMode(gin.TestMode)
  759. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  760. defer cancelControl(context.Canceled)
  761. upstream := newStagedPassthroughConn()
  762. upstream.Send(`{"type":"response.created","response":{"id":"resp_preamble","model":"gpt-5.1"}}`)
  763. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount())
  764. defer server.Close()
  765. clientConn := dialPassthroughLifecycleClient(t, server)
  766. defer func() { _ = clientConn.CloseNow() }()
  767. created, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  768. require.NoError(t, err)
  769. require.Equal(t, "response.created", gjson.GetBytes(created, "type").String())
  770. _, err = readPassthroughLifecycleFrame(t, clientConn, 2500*time.Millisecond)
  771. var websocketCloseErr coderws.CloseError
  772. require.ErrorAs(t, err, &websocketCloseErr)
  773. require.Equal(t, coderws.StatusGoingAway, websocketCloseErr.Code)
  774. require.Equal(t, "upstream produced no semantic output; please reconnect", websocketCloseErr.Reason)
  775. select {
  776. case err := <-serverErr:
  777. var failoverErr *UpstreamFailoverError
  778. require.NotErrorAs(t, err, &failoverErr)
  779. var closeErr *OpenAIWSClientCloseError
  780. require.ErrorAs(t, err, &closeErr)
  781. require.Equal(t, coderws.StatusGoingAway, closeErr.StatusCode())
  782. require.Equal(t, "upstream produced no semantic output; please reconnect", closeErr.Reason())
  783. case <-time.After(2500 * time.Millisecond):
  784. t.Fatal("response.created timeout did not close the passthrough connection")
  785. }
  786. }
  787. func TestPassthroughLifecycle_SecondTurnTimeoutIsNotFailoverSafe(t *testing.T) {
  788. gin.SetMode(gin.TestMode)
  789. controlCtx, cancelControl := context.WithCancelCause(context.Background())
  790. defer cancelControl(context.Canceled)
  791. upstream := newStagedPassthroughConn()
  792. upstream.Send(`{"type":"response.completed","response":{"id":"resp_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`)
  793. server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount())
  794. defer server.Close()
  795. clientConn := dialPassthroughLifecycleClient(t, server)
  796. defer func() { _ = clientConn.CloseNow() }()
  797. completed, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  798. require.NoError(t, err)
  799. require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String())
  800. writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
  801. err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","previous_response_id":"resp_first"}`))
  802. cancelWrite()
  803. require.NoError(t, err)
  804. upstream.Send(`{"type":"response.created","response":{"id":"resp_second","model":"gpt-5.1"}}`)
  805. created, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
  806. require.NoError(t, err)
  807. require.Equal(t, "response.created", gjson.GetBytes(created, "type").String())
  808. _, err = readPassthroughLifecycleFrame(t, clientConn, 2500*time.Millisecond)
  809. var websocketCloseErr coderws.CloseError
  810. require.ErrorAs(t, err, &websocketCloseErr)
  811. require.Equal(t, coderws.StatusGoingAway, websocketCloseErr.Code)
  812. require.Equal(t, "upstream produced no semantic output; please reconnect", websocketCloseErr.Reason)
  813. select {
  814. case err := <-serverErr:
  815. var failoverErr *UpstreamFailoverError
  816. require.NotErrorAs(t, err, &failoverErr, "handler must not replay the initial request on another account for a later-turn timeout")
  817. var closeErr *OpenAIWSClientCloseError
  818. require.ErrorAs(t, err, &closeErr)
  819. require.Equal(t, coderws.StatusGoingAway, closeErr.StatusCode())
  820. case <-time.After(2500 * time.Millisecond):
  821. t.Fatal("second turn first semantic output was left unbounded")
  822. }
  823. }