openai_ws_v2_passthrough_adapter.go 57 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612
  1. package service
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "net/http"
  7. "net/url"
  8. "strings"
  9. "sync"
  10. "sync/atomic"
  11. "time"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
  14. openaiwsv2 "github.com/Wei-Shaw/sub2api/internal/service/openai_ws_v2"
  15. coderws "github.com/coder/websocket"
  16. "github.com/gin-gonic/gin"
  17. "github.com/tidwall/gjson"
  18. )
  19. type openAIWSClientFrameConn struct {
  20. conn *coderws.Conn
  21. controlCtx context.Context
  22. interTurnIdleTimeout time.Duration
  23. interTurnStarted chan struct{}
  24. waitingForNextTurn atomic.Bool
  25. // The relay observes upstream payloads, while clients must keep seeing the
  26. // model identifier they supplied for the current turn.
  27. restoreResponseModel func([]byte) []byte
  28. restoreToolNames func([]byte) []byte
  29. }
  30. // openAIWSPolicyEnforcingFrameConn wraps a client-side FrameConn and runs
  31. // every client→upstream frame through the OpenAI Fast Policy. It is the
  32. // passthrough-relay equivalent of the parseClientPayload integration in the
  33. // ingress session path. filter returns:
  34. // - newPayload, nil, nil: forward the (possibly mutated) payload
  35. // - _, *OpenAIFastBlockedError, nil: block — the wrapper sends an error
  36. // event via onBlock and surfaces a transport-level error so the relay
  37. // stops reading from the client.
  38. // - _, _, err: a transport error other than block.
  39. type openAIWSPolicyEnforcingFrameConn struct {
  40. inner openaiwsv2.FrameConn
  41. filter func(msgType coderws.MessageType, payload []byte) ([]byte, *OpenAIFastBlockedError, error)
  42. onBlock func(blocked *OpenAIFastBlockedError)
  43. }
  44. var _ openaiwsv2.FrameConn = (*openAIWSPolicyEnforcingFrameConn)(nil)
  45. func (c *openAIWSPolicyEnforcingFrameConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) {
  46. if c == nil || c.inner == nil {
  47. return coderws.MessageText, nil, errOpenAIWSConnClosed
  48. }
  49. msgType, payload, err := c.inner.ReadFrame(ctx)
  50. if err != nil {
  51. return msgType, payload, err
  52. }
  53. if c.filter == nil {
  54. return msgType, payload, nil
  55. }
  56. updated, blocked, filterErr := c.filter(msgType, payload)
  57. if filterErr != nil {
  58. return msgType, payload, filterErr
  59. }
  60. if blocked != nil {
  61. if c.onBlock != nil {
  62. c.onBlock(blocked)
  63. }
  64. return msgType, nil, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, blocked.Message, blocked)
  65. }
  66. return msgType, updated, nil
  67. }
  68. func (c *openAIWSPolicyEnforcingFrameConn) WriteFrame(ctx context.Context, msgType coderws.MessageType, payload []byte) error {
  69. if c == nil || c.inner == nil {
  70. return errOpenAIWSConnClosed
  71. }
  72. return c.inner.WriteFrame(ctx, msgType, payload)
  73. }
  74. func (c *openAIWSPolicyEnforcingFrameConn) Close() error {
  75. if c == nil || c.inner == nil {
  76. return nil
  77. }
  78. return c.inner.Close()
  79. }
  80. // openAIWSPassthroughPolicyModelForFrame returns the model actually forwarded
  81. // by a passthrough WS frame. Passthrough only replaces authentication, so an
  82. // account's ordinary model_mapping and Codex aliases must not rewrite it.
  83. func openAIWSPassthroughPolicyModelForFrame(account *Account, payload []byte) string {
  84. if account == nil || len(payload) == 0 {
  85. return ""
  86. }
  87. original := strings.TrimSpace(gjson.GetBytes(payload, "model").String())
  88. if original == "" {
  89. return ""
  90. }
  91. return original
  92. }
  93. // openAIWSPassthroughPolicyModelFromSessionFrame returns the upstream model
  94. // derived from a session.update frame's session.model field. Returns "" when
  95. // the frame is not a session.update event or carries no session.model. Used
  96. // by the per-frame policy filter (client→upstream direction) to keep
  97. // capturedSessionModel in sync with the session-level model the client may
  98. // rotate mid-session.
  99. //
  100. // Realtime / Responses WS lets the client change the session model after
  101. // the WS handshake via:
  102. //
  103. // {"type":"session.update","session":{"model":"gpt-5.5", ...}}
  104. //
  105. // If we only capture the model from the very first frame, a client can ship
  106. // gpt-4o on the first response.create (whitelisted as pass), then
  107. // session.update to gpt-5.5, then send response.create without "model" so
  108. // the per-frame resolver returns "" and the stale capturedSessionModel falls
  109. // back to gpt-4o — defeating the gpt-5.5 fast-policy filter.
  110. func openAIWSPassthroughPolicyModelFromSessionFrame(account *Account, payload []byte) string {
  111. if account == nil || len(payload) == 0 {
  112. return ""
  113. }
  114. frameType := strings.TrimSpace(gjson.GetBytes(payload, "type").String())
  115. if frameType != "session.update" {
  116. return ""
  117. }
  118. original := strings.TrimSpace(gjson.GetBytes(payload, "session.model").String())
  119. if original == "" {
  120. return ""
  121. }
  122. return original
  123. }
  124. type openAIWSPassthroughUsageMeta struct {
  125. serviceTier atomic.Pointer[string]
  126. reasoningEffort atomic.Pointer[string]
  127. requestedReasoningEffort atomic.Pointer[string]
  128. requestModel atomic.Pointer[string]
  129. upstreamModel atomic.Pointer[string]
  130. // 仅在 client->upstream filter goroutine 中读写;Load 侧通过上方原子指针同步。
  131. sessionRequestModel string
  132. }
  133. func newOpenAIWSPassthroughUsageMeta(initialRequestModel string, firstFrame []byte) *openAIWSPassthroughUsageMeta {
  134. meta := &openAIWSPassthroughUsageMeta{
  135. sessionRequestModel: strings.TrimSpace(initialRequestModel),
  136. }
  137. if meta.sessionRequestModel == "" {
  138. meta.sessionRequestModel = openAIWSPassthroughRequestModelForFrame(firstFrame)
  139. }
  140. return meta
  141. }
  142. func (m *openAIWSPassthroughUsageMeta) initFromFirstFrame(policyOutput []byte, mappedModel string) {
  143. if m == nil {
  144. return
  145. }
  146. m.serviceTier.Store(extractOpenAIServiceTierFromBody(policyOutput))
  147. m.reasoningEffort.Store(extractOpenAIReasoningEffortFromBody(policyOutput, mappedModel, m.sessionRequestModel))
  148. m.storeTurnModels(m.sessionRequestModel, policyOutput)
  149. }
  150. func (m *openAIWSPassthroughUsageMeta) captureRequestedReasoningEffort(originalBody []byte, modelCandidates ...string) {
  151. if m == nil {
  152. return
  153. }
  154. candidates := append([]string{m.sessionRequestModel}, modelCandidates...)
  155. m.requestedReasoningEffort.Store(CanonicalRequestedReasoningEffort(originalBody, candidates...))
  156. }
  157. func (m *openAIWSPassthroughUsageMeta) updateSessionRequestModel(payload []byte) {
  158. if m == nil {
  159. return
  160. }
  161. if model := openAIWSPassthroughRequestModelFromSessionFrame(payload); model != "" {
  162. m.sessionRequestModel = model
  163. }
  164. }
  165. func (m *openAIWSPassthroughUsageMeta) requestModelForFrame(payload []byte) string {
  166. if m == nil {
  167. return openAIWSPassthroughRequestModelForFrame(payload)
  168. }
  169. if model := openAIWSPassthroughRequestModelForFrame(payload); model != "" {
  170. return model
  171. }
  172. return m.sessionRequestModel
  173. }
  174. func (m *openAIWSPassthroughUsageMeta) updateFromResponseCreate(policyOutput []byte, mappedModel string, requestModelForFrame string) {
  175. if m == nil {
  176. return
  177. }
  178. m.serviceTier.Store(extractOpenAIServiceTierFromBody(policyOutput))
  179. m.reasoningEffort.Store(extractOpenAIReasoningEffortFromBody(policyOutput, mappedModel, requestModelForFrame))
  180. m.storeTurnModels(requestModelForFrame, policyOutput)
  181. }
  182. func (m *openAIWSPassthroughUsageMeta) storeTurnModels(requestModel string, upstreamPayload []byte) {
  183. if m == nil {
  184. return
  185. }
  186. requestModel = strings.TrimSpace(requestModel)
  187. upstreamModel := strings.TrimSpace(gjson.GetBytes(upstreamPayload, "model").String())
  188. if upstreamModel == "" {
  189. upstreamModel = requestModel
  190. }
  191. m.requestModel.Store(openAIWSTrimmedStringPtr(requestModel))
  192. m.upstreamModel.Store(openAIWSTrimmedStringPtr(upstreamModel))
  193. }
  194. func (m *openAIWSPassthroughUsageMeta) turnModels(fallback string) (string, string) {
  195. requestModel := strings.TrimSpace(fallback)
  196. upstreamModel := requestModel
  197. if m == nil {
  198. return requestModel, upstreamModel
  199. }
  200. if current := m.requestModel.Load(); current != nil && strings.TrimSpace(*current) != "" {
  201. requestModel = strings.TrimSpace(*current)
  202. }
  203. if current := m.upstreamModel.Load(); current != nil && strings.TrimSpace(*current) != "" {
  204. upstreamModel = strings.TrimSpace(*current)
  205. }
  206. return requestModel, upstreamModel
  207. }
  208. func openAIWSTrimmedStringPtr(value string) *string {
  209. value = strings.TrimSpace(value)
  210. if value == "" {
  211. return nil
  212. }
  213. return &value
  214. }
  215. func openAIWSDifferentModel(requestModel, upstreamModel string) string {
  216. upstreamModel = strings.TrimSpace(upstreamModel)
  217. if upstreamModel == "" || upstreamModel == strings.TrimSpace(requestModel) {
  218. return ""
  219. }
  220. return upstreamModel
  221. }
  222. func openAIWSPassthroughRequestModelForFrame(payload []byte) string {
  223. if len(payload) == 0 || strings.TrimSpace(gjson.GetBytes(payload, "type").String()) != "response.create" {
  224. return ""
  225. }
  226. return strings.TrimSpace(gjson.GetBytes(payload, "model").String())
  227. }
  228. func openAIWSPassthroughRequestModelFromSessionFrame(payload []byte) string {
  229. if len(payload) == 0 || strings.TrimSpace(gjson.GetBytes(payload, "type").String()) != "session.update" {
  230. return ""
  231. }
  232. return strings.TrimSpace(gjson.GetBytes(payload, "session.model").String())
  233. }
  234. const openaiWSV2PassthroughModeFields = "ws_mode=passthrough ws_router=v2"
  235. var errOpenAIWSPassthroughFirstOutputTimeout = errors.New("openai websocket passthrough first output timeout")
  236. var errOpenAIWSPassthroughActiveTurnTimeout = errors.New("openai websocket passthrough active turn read timeout")
  237. type openAIWSPassthroughDeadlinePhase uint8
  238. const (
  239. openAIWSPassthroughDeadlinePhaseFirstSemantic openAIWSPassthroughDeadlinePhase = iota + 1
  240. openAIWSPassthroughDeadlinePhaseActiveRead
  241. )
  242. type openAIWSPassthroughFirstOutputDeadline struct {
  243. timeout time.Duration
  244. startedAt time.Time
  245. requestModel string
  246. reasoningEffort string
  247. phase openAIWSPassthroughDeadlinePhase
  248. }
  249. type openAIWSPassthroughFirstOutputTimeoutError struct {
  250. deadline openAIWSPassthroughFirstOutputDeadline
  251. }
  252. func (e *openAIWSPassthroughFirstOutputTimeoutError) Error() string {
  253. return errOpenAIWSPassthroughFirstOutputTimeout.Error()
  254. }
  255. func (e *openAIWSPassthroughFirstOutputTimeoutError) Unwrap() error {
  256. return errOpenAIWSPassthroughFirstOutputTimeout
  257. }
  258. type openAIWSPassthroughActiveTurnTimeoutError struct{}
  259. func (e *openAIWSPassthroughActiveTurnTimeoutError) Error() string {
  260. return errOpenAIWSPassthroughActiveTurnTimeout.Error()
  261. }
  262. func (e *openAIWSPassthroughActiveTurnTimeoutError) Unwrap() error {
  263. return errOpenAIWSPassthroughActiveTurnTimeout
  264. }
  265. type openAIWSPassthroughFirstOutputDeadlineState struct {
  266. armed bool
  267. generation uint64
  268. deadline openAIWSPassthroughFirstOutputDeadline
  269. }
  270. type openAIWSPassthroughTurnLifecycle struct {
  271. mu sync.Mutex
  272. inFlight bool
  273. }
  274. func newOpenAIWSPassthroughTurnLifecycle(inFlight bool) *openAIWSPassthroughTurnLifecycle {
  275. return &openAIWSPassthroughTurnLifecycle{inFlight: inFlight}
  276. }
  277. func (l *openAIWSPassthroughTurnLifecycle) beginResponseCreate(onAccepted func()) bool {
  278. if l == nil {
  279. return false
  280. }
  281. l.mu.Lock()
  282. defer l.mu.Unlock()
  283. if l.inFlight {
  284. return false
  285. }
  286. l.inFlight = true
  287. if onAccepted != nil {
  288. onAccepted()
  289. }
  290. return true
  291. }
  292. func (l *openAIWSPassthroughTurnLifecycle) cancelResponseCreate() {
  293. if l == nil {
  294. return
  295. }
  296. l.mu.Lock()
  297. l.inFlight = false
  298. l.mu.Unlock()
  299. }
  300. func (l *openAIWSPassthroughTurnLifecycle) beginTerminalWrite() {
  301. if l != nil {
  302. l.mu.Lock()
  303. }
  304. }
  305. func (l *openAIWSPassthroughTurnLifecycle) finishTerminalWrite(succeeded bool, onSucceeded func()) {
  306. if l == nil {
  307. return
  308. }
  309. if succeeded {
  310. if onSucceeded != nil {
  311. onSucceeded()
  312. }
  313. l.inFlight = false
  314. }
  315. l.mu.Unlock()
  316. }
  317. type openAIWSPassthroughFirstOutputFrameConn struct {
  318. inner openaiwsv2.FrameConn
  319. resolveDeadline func(payload []byte) openAIWSPassthroughFirstOutputDeadline
  320. activeReadTimeout time.Duration
  321. mu sync.Mutex
  322. state openAIWSPassthroughFirstOutputDeadlineState
  323. deadlineChanged chan struct{}
  324. }
  325. func (c *openAIWSPassthroughFirstOutputFrameConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) {
  326. if c == nil || c.inner == nil {
  327. return coderws.MessageText, nil, errOpenAIWSConnClosed
  328. }
  329. if ctx == nil {
  330. ctx = context.Background()
  331. }
  332. type readResult struct {
  333. msgType coderws.MessageType
  334. payload []byte
  335. err error
  336. }
  337. readCtx, cancelRead := context.WithCancel(ctx)
  338. readResultCh := make(chan readResult, 1)
  339. go func() {
  340. msgType, payload, err := c.inner.ReadFrame(readCtx)
  341. readResultCh <- readResult{msgType: msgType, payload: payload, err: err}
  342. }()
  343. var timer *time.Timer
  344. var timerCh <-chan time.Time
  345. resetTimer := func() {
  346. state := c.deadlineState()
  347. if timer != nil {
  348. if !timer.Stop() {
  349. select {
  350. case <-timer.C:
  351. default:
  352. }
  353. }
  354. }
  355. if !state.armed || state.deadline.timeout <= 0 {
  356. timerCh = nil
  357. return
  358. }
  359. remaining := time.Until(state.deadline.startedAt.Add(state.deadline.timeout))
  360. if remaining < 0 {
  361. remaining = 0
  362. }
  363. if timer == nil {
  364. timer = time.NewTimer(remaining)
  365. } else {
  366. timer.Reset(remaining)
  367. }
  368. timerCh = timer.C
  369. }
  370. resetTimer()
  371. defer func() {
  372. cancelRead()
  373. if timer != nil {
  374. timer.Stop()
  375. }
  376. }()
  377. for {
  378. select {
  379. case result := <-readResultCh:
  380. if result.err == nil {
  381. c.observeUpstreamActivity(result.msgType, result.payload)
  382. }
  383. return result.msgType, result.payload, result.err
  384. case <-c.deadlineChanged:
  385. resetTimer()
  386. case <-timerCh:
  387. state := c.deadlineState()
  388. if !state.armed || state.deadline.timeout <= 0 || time.Now().Before(state.deadline.startedAt.Add(state.deadline.timeout)) {
  389. resetTimer()
  390. continue
  391. }
  392. if ctx.Err() != nil {
  393. cancelRead()
  394. <-readResultCh
  395. return coderws.MessageText, nil, ctx.Err()
  396. }
  397. cancelRead()
  398. <-readResultCh
  399. if state.deadline.phase == openAIWSPassthroughDeadlinePhaseActiveRead {
  400. return coderws.MessageText, nil, &openAIWSPassthroughActiveTurnTimeoutError{}
  401. }
  402. return coderws.MessageText, nil, &openAIWSPassthroughFirstOutputTimeoutError{deadline: state.deadline}
  403. case <-ctx.Done():
  404. cancelRead()
  405. <-readResultCh
  406. return coderws.MessageText, nil, ctx.Err()
  407. }
  408. }
  409. }
  410. func (c *openAIWSPassthroughFirstOutputFrameConn) WriteFrame(ctx context.Context, msgType coderws.MessageType, payload []byte) error {
  411. if c == nil || c.inner == nil {
  412. return errOpenAIWSConnClosed
  413. }
  414. generation := uint64(0)
  415. if msgType == coderws.MessageText && strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.create" {
  416. generation = c.armDeadline(payload)
  417. }
  418. if err := c.inner.WriteFrame(ctx, msgType, payload); err != nil {
  419. c.disarmDeadline(generation)
  420. return err
  421. }
  422. return nil
  423. }
  424. func (c *openAIWSPassthroughFirstOutputFrameConn) Close() error {
  425. if c == nil || c.inner == nil {
  426. return nil
  427. }
  428. return c.inner.Close()
  429. }
  430. func (c *openAIWSPassthroughFirstOutputFrameConn) armDeadline(payload []byte) uint64 {
  431. if c == nil || c.resolveDeadline == nil {
  432. return 0
  433. }
  434. deadline := c.resolveDeadline(payload)
  435. if deadline.timeout <= 0 {
  436. return 0
  437. }
  438. if deadline.startedAt.IsZero() {
  439. deadline.startedAt = time.Now()
  440. }
  441. deadline.phase = openAIWSPassthroughDeadlinePhaseFirstSemantic
  442. c.mu.Lock()
  443. c.state.generation++
  444. generation := c.state.generation
  445. c.state.armed = true
  446. c.state.deadline = deadline
  447. c.mu.Unlock()
  448. c.notifyDeadlineChanged()
  449. return generation
  450. }
  451. func (c *openAIWSPassthroughFirstOutputFrameConn) observeUpstreamActivity(msgType coderws.MessageType, payload []byte) {
  452. if c == nil {
  453. return
  454. }
  455. if msgType == coderws.MessageText && openAIWSPassthroughIsTerminalOutput(payload) {
  456. c.disarmDeadline(0)
  457. return
  458. }
  459. state := c.deadlineState()
  460. if state.armed && state.deadline.phase == openAIWSPassthroughDeadlinePhaseActiveRead {
  461. c.armActiveReadDeadline()
  462. return
  463. }
  464. if msgType == coderws.MessageText && openAIWSPassthroughStartsSemanticOutput(payload) {
  465. c.armActiveReadDeadline()
  466. }
  467. }
  468. func (c *openAIWSPassthroughFirstOutputFrameConn) armActiveReadDeadline() {
  469. if c == nil {
  470. return
  471. }
  472. if c.activeReadTimeout <= 0 {
  473. c.disarmDeadline(0)
  474. return
  475. }
  476. c.mu.Lock()
  477. c.state.generation++
  478. c.state.armed = true
  479. c.state.deadline = openAIWSPassthroughFirstOutputDeadline{
  480. timeout: c.activeReadTimeout,
  481. startedAt: time.Now(),
  482. phase: openAIWSPassthroughDeadlinePhaseActiveRead,
  483. }
  484. c.mu.Unlock()
  485. c.notifyDeadlineChanged()
  486. }
  487. func (c *openAIWSPassthroughFirstOutputFrameConn) disarmDeadline(generation uint64) {
  488. if c == nil {
  489. return
  490. }
  491. c.mu.Lock()
  492. if !c.state.armed || (generation != 0 && generation != c.state.generation) {
  493. c.mu.Unlock()
  494. return
  495. }
  496. c.state.armed = false
  497. c.mu.Unlock()
  498. c.notifyDeadlineChanged()
  499. }
  500. func (c *openAIWSPassthroughFirstOutputFrameConn) deadlineState() openAIWSPassthroughFirstOutputDeadlineState {
  501. if c == nil {
  502. return openAIWSPassthroughFirstOutputDeadlineState{}
  503. }
  504. c.mu.Lock()
  505. defer c.mu.Unlock()
  506. return c.state
  507. }
  508. func (c *openAIWSPassthroughFirstOutputFrameConn) notifyDeadlineChanged() {
  509. if c == nil || c.deadlineChanged == nil {
  510. return
  511. }
  512. select {
  513. case c.deadlineChanged <- struct{}{}:
  514. default:
  515. }
  516. }
  517. func openAIWSPassthroughStartsSemanticOutput(payload []byte) bool {
  518. eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String())
  519. switch eventType {
  520. case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled":
  521. return true
  522. case "", "response.created", "response.in_progress", "response.output_item.added", "response.output_item.done":
  523. return false
  524. }
  525. return strings.Contains(eventType, ".delta") ||
  526. strings.HasPrefix(eventType, "response.output_text") ||
  527. strings.HasPrefix(eventType, "response.output")
  528. }
  529. func openAIWSPassthroughIsTerminalOutput(payload []byte) bool {
  530. switch strings.TrimSpace(gjson.GetBytes(payload, "type").String()) {
  531. case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled":
  532. return true
  533. default:
  534. return false
  535. }
  536. }
  537. var _ openaiwsv2.FrameConn = (*openAIWSClientFrameConn)(nil)
  538. var _ openaiwsv2.FrameConn = (*openAIWSPassthroughFirstOutputFrameConn)(nil)
  539. func (c *openAIWSClientFrameConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) {
  540. if c == nil || c.conn == nil {
  541. return coderws.MessageText, nil, errOpenAIWSConnClosed
  542. }
  543. controlCtx := ctx
  544. if c.controlCtx != nil {
  545. controlCtx = c.controlCtx
  546. }
  547. msgType, payload, err := readOpenAIWSClientMessageWithTimeoutStart(
  548. controlCtx,
  549. c.conn,
  550. c.interTurnIdleTimeout,
  551. coderws.StatusNormalClosure,
  552. "websocket idle timeout",
  553. c.interTurnStarted,
  554. func() bool { return c.waitingForNextTurn.Load() },
  555. )
  556. return msgType, payload, err
  557. }
  558. func (c *openAIWSClientFrameConn) markTurnStarted() {
  559. if c != nil {
  560. c.waitingForNextTurn.Store(false)
  561. }
  562. }
  563. func (c *openAIWSClientFrameConn) markTurnCompleted() {
  564. if c == nil {
  565. return
  566. }
  567. c.waitingForNextTurn.Store(true)
  568. select {
  569. case c.interTurnStarted <- struct{}{}:
  570. default:
  571. }
  572. }
  573. func (c *openAIWSClientFrameConn) WriteFrame(ctx context.Context, msgType coderws.MessageType, payload []byte) error {
  574. if c == nil || c.conn == nil {
  575. return errOpenAIWSConnClosed
  576. }
  577. if ctx == nil {
  578. ctx = context.Background()
  579. }
  580. if msgType == coderws.MessageText {
  581. if normalized, changed := normalizeCompletedImageGenerationStatus(payload); changed {
  582. payload = normalized
  583. }
  584. if c.restoreResponseModel != nil {
  585. payload = c.restoreResponseModel(payload)
  586. }
  587. if c.restoreToolNames != nil {
  588. payload = c.restoreToolNames(payload)
  589. }
  590. }
  591. return c.conn.Write(ctx, msgType, payload)
  592. }
  593. func (c *openAIWSClientFrameConn) Close() error {
  594. if c == nil || c.conn == nil {
  595. return nil
  596. }
  597. _ = c.conn.Close(coderws.StatusNormalClosure, "")
  598. _ = c.conn.CloseNow()
  599. return nil
  600. }
  601. func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
  602. ctx context.Context,
  603. c *gin.Context,
  604. clientConn *coderws.Conn,
  605. account *Account,
  606. token string,
  607. firstClientMessage []byte,
  608. hooks *OpenAIWSIngressHooks,
  609. wsDecision OpenAIWSProtocolDecision,
  610. ) error {
  611. // The same request context survives scheduler failover attempts. Clear any
  612. // reverse mapping installed by the prior account before processing frames.
  613. setCodexToolNameReverse(c, nil)
  614. if s == nil {
  615. return errors.New("service is nil")
  616. }
  617. if clientConn == nil {
  618. return errors.New("client websocket is nil")
  619. }
  620. if account == nil {
  621. return errors.New("account is nil")
  622. }
  623. if err := validateOpenAIWSBearerToken(account, token); err != nil {
  624. return err
  625. }
  626. mode1OriginalFirst := append([]byte(nil), firstClientMessage...)
  627. if isOpenAIResponsesLiteWebSocketPayload(firstClientMessage) {
  628. liteFirstMessage, _, liteErr := normalizeOpenAIResponsesLitePayloadForAccount(firstClientMessage, account)
  629. if liteErr != nil {
  630. return NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, liteErr.Error(), liteErr)
  631. }
  632. firstClientMessage = liteFirstMessage
  633. }
  634. originalFirstClientMessage := firstClientMessage
  635. if next, policyErr := applyOpenAIWSReasoningEffortPolicy(firstClientMessage, hooks); policyErr != nil {
  636. return NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, policyErr.Error(), policyErr)
  637. } else {
  638. firstClientMessage = next
  639. }
  640. requestModel := strings.TrimSpace(gjson.GetBytes(firstClientMessage, "model").String())
  641. requestPreviousResponseID := strings.TrimSpace(gjson.GetBytes(firstClientMessage, "previous_response_id").String())
  642. promptCacheKey := strings.TrimSpace(gjson.GetBytes(firstClientMessage, "prompt_cache_key").String())
  643. logOpenAIWSV2Passthrough(
  644. "relay_start account_id=%d model=%s previous_response_id=%s first_message_type=%s first_message_bytes=%d",
  645. account.ID,
  646. truncateOpenAIWSLogValue(requestModel, openAIWSLogValueMaxLen),
  647. truncateOpenAIWSLogValue(requestPreviousResponseID, openAIWSIDValueMaxLen),
  648. openaiwsv2RelayMessageTypeName(coderws.MessageText),
  649. len(firstClientMessage),
  650. )
  651. // Apply OpenAI Fast Policy on the first response.create frame. Subsequent
  652. // frames are filtered via a wrapping FrameConn below so every client→
  653. // upstream frame goes through the same policy evaluator/normalize/scope as
  654. // HTTP entrypoints.
  655. //
  656. // We capture the session-level model from the first frame here so the
  657. // per-frame filter (below) can fall back to it when a follow-up frame
  658. // omits "model" — Realtime clients are allowed to send response.create
  659. // without re-stating the model, in which case the upstream uses the model
  660. // negotiated at session.update time. Without this fallback, an empty
  661. // model would miss any admin-configured model whitelist and be silently
  662. // passed through, defeating that policy on every frame after the first.
  663. initialRequestModel := ""
  664. if hooks != nil {
  665. initialRequestModel = strings.TrimSpace(hooks.InitialRequestModel)
  666. }
  667. if initialRequestModel == "" {
  668. initialRequestModel = openAIWSPassthroughRequestModelForFrame(firstClientMessage)
  669. }
  670. if hooks != nil && hooks.MapRequestModel != nil {
  671. mappedModel, mapErr := hooks.MapRequestModel(1, initialRequestModel)
  672. if mapErr != nil {
  673. return mapErr
  674. }
  675. if mappedModel = strings.TrimSpace(mappedModel); mappedModel != "" {
  676. firstClientMessage = s.ReplaceModelInBody(firstClientMessage, mappedModel)
  677. }
  678. }
  679. capturedSessionModel := openAIWSPassthroughPolicyModelForFrame(account, firstClientMessage)
  680. if capturedSessionModel != "" && capturedSessionModel != strings.TrimSpace(gjson.GetBytes(firstClientMessage, "model").String()) {
  681. firstClientMessage = s.ReplaceModelInBody(firstClientMessage, capturedSessionModel)
  682. }
  683. firstMessageResponsesLite := isOpenAIResponsesLiteWebSocketPayload(firstClientMessage)
  684. if normalized, compatibilityChanged, normalizeErr := normalizeOpenAIResponsesWebSocketCompatibilityBody(firstClientMessage, account, firstMessageResponsesLite); normalizeErr != nil {
  685. return fmt.Errorf("normalize first websocket response.create: %w", normalizeErr)
  686. } else if compatibilityChanged {
  687. firstClientMessage = normalized
  688. }
  689. if account.IsOpenAIOAuthLike() {
  690. aliasedBody, reverse, aliased, aliasErr := aliasOpenAIOAuthReservedToolNamesBody(firstClientMessage)
  691. if aliasErr != nil {
  692. return aliasErr
  693. }
  694. updateCodexToolNameReverseForWSFrame(c, firstClientMessage, reverse)
  695. if aliased {
  696. firstClientMessage = aliasedBody
  697. }
  698. }
  699. accountScopedFirst, accountScoped, scopeErr := applyCodexAccountAndProtectionIdentityRaw(c, account, firstClientMessage)
  700. if scopeErr != nil {
  701. return NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket identity metadata", scopeErr)
  702. }
  703. if accountScoped {
  704. firstClientMessage = accountScopedFirst
  705. }
  706. usageMeta := newOpenAIWSPassthroughUsageMeta(initialRequestModel, firstClientMessage)
  707. updatedFirst, blocked, policyErr := s.applyOpenAIFastPolicyToWSResponseCreate(ctx, account, capturedSessionModel, firstClientMessage)
  708. if policyErr != nil {
  709. return fmt.Errorf("apply openai fast policy on first ws frame: %w", policyErr)
  710. }
  711. if blocked != nil {
  712. MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalPolicyDenied)
  713. // coder/websocket@v1.8.14 Conn.Write is synchronous: it acquires
  714. // writeFrameMu, writes the entire frame, and Flushes the underlying
  715. // bufio writer before returning (write.go:42 → write.go:307-311).
  716. // The subsequent close handshake re-acquires the same writeFrameMu
  717. // to send the close frame, so the error event is guaranteed to
  718. // reach the kernel send buffer before any close frame is queued.
  719. // No explicit flush hop is required here.
  720. eventBytes := buildOpenAIFastPolicyBlockedWSEvent(blocked)
  721. if eventBytes != nil {
  722. writeCtx, cancelWrite := context.WithTimeout(ctx, s.openAIWSWriteTimeout())
  723. _ = clientConn.Write(writeCtx, coderws.MessageText, eventBytes)
  724. cancelWrite()
  725. }
  726. return NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, blocked.Message, blocked)
  727. }
  728. firstClientMessage = updatedFirst
  729. if account.RequestIntegrityMode() != "off" {
  730. if err := checkAccountRequestIntegrity(c, account, mode1OriginalFirst, firstClientMessage); err != nil {
  731. return NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, err.Error(), err)
  732. }
  733. }
  734. // 在 policy filter 之后再提取 service_tier / reasoning_effort 用于
  735. // usage 上报:filter 命中时 service_tier 已经从 firstClientMessage 中删除,
  736. // 最终出站 tier 应为 nil,而不是用户最初请求的 "priority"。观察到的回包
  737. // tier 单独保存在 UpstreamResponseServiceTier,由 usage 阶段统一决策。
  738. // HTTP 入口(line ~2728 extractOpenAIServiceTier(reqBody))
  739. // 与 WS ingress(openai_ws_forwarder.go:2991 取自 payload)的语义一致。
  740. //
  741. // 多轮 passthrough:OpenAI Realtime / Responses WS 协议允许客户端在
  742. // 同一连接的不同 response.create 帧上发送不同 service_tier(参考
  743. // codex-rs/core/src/client.rs build_responses_request 每次重新填值)。
  744. // 因此使用 atomic.Pointer[string] 在 filter(runClientToUpstream
  745. // goroutine)和 OnTurnComplete / final result(runUpstreamToClient
  746. // goroutine)之间同步当前 turn 的 usage metadata。
  747. usageMeta.initFromFirstFrame(firstClientMessage, capturedSessionModel)
  748. usageMeta.captureRequestedReasoningEffort(originalFirstClientMessage, capturedSessionModel)
  749. _, initialUpstreamModel := usageMeta.turnModels(initialRequestModel)
  750. SetOpsUpstreamModel(c, initialUpstreamModel)
  751. wsURL, err := s.buildOpenAIResponsesWSURL(account)
  752. if err != nil {
  753. return fmt.Errorf("build ws url: %w", err)
  754. }
  755. wsHost := "-"
  756. wsPath := "-"
  757. if parsedURL, parseErr := url.Parse(wsURL); parseErr == nil && parsedURL != nil {
  758. wsHost = normalizeOpenAIWSLogValue(parsedURL.Host)
  759. wsPath = normalizeOpenAIWSLogValue(parsedURL.Path)
  760. }
  761. logOpenAIWSV2Passthrough(
  762. "relay_dial_start account_id=%d ws_host=%s ws_path=%s proxy_enabled=%v",
  763. account.ID,
  764. wsHost,
  765. wsPath,
  766. account.ProxyID != nil && account.Proxy != nil,
  767. )
  768. isCodexCLI := false
  769. if c != nil {
  770. isCodexCLI = openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator"))
  771. }
  772. if s.cfg != nil && s.cfg.Gateway.ForceCodexCLI {
  773. isCodexCLI = true
  774. }
  775. turnState := ""
  776. turnMetadata := ""
  777. if c != nil {
  778. turnState = strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader))
  779. turnMetadata = strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader))
  780. }
  781. headers, _, buildHdrErr := s.buildOpenAIWSHeaders(
  782. ctx,
  783. c,
  784. account,
  785. token,
  786. wsDecision,
  787. isCodexCLI,
  788. turnState,
  789. turnMetadata,
  790. promptCacheKey,
  791. gjson.GetBytes(firstClientMessage, "model").String(),
  792. gjson.GetBytes(firstClientMessage, "service_tier").String(),
  793. )
  794. if buildHdrErr != nil {
  795. return fmt.Errorf("build ws headers: %w", buildHdrErr)
  796. }
  797. proxyURL := ""
  798. if account.ProxyID != nil && account.Proxy != nil {
  799. proxyURL = account.Proxy.URL()
  800. }
  801. dialer := s.getOpenAIWSPassthroughDialer()
  802. if dialer == nil {
  803. return errors.New("openai ws passthrough dialer is nil")
  804. }
  805. tlsProfile, err := resolveMode1TLSProfile(account)
  806. if err != nil {
  807. return err
  808. }
  809. if tlsProfile != nil && s.cfg != nil && !s.cfg.Gateway.TLSFingerprint.Enabled {
  810. tlsProfile = nil
  811. }
  812. agentTaskRecoveryTried := false
  813. var upstreamConn openAIWSClientConn
  814. statusCode := 0
  815. var handshakeHeaders http.Header
  816. for {
  817. headers, err = s.refreshOpenAIAgentIdentityHeaders(ctx, account, headers)
  818. if err != nil {
  819. return fmt.Errorf("refresh ws authentication headers: %w", err)
  820. }
  821. dialCtx, cancelDial := context.WithTimeout(ctx, s.openAIWSDialTimeout())
  822. upstreamConn, statusCode, handshakeHeaders, err = dialer.Dial(withOpenAIWSTLSProfile(dialCtx, account.ID, tlsProfile), wsURL, headers, proxyURL)
  823. cancelDial()
  824. if err == nil {
  825. break
  826. }
  827. var handshakeErr *openAIWSHandshakeError
  828. responseBody := []byte(nil)
  829. if errors.As(err, &handshakeErr) && handshakeErr != nil {
  830. responseBody = handshakeErr.Body
  831. }
  832. dialErr := &openAIWSDialError{StatusCode: statusCode, ResponseHeaders: cloneHeader(handshakeHeaders), ResponseBody: responseBody, Err: err}
  833. if s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried {
  834. agentTaskRecoveryTried = true
  835. if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil {
  836. return fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
  837. }
  838. continue
  839. }
  840. logOpenAIWSV2Passthrough(
  841. "relay_dial_failed account_id=%d status_code=%d err=%s",
  842. account.ID,
  843. statusCode,
  844. truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen),
  845. )
  846. s.handleOpenAIWSDialTransientFailure(ctx, account, capturedSessionModel, dialErr)
  847. if statusCode == http.StatusTooManyRequests {
  848. s.persistOpenAIWSRateLimitSignal(ctx, account, handshakeHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error()), capturedSessionModel)
  849. return s.newOpenAIWSRateLimitFailoverError(account, handshakeHeaders, nil, err.Error())
  850. }
  851. return s.mapOpenAIWSPassthroughDialError(err, statusCode, handshakeHeaders)
  852. }
  853. defer func() {
  854. _ = upstreamConn.Close()
  855. }()
  856. logOpenAIWSV2Passthrough(
  857. "relay_dial_ok account_id=%d status_code=%d upstream_request_id=%s",
  858. account.ID,
  859. statusCode,
  860. openAIWSHeaderValueForLog(handshakeHeaders, "x-request-id"),
  861. )
  862. upstreamFrameConn, ok := upstreamConn.(openaiwsv2.FrameConn)
  863. if !ok {
  864. return errors.New("openai ws passthrough upstream connection does not support frame relay")
  865. }
  866. upstreamFrameConn, trafficWrapErr := wrapAccountTrafficFrameConn(ctx, s.httpUpstream, account, upstreamFrameConn)
  867. if trafficWrapErr != nil {
  868. return trafficWrapErr
  869. }
  870. if controlled, ok := upstreamFrameConn.(*accountTrafficFrameConn); ok {
  871. defer controlled.finish(0)
  872. }
  873. relayUpstreamFrameConn := &openAIWSPassthroughFirstOutputFrameConn{
  874. inner: upstreamFrameConn,
  875. activeReadTimeout: s.openAIWSPassthroughIdleTimeout(),
  876. deadlineChanged: make(chan struct{}, 1),
  877. resolveDeadline: func(payload []byte) openAIWSPassthroughFirstOutputDeadline {
  878. reasoningEffort := ""
  879. if current := usageMeta.reasoningEffort.Load(); current != nil {
  880. reasoningEffort = *current
  881. }
  882. timeout := s.openAIFirstOutputTimeout(reasoningEffort)
  883. if timeout <= 0 {
  884. timeout = s.openAIWSPassthroughIdleTimeout()
  885. }
  886. model := openAIWSPassthroughRequestModelForFrame(payload)
  887. if model == "" {
  888. model = usageMeta.requestModelForFrame(payload)
  889. }
  890. if model == "" {
  891. model = requestModel
  892. }
  893. return openAIWSPassthroughFirstOutputDeadline{
  894. timeout: timeout,
  895. startedAt: time.Now(),
  896. requestModel: model,
  897. reasoningEffort: reasoningEffort,
  898. }
  899. },
  900. }
  901. completedTurns := atomic.Int32{}
  902. turnLifecycle := newOpenAIWSPassthroughTurnLifecycle(true)
  903. var acceptedTurnStartedAt atomic.Pointer[time.Time]
  904. clientFrameConn := &openAIWSClientFrameConn{
  905. conn: clientConn,
  906. controlCtx: ctx,
  907. interTurnIdleTimeout: s.openAIWSIngressInterTurnIdleTimeout(),
  908. interTurnStarted: make(chan struct{}, 1),
  909. restoreResponseModel: func(payload []byte) []byte {
  910. eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String())
  911. if !openAIWSEventMayContainModel(eventType) {
  912. return payload
  913. }
  914. requestModel, upstreamModel := usageMeta.turnModels("")
  915. return replaceOpenAIWSMessageModel(payload, upstreamModel, requestModel)
  916. },
  917. restoreToolNames: func(payload []byte) []byte {
  918. return restoreCodexToolNamesFromContext(c, payload)
  919. },
  920. }
  921. policyClientConn := &openAIWSPolicyEnforcingFrameConn{
  922. inner: clientFrameConn,
  923. // 注意线程安全:filter 仅在 runClientToUpstream 这一条
  924. // goroutine 中被调用(passthrough_relay.go: ReadFrame loop),
  925. // capturedSessionModel 的读写都发生在该 goroutine 内,因此无需
  926. // 加锁/原子化。
  927. filter: func(msgType coderws.MessageType, payload []byte) (out []byte, blocked *OpenAIFastBlockedError, filterErr error) {
  928. if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
  929. return payload, nil, nil
  930. }
  931. mode1OriginalFrame := append([]byte(nil), payload...)
  932. eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String())
  933. isResponseCreate := eventType == "response.create"
  934. responseCreateAt := time.Time{}
  935. acceptedTurn := false
  936. if isResponseCreate {
  937. responseCreateAt = time.Now()
  938. if !turnLifecycle.beginResponseCreate(clientFrameConn.markTurnStarted) {
  939. err := errors.New("overlapping response.create is not supported")
  940. return payload, nil, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, err.Error(), err)
  941. }
  942. defer func() {
  943. if !acceptedTurn {
  944. turnLifecycle.cancelResponseCreate()
  945. }
  946. }()
  947. }
  948. responsesLite := isResponseCreate && isOpenAIResponsesLiteWebSocketPayload(payload)
  949. if isResponseCreate {
  950. if normalized, compatibilityChanged, normalizeErr := normalizeOpenAIResponsesWebSocketCompatibilityBody(payload, account, responsesLite); normalizeErr != nil {
  951. return payload, nil, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", normalizeErr)
  952. } else if compatibilityChanged {
  953. payload = normalized
  954. }
  955. }
  956. if account.IsOpenAIOAuthLike() && (isResponseCreate || eventType == "session.update") {
  957. aliasedBody, reverse, aliased, aliasErr := aliasOpenAIOAuthReservedToolNamesBody(payload)
  958. if aliasErr != nil {
  959. return payload, nil, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, aliasErr.Error(), aliasErr)
  960. }
  961. updateCodexToolNameReverseForWSFrame(c, payload, reverse)
  962. if aliased {
  963. payload = aliasedBody
  964. }
  965. }
  966. if isResponseCreate || eventType == "session.update" {
  967. accountScopedPayload, accountScoped, scopeErr := applyCodexAccountAndProtectionIdentityRaw(c, account, payload)
  968. if scopeErr != nil {
  969. return payload, nil, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket identity metadata", scopeErr)
  970. }
  971. if accountScoped {
  972. payload = accountScopedPayload
  973. }
  974. }
  975. if isResponseCreate {
  976. if responsesLite {
  977. litePayload, _, liteErr := normalizeOpenAIResponsesLitePayloadForAccount(payload, account)
  978. if liteErr != nil {
  979. return payload, nil, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, liteErr.Error(), liteErr)
  980. }
  981. payload = litePayload
  982. }
  983. originalResponseCreate := payload
  984. if next, policyErr := applyOpenAIWSReasoningEffortPolicy(payload, hooks); policyErr != nil {
  985. return payload, nil, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, policyErr.Error(), policyErr)
  986. } else {
  987. payload = next
  988. }
  989. usageMeta.captureRequestedReasoningEffort(originalResponseCreate)
  990. }
  991. turnNo := int(completedTurns.Load()) + 1
  992. if turnNo < 2 {
  993. turnNo = 2
  994. }
  995. requestModelForThisFrame := ""
  996. if isResponseCreate {
  997. requestModelForThisFrame = usageMeta.requestModelForFrame(payload)
  998. if requestModelForThisFrame == "" {
  999. requestModelForThisFrame = capturedSessionModel
  1000. }
  1001. if hooks != nil && hooks.BeforeRequest != nil {
  1002. if err := hooks.BeforeRequest(turnNo, payload, requestModelForThisFrame); err != nil {
  1003. return payload, nil, err
  1004. }
  1005. }
  1006. if hooks != nil && hooks.BeforeTurn != nil {
  1007. if err := hooks.BeforeTurn(turnNo); err != nil {
  1008. return payload, nil, err
  1009. }
  1010. }
  1011. if hooks != nil && hooks.MapRequestModel != nil {
  1012. upstreamModel, err := hooks.MapRequestModel(turnNo, requestModelForThisFrame)
  1013. if err != nil {
  1014. return payload, nil, err
  1015. }
  1016. if upstreamModel = strings.TrimSpace(upstreamModel); upstreamModel != "" {
  1017. payload = s.ReplaceModelInBody(payload, upstreamModel)
  1018. }
  1019. }
  1020. }
  1021. // 在评估策略前先刷新 capturedSessionModel:客户端可能通过
  1022. // session.update 修改 session-level model(Realtime /
  1023. // Responses WS 协议允许),如果不刷新就会出现
  1024. // "首帧 model=gpt-4o(pass)→ session.update 改成 gpt-5.5
  1025. // → 不带 model 的 response.create fallback 到 gpt-4o" 的
  1026. // 绕过路径。这里只看 session.update 事件中的 session.model
  1027. // 字段,response.create 自己的 model 仍然由其本帧字段决定。
  1028. if updated := openAIWSPassthroughPolicyModelFromSessionFrame(account, payload); updated != "" {
  1029. capturedSessionModel = updated
  1030. }
  1031. usageMeta.updateSessionRequestModel(payload)
  1032. if requestModelForThisFrame == "" {
  1033. requestModelForThisFrame = usageMeta.requestModelForFrame(payload)
  1034. }
  1035. // Per-frame model first; if the client omits "model" on a
  1036. // follow-up frame (legal in Realtime), fall back to the
  1037. // session-level model captured from the first frame so the
  1038. // model whitelist still resolves. An empty model would miss
  1039. // any whitelist and silently fall back to pass.
  1040. model := openAIWSPassthroughPolicyModelForFrame(account, payload)
  1041. if model == "" {
  1042. model = capturedSessionModel
  1043. }
  1044. if isResponseCreate && model != "" && model != strings.TrimSpace(gjson.GetBytes(payload, "model").String()) {
  1045. payload = s.ReplaceModelInBody(payload, model)
  1046. }
  1047. out, blocked, policyErr := s.applyOpenAIFastPolicyToWSResponseCreate(ctx, account, model, payload)
  1048. if account.RequestIntegrityMode() != "off" && policyErr == nil && blocked == nil {
  1049. if err := checkAccountRequestIntegrity(c, account, mode1OriginalFrame, out); err != nil {
  1050. return payload, nil, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, err.Error(), err)
  1051. }
  1052. }
  1053. // 多轮 passthrough usage:仅在成功(non-block / non-err)
  1054. // 的 response.create 帧上更新 usageMeta,使用
  1055. // filter 处理后的 payload,与首帧 policy-after-extract 语义
  1056. // 保持一致(参见上方 extractOpenAIServiceTierFromBody 注释)。
  1057. // - 非 response.create 帧(response.cancel /
  1058. // conversation.item.create / session.update 等)不携带
  1059. // per-response metadata,不应覆盖前一轮值。
  1060. // - blocked != nil:该帧不会发送上游,usage metadata 应保持
  1061. // 上一轮值。
  1062. // - policyErr != nil:异常路径,保持上一轮值。
  1063. // - 不带 service_tier 的 response.create 会让
  1064. // extractOpenAIServiceTierFromBody 返回 nil;这里有意
  1065. // 覆盖(Store(nil)),因为 OpenAI 上游对该帧实际不传
  1066. // service_tier 时按 default 处理,billing 应如实反映。
  1067. if policyErr == nil && blocked == nil && isResponseCreate {
  1068. usageMeta.updateFromResponseCreate(out, model, requestModelForThisFrame)
  1069. _, actualModel := usageMeta.turnModels(requestModelForThisFrame)
  1070. SetOpsUpstreamModel(c, actualModel)
  1071. responseCreateAtCopy := responseCreateAt
  1072. acceptedTurnStartedAt.Store(&responseCreateAtCopy)
  1073. acceptedTurn = true
  1074. }
  1075. return out, blocked, policyErr
  1076. },
  1077. onBlock: func(blocked *OpenAIFastBlockedError) {
  1078. MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalPolicyDenied)
  1079. // See note above on Conn.Write being synchronous w.r.t. flush;
  1080. // no explicit flush is required to ensure the error event lands
  1081. // before the close frame.
  1082. eventBytes := buildOpenAIFastPolicyBlockedWSEvent(blocked)
  1083. if eventBytes == nil {
  1084. return
  1085. }
  1086. writeCtx, cancel := context.WithTimeout(ctx, s.openAIWSWriteTimeout())
  1087. _ = clientConn.Write(writeCtx, coderws.MessageText, eventBytes)
  1088. cancel()
  1089. },
  1090. }
  1091. upstreamFirstMessageSent := false
  1092. firstWriteCtx, cancelFirstWrite := context.WithTimeout(ctx, s.openAIWSWriteTimeout())
  1093. firstWriteErr := relayUpstreamFrameConn.WriteFrame(firstWriteCtx, coderws.MessageText, firstClientMessage)
  1094. cancelFirstWrite()
  1095. if firstWriteErr != nil {
  1096. return wrapOpenAIWSIngressTurnError(
  1097. "write_upstream",
  1098. fmt.Errorf("write first upstream websocket request: %w", firstWriteErr),
  1099. false,
  1100. )
  1101. }
  1102. upstreamFirstMessageSent = true
  1103. readNextClientFrame := func(readCtx context.Context, conn openaiwsv2.FrameConn) (coderws.MessageType, []byte, error) {
  1104. for {
  1105. msgType, payload, readErr := conn.ReadFrame(readCtx)
  1106. if readErr != nil {
  1107. return msgType, payload, readErr
  1108. }
  1109. if (msgType == coderws.MessageText || msgType == coderws.MessageBinary) && strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.create" {
  1110. return msgType, payload, nil
  1111. }
  1112. if writeErr := upstreamFrameConn.WriteFrame(readCtx, msgType, payload); writeErr != nil {
  1113. return msgType, payload, writeErr
  1114. }
  1115. }
  1116. }
  1117. firstTurnStartedAt := time.Time{}
  1118. if hooks != nil {
  1119. firstTurnStartedAt = hooks.InitialTurnStartedAt
  1120. }
  1121. failureAccountSideEffectsApplied := false
  1122. relayResult, relayExit := openaiwsv2.RunEntry(openaiwsv2.EntryInput{
  1123. Ctx: ctx,
  1124. ClientConn: policyClientConn,
  1125. UpstreamConn: relayUpstreamFrameConn,
  1126. FirstClientMessage: firstClientMessage,
  1127. Options: openaiwsv2.RelayOptions{
  1128. WriteTimeout: s.openAIWSWriteTimeout(),
  1129. FirstTurnStartedAt: firstTurnStartedAt,
  1130. TakeNextTurnStartedAt: func() time.Time {
  1131. startedAt := acceptedTurnStartedAt.Swap(nil)
  1132. if startedAt == nil {
  1133. return time.Time{}
  1134. }
  1135. return *startedAt
  1136. },
  1137. // Passthrough idle is enforced only after a completed turn by
  1138. // clientFrameConn. The relay-wide activity watchdog would also
  1139. // terminate a healthy active upstream turn.
  1140. IdleTimeout: 0,
  1141. FirstMessageType: coderws.MessageText,
  1142. FirstMessageSent: upstreamFirstMessageSent,
  1143. StartClientAfterFirstDownstream: true,
  1144. ReadClientFrame: readNextClientFrame,
  1145. OnUsageParseFailure: func(eventType string, usageRaw string) {
  1146. logOpenAIWSV2Passthrough(
  1147. "usage_parse_failed event_type=%s usage_raw=%s",
  1148. truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen),
  1149. truncateOpenAIWSLogValue(usageRaw, openAIWSLogValueMaxLen),
  1150. )
  1151. },
  1152. OnTurnComplete: func(turn openaiwsv2.RelayTurnResult) {
  1153. turnNo := int(completedTurns.Add(1))
  1154. if hooks != nil && hooks.TurnStarted != nil && !turn.StartedAt.IsZero() {
  1155. hooks.TurnStarted(turnNo, turn.StartedAt)
  1156. }
  1157. turnRequestModel, turnUpstreamModel := usageMeta.turnModels(turn.RequestModel)
  1158. turnResult := &OpenAIForwardResult{
  1159. RequestID: turn.RequestID,
  1160. Usage: OpenAIUsage{
  1161. InputTokens: turn.Usage.InputTokens,
  1162. OutputTokens: turn.Usage.OutputTokens,
  1163. CacheCreationInputTokens: turn.Usage.CacheCreationInputTokens,
  1164. CacheReadInputTokens: turn.Usage.CacheReadInputTokens,
  1165. ImageOutputTokens: turn.Usage.ImageOutputTokens,
  1166. },
  1167. Model: turnRequestModel,
  1168. UpstreamModel: openAIWSDifferentModel(turnRequestModel, turnUpstreamModel),
  1169. UpstreamResponseModel: turn.ResponseModel,
  1170. UpstreamResponseModelConflict: turn.ResponseModelConflict,
  1171. UpstreamResponseServiceTier: normalizeObservedOpenAIServiceTier(turn.ResponseServiceTier),
  1172. ServiceTier: usageMeta.serviceTier.Load(),
  1173. ReasoningEffort: usageMeta.reasoningEffort.Load(),
  1174. RequestedReasoningEffort: usageMeta.requestedReasoningEffort.Load(),
  1175. Stream: true,
  1176. OpenAIWSMode: true,
  1177. UpstreamTerminalEvent: normalizeOpenAIWSTerminalEvent(turn.TerminalEventType),
  1178. ResponseHeaders: cloneHeader(handshakeHeaders),
  1179. Duration: turn.Duration,
  1180. FirstTokenMs: turn.FirstTokenMs,
  1181. }
  1182. logOpenAIWSV2Passthrough(
  1183. "relay_turn_completed account_id=%d turn=%d request_id=%s terminal_event=%s turn_requested_model=%s turn_upstream_model=%s duration_ms=%d first_token_ms=%d input_tokens=%d output_tokens=%d cache_read_tokens=%d",
  1184. account.ID,
  1185. turnNo,
  1186. truncateOpenAIWSLogValue(turnResult.RequestID, openAIWSIDValueMaxLen),
  1187. truncateOpenAIWSLogValue(turn.TerminalEventType, openAIWSLogValueMaxLen),
  1188. truncateOpenAIWSLogValue(turnRequestModel, openAIWSLogValueMaxLen),
  1189. truncateOpenAIWSLogValue(turnUpstreamModel, openAIWSLogValueMaxLen),
  1190. turnResult.Duration.Milliseconds(),
  1191. openAIWSFirstTokenMsForLog(turnResult.FirstTokenMs),
  1192. turnResult.Usage.InputTokens,
  1193. turnResult.Usage.OutputTokens,
  1194. turnResult.Usage.CacheReadInputTokens,
  1195. )
  1196. if hooks != nil && hooks.AfterTurn != nil {
  1197. hooks.AfterTurn(turnNo, turnResult, nil)
  1198. }
  1199. },
  1200. BeforeClientWrite: func(msgType coderws.MessageType, payload []byte) {
  1201. if msgType == coderws.MessageText && openAIWSPassthroughIsTerminalOutput(payload) {
  1202. turnLifecycle.beginTerminalWrite()
  1203. }
  1204. },
  1205. AfterClientWrite: func(msgType coderws.MessageType, payload []byte, writeErr error) {
  1206. if msgType == coderws.MessageText && writeErr == nil {
  1207. eventType, _, _ := parseOpenAIWSEventEnvelope(payload)
  1208. markOpenAIWSClientVisibleFailure(c, eventType, payload)
  1209. }
  1210. if msgType == coderws.MessageText && openAIWSPassthroughIsTerminalOutput(payload) {
  1211. turnLifecycle.finishTerminalWrite(writeErr == nil, clientFrameConn.markTurnCompleted)
  1212. }
  1213. },
  1214. BeforeRelayCancel: func(exit openaiwsv2.RelayExit) {
  1215. if context.Cause(ctx) != nil {
  1216. return
  1217. }
  1218. status, reason, ok := openAIWSPassthroughRelayClientClose(exit, int(completedTurns.Load()))
  1219. if !ok {
  1220. return
  1221. }
  1222. // Match the handler close path and stay within the WebSocket control
  1223. // frame limit; an oversized reason makes coder/websocket skip the
  1224. // close frame, leaving the client with EOF instead of the status code.
  1225. reason = truncateString(reason, 120)
  1226. _ = clientConn.Close(status, reason)
  1227. _ = clientConn.CloseNow()
  1228. },
  1229. BeforeWriteClient: func(msgType coderws.MessageType, payload []byte, wroteDownstream bool) error {
  1230. if msgType != coderws.MessageText {
  1231. return nil
  1232. }
  1233. eventType, _, _ := parseOpenAIWSEventEnvelope(payload)
  1234. if eventType == "response.created" {
  1235. failureAccountSideEffectsApplied = false
  1236. }
  1237. if (eventType == "error" || eventType == "response.failed") && markOpenAIWSV2PassthroughCyberPolicy(c, payload) {
  1238. return nil
  1239. }
  1240. errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(payload)
  1241. isPreOutputRateLimit := eventType == "error" && !wroteDownstream && isOpenAIWSRateLimitError(errCodeRaw, errTypeRaw, errMsgRaw)
  1242. if (eventType == "error" || eventType == "response.failed") && !failureAccountSideEffectsApplied && !isPreOutputRateLimit {
  1243. failureAccountSideEffectsApplied = s.handleOpenAIWSFailureAccountSideEffects(ctx, account, capturedSessionModel, handshakeHeaders, payload)
  1244. }
  1245. if eventType != "error" {
  1246. return nil
  1247. }
  1248. if wroteDownstream || !isOpenAIWSRateLimitError(errCodeRaw, errTypeRaw, errMsgRaw) {
  1249. return nil
  1250. }
  1251. s.persistOpenAIWSRateLimitSignal(ctx, account, handshakeHeaders, payload, errCodeRaw, errTypeRaw, errMsgRaw, capturedSessionModel)
  1252. logOpenAIWSV2Passthrough(
  1253. "relay_rate_limit_failover account_id=%d err_code=%s err_type=%s err_message=%s",
  1254. account.ID,
  1255. truncateOpenAIWSLogValue(errCodeRaw, openAIWSLogValueMaxLen),
  1256. truncateOpenAIWSLogValue(errTypeRaw, openAIWSLogValueMaxLen),
  1257. truncateOpenAIWSLogValue(errMsgRaw, openAIWSLogValueMaxLen),
  1258. )
  1259. if completedTurns.Load() > 0 {
  1260. return NewOpenAIWSClientCloseError(
  1261. coderws.StatusTryAgainLater,
  1262. "upstream rate limit exceeded; please reconnect",
  1263. errors.New("later passthrough turn was rate limited before output"),
  1264. )
  1265. }
  1266. return s.newOpenAIWSRateLimitFailoverError(account, handshakeHeaders, payload, errMsgRaw)
  1267. },
  1268. OnTrace: func(event openaiwsv2.RelayTraceEvent) {
  1269. logOpenAIWSV2Passthrough(
  1270. "relay_trace account_id=%d stage=%s direction=%s msg_type=%s bytes=%d graceful=%v wrote_downstream=%v err=%s",
  1271. account.ID,
  1272. truncateOpenAIWSLogValue(event.Stage, openAIWSLogValueMaxLen),
  1273. truncateOpenAIWSLogValue(event.Direction, openAIWSLogValueMaxLen),
  1274. truncateOpenAIWSLogValue(event.MessageType, openAIWSLogValueMaxLen),
  1275. event.PayloadBytes,
  1276. event.Graceful,
  1277. event.WroteDownstream,
  1278. truncateOpenAIWSLogValue(event.Error, openAIWSLogValueMaxLen),
  1279. )
  1280. },
  1281. },
  1282. })
  1283. if cause := context.Cause(ctx); cause != nil {
  1284. if isOpenAIWSSessionPreempted(ctx) {
  1285. return errOpenAIWSSessionPreempted
  1286. }
  1287. status := coderws.StatusGoingAway
  1288. reason := "websocket request canceled"
  1289. if errors.Is(cause, ErrOpenAIWSIngressLeaseLost) {
  1290. status = coderws.StatusTryAgainLater
  1291. reason = "websocket ingress capacity lease lost; please reconnect"
  1292. }
  1293. _ = clientConn.Close(status, reason)
  1294. _ = clientConn.CloseNow()
  1295. return NewOpenAIWSClientCloseError(status, reason, cause)
  1296. }
  1297. resultRequestModel, resultUpstreamModel := usageMeta.turnModels(relayResult.RequestModel)
  1298. result := &OpenAIForwardResult{
  1299. RequestID: relayResult.RequestID,
  1300. Usage: OpenAIUsage{
  1301. InputTokens: relayResult.Usage.InputTokens,
  1302. OutputTokens: relayResult.Usage.OutputTokens,
  1303. CacheCreationInputTokens: relayResult.Usage.CacheCreationInputTokens,
  1304. CacheReadInputTokens: relayResult.Usage.CacheReadInputTokens,
  1305. ImageOutputTokens: relayResult.Usage.ImageOutputTokens,
  1306. },
  1307. Model: resultRequestModel,
  1308. UpstreamModel: openAIWSDifferentModel(resultRequestModel, resultUpstreamModel),
  1309. UpstreamResponseModel: relayResult.ResponseModel,
  1310. UpstreamResponseModelConflict: relayResult.ResponseModelConflict,
  1311. UpstreamResponseServiceTier: normalizeObservedOpenAIServiceTier(relayResult.ResponseServiceTier),
  1312. ServiceTier: usageMeta.serviceTier.Load(),
  1313. ReasoningEffort: usageMeta.reasoningEffort.Load(),
  1314. RequestedReasoningEffort: usageMeta.requestedReasoningEffort.Load(),
  1315. Stream: true,
  1316. OpenAIWSMode: true,
  1317. UpstreamTerminalEvent: normalizeOpenAIWSTerminalEvent(relayResult.TerminalEventType),
  1318. ResponseHeaders: cloneHeader(handshakeHeaders),
  1319. Duration: relayResult.Duration,
  1320. FirstTokenMs: relayResult.FirstTokenMs,
  1321. }
  1322. turnCount := int(completedTurns.Load())
  1323. if relayExit == nil {
  1324. logOpenAIWSV2Passthrough(
  1325. "relay_completed account_id=%d request_id=%s terminal_event=%s duration_ms=%d c2u_frames=%d u2c_frames=%d dropped_frames=%d turns=%d",
  1326. account.ID,
  1327. truncateOpenAIWSLogValue(result.RequestID, openAIWSIDValueMaxLen),
  1328. truncateOpenAIWSLogValue(relayResult.TerminalEventType, openAIWSLogValueMaxLen),
  1329. result.Duration.Milliseconds(),
  1330. relayResult.ClientToUpstreamFrames,
  1331. relayResult.UpstreamToClientFrames,
  1332. relayResult.DroppedDownstreamFrames,
  1333. turnCount,
  1334. )
  1335. // 正常路径按 terminal 事件逐 turn 已回调;仅在零 turn 场景兜底回调一次。
  1336. if turnCount == 0 && hooks != nil && hooks.AfterTurn != nil {
  1337. if hooks.TurnStarted != nil {
  1338. hooks.TurnStarted(1, time.Now().Add(-result.Duration))
  1339. }
  1340. hooks.AfterTurn(1, result, nil)
  1341. }
  1342. return nil
  1343. }
  1344. logOpenAIWSV2Passthrough(
  1345. "relay_failed account_id=%d stage=%s wrote_downstream=%v err=%s duration_ms=%d c2u_frames=%d u2c_frames=%d dropped_frames=%d turns=%d",
  1346. account.ID,
  1347. truncateOpenAIWSLogValue(relayExit.Stage, openAIWSLogValueMaxLen),
  1348. relayExit.WroteDownstream,
  1349. truncateOpenAIWSLogValue(relayErrorText(relayExit.Err), openAIWSLogValueMaxLen),
  1350. result.Duration.Milliseconds(),
  1351. relayResult.ClientToUpstreamFrames,
  1352. relayResult.UpstreamToClientFrames,
  1353. relayResult.DroppedDownstreamFrames,
  1354. turnCount,
  1355. )
  1356. relayErr := relayExit.Err
  1357. var firstOutputTimeoutErr *openAIWSPassthroughFirstOutputTimeoutError
  1358. if errors.As(relayErr, &firstOutputTimeoutErr) {
  1359. deadline := firstOutputTimeoutErr.deadline
  1360. // The relay ran over the WebSocket transport, so a missing managed
  1361. // proxy is an unknown route (http.DefaultClient), not a direct one.
  1362. wsProxyID, wsProxyName := opsUpstreamWSProxyAttribution(account)
  1363. failoverErr := s.newOpenAIFirstOutputTimeoutError(
  1364. ctx,
  1365. c,
  1366. account,
  1367. wsProxyID,
  1368. wsProxyName,
  1369. deadline.startedAt,
  1370. deadline.requestModel,
  1371. deadline.reasoningEffort,
  1372. deadline.timeout,
  1373. "websocket_first_semantic_output",
  1374. handshakeHeaders,
  1375. )
  1376. if turnCount == 0 && !relayExit.WroteDownstream {
  1377. relayErr = failoverErr
  1378. } else {
  1379. // The handler only retains the initial response.create across
  1380. // account attempts. Replaying it after a later-turn timeout would
  1381. // duplicate the first turn, so later turns end the client session.
  1382. relayErr = NewOpenAIWSClientCloseError(
  1383. coderws.StatusGoingAway,
  1384. "upstream produced no semantic output; please reconnect",
  1385. firstOutputTimeoutErr,
  1386. )
  1387. }
  1388. }
  1389. var activeTurnTimeoutErr *openAIWSPassthroughActiveTurnTimeoutError
  1390. if errors.As(relayErr, &activeTurnTimeoutErr) {
  1391. relayErr = NewOpenAIWSClientCloseError(
  1392. coderws.StatusGoingAway,
  1393. "upstream websocket read timeout; please reconnect",
  1394. activeTurnTimeoutErr,
  1395. )
  1396. }
  1397. if relayExit.Stage == "idle_timeout" {
  1398. relayErr = NewOpenAIWSClientCloseError(
  1399. coderws.StatusPolicyViolation,
  1400. "client websocket idle timeout",
  1401. relayErr,
  1402. )
  1403. }
  1404. turnErr := wrapOpenAIWSIngressTurnError(
  1405. relayExit.Stage,
  1406. relayErr,
  1407. relayExit.WroteDownstream,
  1408. )
  1409. if hooks != nil && hooks.AfterTurn != nil {
  1410. if hooks.TurnStarted != nil {
  1411. hooks.TurnStarted(turnCount+1, time.Now().Add(-result.Duration))
  1412. }
  1413. hooks.AfterTurn(turnCount+1, nil, turnErr)
  1414. }
  1415. return turnErr
  1416. }
  1417. func openAIWSPassthroughRelayClientClose(exit openaiwsv2.RelayExit, completedTurns int) (coderws.StatusCode, string, bool) {
  1418. var closeErr *OpenAIWSClientCloseError
  1419. if errors.As(exit.Err, &closeErr) {
  1420. return closeErr.StatusCode(), closeErr.Reason(), true
  1421. }
  1422. var activeTurnTimeoutErr *openAIWSPassthroughActiveTurnTimeoutError
  1423. if errors.As(exit.Err, &activeTurnTimeoutErr) {
  1424. return coderws.StatusGoingAway, "upstream websocket read timeout; please reconnect", true
  1425. }
  1426. var firstOutputTimeoutErr *openAIWSPassthroughFirstOutputTimeoutError
  1427. if errors.As(exit.Err, &firstOutputTimeoutErr) {
  1428. if completedTurns > 0 || exit.WroteDownstream {
  1429. return coderws.StatusGoingAway, "upstream produced no semantic output; please reconnect", true
  1430. }
  1431. return 0, "", false
  1432. }
  1433. if !exit.Graceful && exit.Stage == "read_upstream" {
  1434. return coderws.StatusInternalError, "upstream websocket proxy failed", true
  1435. }
  1436. return 0, "", false
  1437. }
  1438. func markOpenAIWSV2PassthroughCyberPolicy(c *gin.Context, payload []byte) bool {
  1439. usage := OpenAIUsage{}
  1440. parseOpenAIWSResponseUsageFromCompletedEvent(payload, &usage)
  1441. return markOpenAICyberPolicyEvent(c, payload, http.StatusOK, &usage)
  1442. }
  1443. func (s *OpenAIGatewayService) mapOpenAIWSPassthroughDialError(
  1444. err error,
  1445. statusCode int,
  1446. handshakeHeaders http.Header,
  1447. ) error {
  1448. if err == nil {
  1449. return nil
  1450. }
  1451. wrappedErr := err
  1452. var dialErr *openAIWSDialError
  1453. if !errors.As(err, &dialErr) {
  1454. var handshakeErr *openAIWSHandshakeError
  1455. var responseBody []byte
  1456. if errors.As(err, &handshakeErr) && handshakeErr != nil {
  1457. responseBody = append([]byte(nil), handshakeErr.Body...)
  1458. }
  1459. wrappedErr = &openAIWSDialError{
  1460. StatusCode: statusCode,
  1461. ResponseHeaders: cloneHeader(handshakeHeaders),
  1462. ResponseBody: responseBody,
  1463. Err: err,
  1464. }
  1465. }
  1466. if errors.Is(err, context.Canceled) {
  1467. return err
  1468. }
  1469. if errors.Is(err, context.DeadlineExceeded) {
  1470. return NewOpenAIWSClientCloseError(
  1471. coderws.StatusTryAgainLater,
  1472. "upstream websocket connect timeout",
  1473. wrappedErr,
  1474. )
  1475. }
  1476. if statusCode == http.StatusTooManyRequests {
  1477. return NewOpenAIWSClientCloseError(
  1478. coderws.StatusTryAgainLater,
  1479. "upstream websocket is busy, please retry later",
  1480. wrappedErr,
  1481. )
  1482. }
  1483. if statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden {
  1484. return NewOpenAIWSClientCloseError(
  1485. coderws.StatusPolicyViolation,
  1486. "upstream websocket authentication failed",
  1487. wrappedErr,
  1488. )
  1489. }
  1490. if statusCode >= http.StatusBadRequest && statusCode < http.StatusInternalServerError {
  1491. return NewOpenAIWSClientCloseError(
  1492. coderws.StatusPolicyViolation,
  1493. "upstream websocket handshake rejected",
  1494. wrappedErr,
  1495. )
  1496. }
  1497. return fmt.Errorf("openai ws passthrough dial: %w", wrappedErr)
  1498. }
  1499. func openaiwsv2RelayMessageTypeName(msgType coderws.MessageType) string {
  1500. switch msgType {
  1501. case coderws.MessageText:
  1502. return "text"
  1503. case coderws.MessageBinary:
  1504. return "binary"
  1505. default:
  1506. return fmt.Sprintf("unknown(%d)", msgType)
  1507. }
  1508. }
  1509. func relayErrorText(err error) string {
  1510. if err == nil {
  1511. return ""
  1512. }
  1513. return err.Error()
  1514. }
  1515. func openAIWSFirstTokenMsForLog(firstTokenMs *int) int {
  1516. if firstTokenMs == nil {
  1517. return -1
  1518. }
  1519. return *firstTokenMs
  1520. }
  1521. func logOpenAIWSV2Passthrough(format string, args ...any) {
  1522. logger.LegacyPrintf(
  1523. "service.openai_ws_v2",
  1524. "[OpenAI WS v2 passthrough] %s "+format,
  1525. append([]any{openaiWSV2PassthroughModeFields}, args...)...,
  1526. )
  1527. }