openai_ws_http_bridge.go 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011
  1. package service
  2. import (
  3. "bufio"
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "net/http"
  10. "slices"
  11. "strings"
  12. "time"
  13. "github.com/Wei-Shaw/sub2api/internal/config"
  14. "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
  15. "github.com/gin-gonic/gin"
  16. "github.com/tidwall/gjson"
  17. )
  18. const (
  19. openAIWSClientReadLimitBytesDefault int64 = 64 * 1024 * 1024
  20. openAIWSHTTPBridgeThresholdBytesDefault int64 = 15 * 1024 * 1024
  21. openAIWSHTTPBridgeErrorBodyLimitBytes = 64 * 1024
  22. )
  23. const openAIWSHTTPBridgeToolStateContextKey = "openai_ws_http_bridge_tool_state"
  24. type openAIWSHTTPBridgeToolState struct {
  25. ClientMapping apicompat.ResponsesClientToolMapping
  26. LoweredTools json.RawMessage
  27. }
  28. func openAIWSHTTPBridgeToolStateFromContext(c *gin.Context) (openAIWSHTTPBridgeToolState, bool) {
  29. if c == nil {
  30. return openAIWSHTTPBridgeToolState{}, false
  31. }
  32. value, ok := c.Get(openAIWSHTTPBridgeToolStateContextKey)
  33. state, typed := value.(openAIWSHTTPBridgeToolState)
  34. return state, ok && typed
  35. }
  36. func setOpenAIWSHTTPBridgeToolState(c *gin.Context, state openAIWSHTTPBridgeToolState) {
  37. if c == nil {
  38. return
  39. }
  40. state.LoweredTools = append(json.RawMessage(nil), state.LoweredTools...)
  41. c.Set(openAIWSHTTPBridgeToolStateContextKey, state)
  42. }
  43. func decodeOpenAIWSHTTPBridgeLoweredTools(raw json.RawMessage) []any {
  44. if len(raw) == 0 {
  45. return nil
  46. }
  47. var tools []any
  48. if err := json.Unmarshal(raw, &tools); err != nil {
  49. return nil
  50. }
  51. return tools
  52. }
  53. func openAIWSHTTPBridgeRawField(body []byte, name string) (json.RawMessage, bool) {
  54. var fields map[string]json.RawMessage
  55. if err := json.Unmarshal(body, &fields); err != nil {
  56. return nil, false
  57. }
  58. raw, present := fields[name]
  59. return append(json.RawMessage(nil), raw...), present
  60. }
  61. func openAIWSHTTPBridgeToolUpstreamName(account *Account) string {
  62. if account != nil && account.Platform == PlatformGrok {
  63. return "Grok WS HTTP bridge"
  64. }
  65. return "OpenAI WS HTTP bridge"
  66. }
  67. // ResolveOpenAIWSClientFirstMessageTimeout returns the effective client ingress deadline.
  68. func ResolveOpenAIWSClientFirstMessageTimeout(cfg *config.Config) time.Duration {
  69. seconds := config.DefaultOpenAIWSClientFirstMessageTimeoutSeconds
  70. if cfg != nil && cfg.Gateway.OpenAIWS.ClientFirstMessageTimeoutSeconds > 0 {
  71. seconds = cfg.Gateway.OpenAIWS.ClientFirstMessageTimeoutSeconds
  72. }
  73. return time.Duration(seconds) * time.Second
  74. }
  75. func ResolveOpenAIWSClientReadLimitBytes(cfg *config.Config) int64 {
  76. if cfg == nil || cfg.Gateway.OpenAIWS.ClientReadLimitBytes <= 0 {
  77. return openAIWSClientReadLimitBytesDefault
  78. }
  79. return cfg.Gateway.OpenAIWS.ClientReadLimitBytes
  80. }
  81. func (s *OpenAIGatewayService) openAIWSHTTPBridgeEnabled() bool {
  82. return s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.HTTPBridgeEnabled
  83. }
  84. func (s *OpenAIGatewayService) openAIWSHTTPBridgeThresholdBytes() int64 {
  85. if s == nil || s.cfg == nil || s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes <= 0 {
  86. return openAIWSHTTPBridgeThresholdBytesDefault
  87. }
  88. return s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes
  89. }
  90. func (s *OpenAIGatewayService) shouldBridgeOpenAIWSHTTP(account *Account, payloadBytes int, previousResponseID string) bool {
  91. if account != nil && account.Platform == PlatformGrok {
  92. return true
  93. }
  94. if !s.openAIWSHTTPBridgeEnabled() {
  95. return false
  96. }
  97. if strings.TrimSpace(previousResponseID) != "" {
  98. return false
  99. }
  100. threshold := s.openAIWSHTTPBridgeThresholdBytes()
  101. return threshold > 0 && int64(payloadBytes) >= threshold
  102. }
  103. func (s *OpenAIGatewayService) shouldBridgeOpenAIWSPassthroughFirstMessage(account *Account, payload []byte) bool {
  104. if account != nil && account.Platform == PlatformGrok {
  105. return true
  106. }
  107. if !s.openAIWSHTTPBridgeEnabled() || int64(len(payload)) < s.openAIWSHTTPBridgeThresholdBytes() {
  108. return false
  109. }
  110. if !json.Valid(payload) {
  111. return false
  112. }
  113. i := skipOpenAIWSJSONSpace(payload, 0)
  114. if i >= len(payload) || payload[i] != '{' {
  115. return false
  116. }
  117. i++
  118. eventType := "response.create"
  119. previousResponseID := ""
  120. typeSeen, previousResponseIDSeen := false, false
  121. for {
  122. i = skipOpenAIWSJSONSpace(payload, i)
  123. if payload[i] == '}' {
  124. break
  125. }
  126. keyStart := i
  127. keyEnd := scanOpenAIWSJSONString(payload, keyStart)
  128. i = skipOpenAIWSJSONSpace(payload, keyEnd)
  129. i++ // json.Valid guarantees the colon.
  130. i = skipOpenAIWSJSONSpace(payload, i)
  131. valueStart := i
  132. i = skipOpenAIWSJSONValue(payload, i)
  133. key := ""
  134. // A critical key is at most 20 decoded bytes. The generous encoded bound
  135. // covers escaped spellings without allocating attacker-sized key strings.
  136. if keyEnd-keyStart <= 128 {
  137. _ = json.Unmarshal(payload[keyStart:keyEnd], &key)
  138. }
  139. switch key {
  140. case "type":
  141. if typeSeen {
  142. return false
  143. }
  144. typeSeen = true
  145. var value *string
  146. if err := json.Unmarshal(payload[valueStart:i], &value); err != nil {
  147. return false
  148. }
  149. if value == nil || strings.TrimSpace(*value) == "" {
  150. eventType = "response.create"
  151. } else {
  152. eventType = strings.TrimSpace(*value)
  153. }
  154. case "previous_response_id":
  155. if previousResponseIDSeen {
  156. return false
  157. }
  158. previousResponseIDSeen = true
  159. var value *string
  160. if err := json.Unmarshal(payload[valueStart:i], &value); err != nil {
  161. return false
  162. }
  163. if value != nil {
  164. previousResponseID = strings.TrimSpace(*value)
  165. }
  166. }
  167. i = skipOpenAIWSJSONSpace(payload, i)
  168. if payload[i] == ',' {
  169. i++
  170. }
  171. }
  172. return eventType == "response.create" && previousResponseID == ""
  173. }
  174. func skipOpenAIWSJSONSpace(payload []byte, i int) int {
  175. for i < len(payload) {
  176. switch payload[i] {
  177. case ' ', '\t', '\r', '\n':
  178. i++
  179. default:
  180. return i
  181. }
  182. }
  183. return i
  184. }
  185. func scanOpenAIWSJSONString(payload []byte, i int) int {
  186. for i++; i < len(payload); i++ {
  187. switch payload[i] {
  188. case '\\':
  189. i++
  190. case '"':
  191. return i + 1
  192. }
  193. }
  194. return len(payload)
  195. }
  196. func skipOpenAIWSJSONValue(payload []byte, i int) int {
  197. if payload[i] == '"' {
  198. return scanOpenAIWSJSONString(payload, i)
  199. }
  200. if payload[i] != '{' && payload[i] != '[' {
  201. for i < len(payload) && payload[i] != ',' && payload[i] != '}' {
  202. i++
  203. }
  204. return i
  205. }
  206. depth := 0
  207. for ; i < len(payload); i++ {
  208. switch payload[i] {
  209. case '"':
  210. i = scanOpenAIWSJSONString(payload, i) - 1
  211. case '{', '[':
  212. depth++
  213. case '}', ']':
  214. depth--
  215. if depth == 0 {
  216. return i + 1
  217. }
  218. }
  219. }
  220. return len(payload)
  221. }
  222. func prepareOpenAIWSHTTPBridgeBody(account *Account, payload []byte) ([]byte, error) {
  223. var body map[string]any
  224. if err := decodeOpenAIJSONUseNumber(payload, &body); err != nil {
  225. return nil, err
  226. }
  227. if body == nil {
  228. return nil, errors.New("response.create payload must be a JSON object")
  229. }
  230. delete(body, "type")
  231. delete(body, "generate")
  232. delete(body, "previous_response_id")
  233. deleteOpenAIResponsesNoneReasoningEffortFromObject(account, body)
  234. body["stream"] = true
  235. return json.Marshal(body)
  236. }
  237. type openAIWSToolCallReplayCollector struct {
  238. items []json.RawMessage
  239. seen map[string]struct{}
  240. allItems []json.RawMessage
  241. allSeen map[string]struct{}
  242. }
  243. func (c *openAIWSToolCallReplayCollector) AddEvent(eventType string, message []byte) {
  244. switch strings.TrimSpace(eventType) {
  245. case "response.output_item.done":
  246. item := gjson.GetBytes(message, "item")
  247. c.addAllItem(item)
  248. c.addItem(item)
  249. case "response.completed", "response.done":
  250. output := gjson.GetBytes(message, "response.output")
  251. if !output.IsArray() {
  252. return
  253. }
  254. for _, item := range output.Array() {
  255. c.addAllItem(item)
  256. c.addItem(item)
  257. }
  258. }
  259. }
  260. // Items/AllItems 返回浅拷贝头数组;正文由 collector 独立分配且此后不可变,
  261. // 调用方按 replay 所有权不变式共享持有。
  262. func (c *openAIWSToolCallReplayCollector) Items() []json.RawMessage {
  263. return slices.Clone(c.items)
  264. }
  265. func (c *openAIWSToolCallReplayCollector) AllItems() []json.RawMessage {
  266. return slices.Clone(c.allItems)
  267. }
  268. func (c *openAIWSToolCallReplayCollector) addAllItem(item gjson.Result) {
  269. if !item.Exists() || item.Type != gjson.JSON {
  270. return
  271. }
  272. raw := strings.TrimSpace(item.Raw)
  273. if raw == "" || !strings.HasPrefix(raw, "{") || strings.TrimSpace(item.Get("type").String()) == "" {
  274. return
  275. }
  276. key := strings.TrimSpace(item.Get("id").String())
  277. if key == "" {
  278. key = strings.TrimSpace(item.Get("call_id").String())
  279. }
  280. if key == "" {
  281. key = raw
  282. }
  283. if c.allSeen == nil {
  284. c.allSeen = make(map[string]struct{})
  285. }
  286. if _, ok := c.allSeen[key]; ok {
  287. return
  288. }
  289. c.allSeen[key] = struct{}{}
  290. c.allItems = append(c.allItems, json.RawMessage(raw))
  291. }
  292. func (c *openAIWSToolCallReplayCollector) addItem(item gjson.Result) {
  293. if !item.Exists() || item.Type != gjson.JSON {
  294. return
  295. }
  296. raw := strings.TrimSpace(item.Raw)
  297. if raw == "" || !strings.HasPrefix(raw, "{") {
  298. return
  299. }
  300. if !isCodexToolCallContextItemType(item.Get("type").String()) {
  301. return
  302. }
  303. key := strings.TrimSpace(item.Get("id").String())
  304. if key == "" {
  305. key = strings.TrimSpace(item.Get("call_id").String())
  306. }
  307. if key == "" {
  308. key = raw
  309. }
  310. if c.seen == nil {
  311. c.seen = make(map[string]struct{})
  312. }
  313. if _, ok := c.seen[key]; ok {
  314. return
  315. }
  316. c.seen[key] = struct{}{}
  317. c.items = append(c.items, json.RawMessage(raw))
  318. }
  319. func buildOpenAIWSHTTPBridgeErrorEvent(statusCode int, message string) []byte {
  320. message = strings.TrimSpace(message)
  321. if message == "" {
  322. message = http.StatusText(statusCode)
  323. }
  324. if message == "" {
  325. message = "upstream request failed"
  326. }
  327. event := map[string]any{
  328. "type": "error",
  329. "status": statusCode,
  330. "error": map[string]any{
  331. "type": "upstream_error",
  332. "message": message,
  333. },
  334. }
  335. body, err := json.Marshal(event)
  336. if err != nil {
  337. return []byte(`{"type":"error","error":{"type":"upstream_error","message":"upstream request failed"}}`)
  338. }
  339. return body
  340. }
  341. func buildOpenAIWSHTTPBridgeFailedEvent(responseID, model string, source []byte, fallbackMessage string) []byte {
  342. errorType := strings.TrimSpace(gjson.GetBytes(source, "error.type").String())
  343. if errorType == "" {
  344. errorType = strings.TrimSpace(gjson.GetBytes(source, "response.error.type").String())
  345. }
  346. code := strings.TrimSpace(gjson.GetBytes(source, "error.code").String())
  347. if code == "" {
  348. code = strings.TrimSpace(gjson.GetBytes(source, "response.error.code").String())
  349. }
  350. if code == "" {
  351. code = "upstream_error"
  352. }
  353. message := extractOpenAISSEErrorMessage(source)
  354. if message == "" {
  355. message = strings.TrimSpace(fallbackMessage)
  356. }
  357. if message == "" {
  358. message = "Upstream response failed"
  359. }
  360. errorBody := map[string]any{"code": code, "message": message}
  361. if errorType != "" {
  362. errorBody["type"] = errorType
  363. }
  364. response := map[string]any{
  365. "id": responseID, "object": "response", "status": "failed",
  366. "output": []any{}, "error": errorBody,
  367. }
  368. if model = strings.TrimSpace(model); model != "" {
  369. response["model"] = model
  370. }
  371. body, err := json.Marshal(map[string]any{"type": "response.failed", "response": response})
  372. if err != nil {
  373. return []byte(`{"type":"response.failed","response":{"status":"failed","output":[],"error":{"code":"upstream_error","message":"Upstream response failed"}}}`)
  374. }
  375. return body
  376. }
  377. func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
  378. ctx context.Context,
  379. c *gin.Context,
  380. account *Account,
  381. token string,
  382. payload []byte,
  383. payloadBytes int,
  384. originalModel string,
  385. imageBillingModel string,
  386. imageSizeTier string,
  387. imageInputSize string,
  388. grokCacheIdentity string,
  389. turn int,
  390. writeClientMessage func([]byte) error,
  391. ) (*OpenAIForwardResult, error) {
  392. if s == nil {
  393. return nil, errors.New("service is nil")
  394. }
  395. if s.httpUpstream == nil {
  396. return nil, errors.New("openai http upstream is nil")
  397. }
  398. if account == nil {
  399. return nil, errors.New("account is nil")
  400. }
  401. if writeClientMessage == nil {
  402. return nil, errors.New("client websocket writer is nil")
  403. }
  404. responseModelObserver := &upstreamResponseModelObserver{}
  405. body, err := prepareOpenAIWSHTTPBridgeBody(account, payload)
  406. if err != nil {
  407. return nil, fmt.Errorf("prepare http bridge body: %w", err)
  408. }
  409. grokIntentSourceBody := append([]byte(nil), body...)
  410. _, grokExplicitToolsField := openAIWSHTTPBridgeRawField(grokIntentSourceBody, "tools")
  411. grokExplicitToolIntent := account.Platform == PlatformGrok && hasGrokResponsesToolIntent(grokIntentSourceBody)
  412. var clientToolMapping apicompat.ResponsesClientToolMapping
  413. functionToolUpstream := (account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey) || account.Platform == PlatformGrok
  414. if functionToolUpstream {
  415. if account.Platform == PlatformGrok {
  416. body, err = sanitizeGrokResponsesInput(body)
  417. if err != nil {
  418. return nil, fmt.Errorf("sanitize Grok WS HTTP bridge input: %w", err)
  419. }
  420. }
  421. inheritedState, _ := openAIWSHTTPBridgeToolStateFromContext(c)
  422. inheritedLoweredTools := decodeOpenAIWSHTTPBridgeLoweredTools(inheritedState.LoweredTools)
  423. body, clientToolMapping, err = adaptResponsesClientToolsForFunctionUpstreamWithMapping(
  424. body,
  425. openAIWSHTTPBridgeToolUpstreamName(account),
  426. inheritedState.ClientMapping,
  427. inheritedLoweredTools,
  428. )
  429. if err != nil {
  430. return nil, fmt.Errorf("adapt %s client tools: %w", openAIWSHTTPBridgeToolUpstreamName(account), err)
  431. }
  432. if account.Platform == PlatformGrok && !grokExplicitToolsField && !grokExplicitToolIntent && len(inheritedLoweredTools) > 0 && hasGrokResponsesToolIntent(body) {
  433. // This continuation omitted tools, so the pre-adapter source cannot
  434. // represent the effective inherited declarations. Cache routing must
  435. // see the rehydrated tool intent or it will replace client functions
  436. // with the native-search tool-free route. Explicit current-turn tool
  437. // intent still uses the original pre-sanitization source above.
  438. grokIntentSourceBody = append(grokIntentSourceBody[:0], body...)
  439. }
  440. loweredTools := inheritedState.LoweredTools
  441. if currentTools, present := openAIWSHTTPBridgeRawField(body, "tools"); present {
  442. loweredTools = currentTools
  443. }
  444. setOpenAIWSHTTPBridgeToolState(c, openAIWSHTTPBridgeToolState{
  445. ClientMapping: clientToolMapping,
  446. LoweredTools: loweredTools,
  447. })
  448. }
  449. if account.Platform != PlatformGrok && isOpenAIResponsesLiteWebSocketPayload(payload) {
  450. liteBody, liteChanged, liteErr := normalizeOpenAIResponsesLitePayloadForAccount(body, account)
  451. if liteErr != nil {
  452. return nil, fmt.Errorf("normalize responses Lite payload: %w", liteErr)
  453. }
  454. if liteChanged {
  455. body = liteBody
  456. }
  457. }
  458. buildUpstreamRequest := func(requestBody []byte) (*http.Request, error) {
  459. upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
  460. defer releaseUpstreamCtx()
  461. var upstreamReq *http.Request
  462. var buildErr error
  463. if account.Platform == PlatformGrok {
  464. upstreamReq, buildErr = buildGrokResponsesRequest(upstreamCtx, c, account, requestBody, token, grokCacheIdentity, s.cfg, s.settingService)
  465. } else {
  466. upstreamReq, buildErr = s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, requestBody, token)
  467. }
  468. if buildErr != nil {
  469. return nil, buildErr
  470. }
  471. if account.Platform != PlatformGrok && isOpenAIResponsesLiteWebSocketPayload(payload) {
  472. upstreamReq.Header.Set(responsesLiteHeader, "true")
  473. }
  474. return upstreamReq, nil
  475. }
  476. if account.Platform == PlatformGrok {
  477. upstreamModel := resolveGrokWSUpstreamModel(account, body, originalModel)
  478. body, err = patchGrokResponsesBody(body, upstreamModel)
  479. if err != nil {
  480. return nil, err
  481. }
  482. grokMixedCacheIntentBody := append([]byte(nil), body...)
  483. body, err = applyGrokResponsesCacheIdentity(body, grokIntentSourceBody, grokCacheIdentity, account.IsGrokOAuth())
  484. if err != nil {
  485. return nil, fmt.Errorf("apply grok prompt cache identity: %w", err)
  486. }
  487. body, err = applyGrokFreeRequestToolCacheRoute(c, body, grokMixedCacheIntentBody, account, grokCacheIdentity)
  488. if err != nil {
  489. return nil, fmt.Errorf("apply grok Free function-tool cache route: %w", err)
  490. }
  491. }
  492. actualModel := strings.TrimSpace(gjson.GetBytes(body, "model").String())
  493. if actualModel == "" {
  494. actualModel = canonicalOpenAIAccountSchedulingModel(account, originalModel)
  495. }
  496. SetOpsUpstreamModel(c, actualModel)
  497. proxyURL := ""
  498. if account.ProxyID != nil && account.Proxy != nil {
  499. proxyURL = account.Proxy.URL()
  500. }
  501. if c != nil {
  502. c.Set("openai_passthrough", true)
  503. c.Set("openai_ws_http_bridge", true)
  504. }
  505. turnStart := time.Now()
  506. rejectedFieldRetryState := newOpenAIResponsesRejectedFieldRetryState(body)
  507. var resp *http.Response
  508. for {
  509. upstreamReq, buildErr := buildUpstreamRequest(body)
  510. if buildErr != nil {
  511. return nil, buildErr
  512. }
  513. resp, err = s.doOpenAIUpstream(upstreamReq, proxyURL, account)
  514. if err != nil {
  515. if turn == 1 {
  516. return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true)
  517. }
  518. safeErr := sanitizeUpstreamErrorMessage(err.Error())
  519. clientError := buildOpenAIWSHTTPBridgeErrorEvent(http.StatusBadGateway, "Upstream request failed")
  520. if writeErr := writeClientMessage(clientError); writeErr == nil {
  521. markOpenAIWSClientVisibleFailure(c, "error", clientError)
  522. }
  523. return nil, fmt.Errorf("upstream http bridge request failed: %s", safeErr)
  524. }
  525. if resp.StatusCode < 400 {
  526. break
  527. }
  528. respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes))
  529. _ = resp.Body.Close()
  530. markOpenAICyberPolicyEvent(c, respBody, resp.StatusCode, nil)
  531. if resp.StatusCode == http.StatusBadRequest &&
  532. extractUpstreamErrorCode(respBody) == openAIWSFallbackReasonInvalidEncryptedContent {
  533. s.markOpenAIWSInvalidEncryptedContentLineageFromPayload(
  534. c, body, "ingress_ws_http_bridge_invalid_encrypted_lineage_mark", account.ID, turn,
  535. )
  536. }
  537. retryBody, retryReason, changed, retryErr := normalizeOpenAIResponsesRejectedFieldRetryBody(resp.StatusCode, body, respBody)
  538. if retryErr != nil {
  539. return nil, fmt.Errorf("normalize websocket http bridge rejected field retry: %w", retryErr)
  540. }
  541. if changed && rejectedFieldRetryState.Allow(retryBody) {
  542. logOpenAIWSModeInfo(
  543. "ingress_ws_http_bridge_rejected_field_retry account_id=%d turn=%d reason=%s",
  544. account.ID,
  545. turn,
  546. truncateOpenAIWSLogValue(retryReason, openAIWSLogValueMaxLen),
  547. )
  548. body = retryBody
  549. payloadBytes = len(body)
  550. continue
  551. }
  552. upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
  553. if upstreamMsg == "" {
  554. upstreamMsg = http.StatusText(resp.StatusCode)
  555. }
  556. shouldFailover := s.shouldFailoverOpenAIUpstreamResponse(account, resp.StatusCode, upstreamMsg, respBody)
  557. if account.Platform == PlatformGrok {
  558. shouldFailover = s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody)
  559. s.handleGrokAccountUpstreamError(withGrokTeamRateLimitModel(ctx, resolveGrokWSUpstreamModel(account, body, originalModel)), account, resp.StatusCode, resp.Header, respBody)
  560. if shouldFailover && (turn == 1 || resp.StatusCode == http.StatusTooManyRequests) {
  561. return nil, newOpenAIUpstreamFailoverError(resp.StatusCode, resp.Header, respBody, upstreamMsg, false)
  562. }
  563. } else if shouldFailover && (turn == 1 || resp.StatusCode == http.StatusTooManyRequests) {
  564. return nil, s.handleFailoverErrorResponsePassthrough(ctx, resp, c, account, body, respBody)
  565. }
  566. if account.Platform != PlatformGrok && (shouldFailover || shouldCooldownOpenAITransientUpstreamError(resp.StatusCode, respBody)) {
  567. s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, actualModel)
  568. }
  569. clientError := buildOpenAIWSHTTPBridgeErrorEvent(resp.StatusCode, upstreamMsg)
  570. if writeErr := writeClientMessage(clientError); writeErr == nil {
  571. markOpenAIWSClientVisibleFailure(c, "error", clientError)
  572. }
  573. return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg)
  574. }
  575. defer func() { _ = resp.Body.Close() }()
  576. stopCancelBody := context.AfterFunc(ctx, func() { _ = resp.Body.Close() })
  577. defer stopCancelBody()
  578. if account.Platform == PlatformGrok {
  579. s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, resolveGrokWSUpstreamModel(account, body, originalModel)), account, resp.Header, resp.StatusCode)
  580. }
  581. responseID := ""
  582. usage := OpenAIUsage{}
  583. imageCounter := newOpenAIImageOutputCounter()
  584. var firstTokenMs *int
  585. reqStream := openAIWSPayloadBoolFromRaw(body, "stream", true)
  586. eventCount := 0
  587. tokenEventCount := 0
  588. terminalEventCount := 0
  589. replayCollector := &openAIWSToolCallReplayCollector{}
  590. firstEventType := ""
  591. lastEventType := ""
  592. upstreamTerminalEvent := ""
  593. sawDone := false
  594. wroteDownstream := false
  595. pendingClientMessages := make([][]byte, 0, 4)
  596. pendingClientMessageBytes := int64(0)
  597. capacityFailoverSuppressedLogged := false
  598. clientDisconnected := false
  599. officialOpenAIResponses := account != nil && account.Platform == PlatformOpenAI
  600. bareErrorPending := false
  601. var bareErrorPayload []byte
  602. bareErrorMessage := ""
  603. failureAccountSideEffectsApplied := false
  604. mappedModel := actualModel
  605. needModelReplace := false
  606. var mappedModelBytes []byte
  607. if originalModel != "" {
  608. needModelReplace = mappedModel != "" && mappedModel != originalModel
  609. if needModelReplace {
  610. mappedModelBytes = []byte(mappedModel)
  611. }
  612. }
  613. resultWithUsage := func() *OpenAIForwardResult {
  614. imageCount := imageCounter.Count()
  615. result := &OpenAIForwardResult{
  616. RequestID: responseID,
  617. Usage: usage,
  618. Model: originalModel,
  619. UpstreamModel: mappedModel,
  620. UpstreamResponseModel: responseModelObserver.Model(),
  621. UpstreamResponseModelConflict: responseModelObserver.Conflict(),
  622. UpstreamResponseServiceTier: responseModelObserver.ServiceTier(),
  623. ServiceTier: resolvedOpenAIUpstreamServiceTierFromObserver(responseModelObserver, extractOpenAIServiceTierFromBody(body)),
  624. ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, mappedModel, originalModel), body, mappedModel),
  625. RequestedReasoningEffort: CanonicalRequestedReasoningEffort(body, originalModel, mappedModel),
  626. Stream: reqStream,
  627. OpenAIWSMode: true,
  628. UpstreamTerminalEvent: upstreamTerminalEvent,
  629. ResponseHeaders: cloneHeader(resp.Header),
  630. Duration: time.Since(turnStart),
  631. FirstTokenMs: firstTokenMs,
  632. }
  633. if replayInput := replayCollector.Items(); len(replayInput) > 0 {
  634. result.wsReplayInput = replayInput
  635. result.wsReplayInputExists = true
  636. }
  637. result.wsAccountFailoverReplayInput = replayCollector.AllItems()
  638. if imageCount > 0 {
  639. result.ImageCount = imageCount
  640. result.ImageSize = imageSizeTier
  641. result.ImageInputSize = imageInputSize
  642. result.ImageOutputSizes = imageCounter.Sizes()
  643. result.BillingModel = imageBillingModel
  644. }
  645. return result
  646. }
  647. maxLineSize := defaultMaxLineSize
  648. if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
  649. maxLineSize = s.cfg.Gateway.MaxLineSize
  650. }
  651. if hasResponsesClientToolMapping(clientToolMapping) {
  652. resp.Body = newResponsesClientToolStreamBody(resp.Body, clientToolMapping, maxLineSize)
  653. }
  654. scanner := bufio.NewScanner(resp.Body)
  655. scanBuf := getSSEScannerBuf64K()
  656. scanner.Buffer(scanBuf[:0], maxLineSize)
  657. defer putSSEScannerBuf64K(scanBuf)
  658. pendingSSEEventType := ""
  659. finalizeBareError := func() error {
  660. if !bareErrorPending {
  661. return nil
  662. }
  663. if !failureAccountSideEffectsApplied {
  664. failureAccountSideEffectsApplied = s.handleOpenAIWSFailureAccountSideEffects(ctx, account, mappedModel, resp.Header, bareErrorPayload)
  665. }
  666. upstreamTerminalEvent = "response.failed"
  667. if clientDisconnected {
  668. return nil
  669. }
  670. clientMessage := buildOpenAIWSHTTPBridgeFailedEvent(responseID, originalModel, bareErrorPayload, bareErrorMessage)
  671. if rewritten, changed := sanitizeOpenAICapacityShedErrorCodeForClient(clientMessage); changed {
  672. clientMessage = rewritten
  673. }
  674. messages := append(pendingClientMessages, clientMessage)
  675. pendingClientMessages = nil
  676. pendingClientMessageBytes = 0
  677. for _, message := range messages {
  678. if err := writeClientMessage(message); err != nil {
  679. if isOpenAIWSClientDisconnectError(err) {
  680. clientDisconnected = true
  681. return nil
  682. }
  683. return fmt.Errorf("write synthesized websocket response.failed: %w", err)
  684. }
  685. wroteDownstream = true
  686. }
  687. markOpenAIWSClientVisibleFailure(c, "response.failed", clientMessage)
  688. return nil
  689. }
  690. for scanner.Scan() {
  691. line := scanner.Text()
  692. if eventType, ok := extractOpenAISSEEventLine(line); ok {
  693. pendingSSEEventType = eventType
  694. continue
  695. }
  696. if strings.TrimSpace(line) == "" {
  697. pendingSSEEventType = ""
  698. continue
  699. }
  700. data, ok := extractOpenAISSEDataLine(line)
  701. if !ok {
  702. continue
  703. }
  704. trimmedData := strings.TrimSpace(data)
  705. if trimmedData == "" {
  706. continue
  707. }
  708. if trimmedData == "[DONE]" {
  709. sawDone = true
  710. continue
  711. }
  712. upstreamMessage := []byte(openAICompatPayloadWithEventType(trimmedData, pendingSSEEventType))
  713. if normalized, changed := normalizeCompletedImageGenerationStatus(upstreamMessage); changed {
  714. upstreamMessage = normalized
  715. }
  716. eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage)
  717. responseModelObserver.ObserveOpenAI(upstreamMessage, eventType)
  718. if responseID == "" && eventResponseID != "" {
  719. responseID = eventResponseID
  720. }
  721. if eventType != "" {
  722. eventCount++
  723. if firstEventType == "" {
  724. firstEventType = eventType
  725. }
  726. lastEventType = eventType
  727. }
  728. if isOpenAIWSTokenEvent(eventType) {
  729. tokenEventCount++
  730. if firstTokenMs == nil {
  731. ms := int(time.Since(turnStart).Milliseconds())
  732. firstTokenMs = &ms
  733. }
  734. }
  735. if openAIWSMessageShouldParseUsage(eventType, upstreamMessage) {
  736. parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage)
  737. }
  738. if eventType == "error" || eventType == "response.failed" {
  739. markOpenAICyberPolicyEvent(c, upstreamMessage, http.StatusOK, &usage)
  740. }
  741. imageCounter.AddSSEData(upstreamMessage)
  742. if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && strings.Contains(trimmedData, mappedModel) {
  743. upstreamMessage = replaceOpenAIWSMessageModel(upstreamMessage, mappedModel, originalModel)
  744. }
  745. if s.toolCorrector != nil && openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(upstreamMessage) {
  746. if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(upstreamMessage); changed {
  747. upstreamMessage = corrected
  748. }
  749. }
  750. replayCollector.AddEvent(eventType, upstreamMessage)
  751. var upstreamEventErr error
  752. if officialOpenAIResponses && bareErrorPending && (eventType == "response.completed" || eventType == "response.done") {
  753. // Some upstreams emit a recoverable bare error before the authoritative
  754. // successful terminal. Do not replace that terminal with a synthetic
  755. // failure or retain side effects from the superseded error.
  756. bareErrorPending = false
  757. bareErrorPayload = nil
  758. bareErrorMessage = ""
  759. }
  760. suppressClientMessage := officialOpenAIResponses && bareErrorPending && eventType != "response.failed"
  761. if eventType == "error" || eventType == "response.failed" {
  762. errMessage := extractOpenAISSEErrorMessage(upstreamMessage)
  763. if errMessage == "" {
  764. errMessage = "upstream error event"
  765. }
  766. statusCode := openAIStreamFailureStatus(upstreamMessage, errMessage)
  767. shouldFailover := openAIStreamFailedEventShouldFailover(upstreamMessage, errMessage)
  768. if eventType == "error" {
  769. errCodeRaw, errTypeRaw, _ := parseOpenAIWSErrorEventFields(upstreamMessage)
  770. shouldFailover = openAIStreamErrorEventShouldFailover(upstreamMessage, errMessage)
  771. if account.Platform == PlatformGrok {
  772. statusCode = openAIWSErrorHTTPStatusFromRaw(errCodeRaw, errTypeRaw)
  773. }
  774. if reason, _ := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMessage); reason == openAIWSFallbackReasonInvalidEncryptedContent {
  775. s.markOpenAIWSInvalidEncryptedContentLineageFromPayload(
  776. c, body, "ingress_ws_http_bridge_invalid_encrypted_lineage_mark", account.ID, turn,
  777. )
  778. }
  779. }
  780. requestScopedCapacity := isOpenAIUpstreamCapacityShedEvent(upstreamMessage)
  781. if account.Platform == PlatformGrok && eventType == "error" {
  782. // SSE error events do not carry an HTTP status. The local status
  783. // mapper therefore defaults unknown xAI codes (for example
  784. // new_sensitive) to 502; classify the body as a request-scoped
  785. // 403 before applying status-based failover or account state.
  786. if isGrokContentPolicyRejection(http.StatusForbidden, upstreamMessage) {
  787. shouldFailover = false
  788. } else {
  789. shouldFailover = s.shouldFailoverGrokUpstreamError(statusCode, upstreamMessage)
  790. s.handleGrokAccountUpstreamError(ctx, account, statusCode, resp.Header, upstreamMessage)
  791. }
  792. }
  793. if !wroteDownstream && shouldFailover && (turn == 1 || statusCode == http.StatusTooManyRequests) {
  794. if account.Platform == PlatformGrok {
  795. return nil, newOpenAIUpstreamFailoverError(statusCode, resp.Header, upstreamMessage, errMessage, false)
  796. }
  797. return nil, s.newOpenAIStreamFailoverErrorWithModel(c, account, true, resp.Header.Get("x-request-id"), upstreamMessage, errMessage, mappedModel, resp.Header)
  798. }
  799. if account.Platform != PlatformGrok && !failureAccountSideEffectsApplied {
  800. if eventType == "response.failed" || (!officialOpenAIResponses && shouldFailover && !requestScopedCapacity) {
  801. failureAccountSideEffectsApplied = s.handleOpenAIWSFailureAccountSideEffects(ctx, account, mappedModel, resp.Header, upstreamMessage)
  802. }
  803. }
  804. if wroteDownstream && requestScopedCapacity && !capacityFailoverSuppressedLogged {
  805. logOpenAICapacityFailoverSuppressed(ctx, account, "ws_http_bridge", resp.Header.Get("x-request-id"), eventType)
  806. capacityFailoverSuppressedLogged = true
  807. }
  808. if eventType == "error" && !officialOpenAIResponses {
  809. upstreamEventErr = errors.New(errMessage)
  810. } else if eventType == "error" {
  811. bareErrorPending = true
  812. bareErrorPayload = append(bareErrorPayload[:0], upstreamMessage...)
  813. bareErrorMessage = errMessage
  814. suppressClientMessage = true
  815. } else {
  816. bareErrorPending = false
  817. }
  818. }
  819. // 客户端写出副本改写容量降载码:Codex 对 error/response.failed 中的
  820. // server_is_overloaded / slow_down 判致命并终止会话,改写后走客户端内置
  821. // 重试。账号状态与终止事件判定(下方 handleOpenAIWSTerminalTransientFailure)
  822. // 仍使用未改写的 upstreamMessage。
  823. clientMessage := upstreamMessage
  824. if eventType == "error" || eventType == "response.failed" {
  825. if rewritten, changed := sanitizeOpenAICapacityShedErrorCodeForClient(clientMessage); changed {
  826. clientMessage = rewritten
  827. }
  828. }
  829. if !clientDisconnected && !suppressClientMessage {
  830. stageBeforeSemanticOutput := turn == 1 && account.Platform == PlatformOpenAI && !wroteDownstream
  831. commitStagedMessages := !stageBeforeSemanticOutput ||
  832. openAIStreamDataStartsClientOutput(string(clientMessage), eventType) ||
  833. isOpenAIWSTerminalEvent(eventType)
  834. if stageBeforeSemanticOutput && !commitStagedMessages {
  835. if pendingClientMessageBytes+int64(len(clientMessage)) > openAIFirstOutputStageMaxBytes {
  836. return nil, s.newOpenAIStreamFailoverError(
  837. c,
  838. account,
  839. true,
  840. resp.Header.Get("x-request-id"),
  841. nil,
  842. "OpenAI WS HTTP bridge first-output staging limit exceeded",
  843. resp.Header,
  844. )
  845. }
  846. pendingClientMessages = append(pendingClientMessages, append([]byte(nil), clientMessage...))
  847. pendingClientMessageBytes += int64(len(clientMessage))
  848. } else {
  849. messages := append(pendingClientMessages, clientMessage)
  850. pendingClientMessages = nil
  851. pendingClientMessageBytes = 0
  852. for _, message := range messages {
  853. if err := writeClientMessage(message); err != nil {
  854. if isOpenAIWSClientDisconnectError(err) {
  855. clientDisconnected = true
  856. closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err)
  857. logOpenAIWSModeInfo(
  858. "ingress_ws_http_bridge_client_disconnected_drain account_id=%d turn=%d close_status=%s close_reason=%s",
  859. account.ID,
  860. turn,
  861. closeStatus,
  862. truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen),
  863. )
  864. break
  865. }
  866. return nil, wrapOpenAIWSIngressTurnError(
  867. "write_client",
  868. fmt.Errorf("write client websocket event: %w", err),
  869. wroteDownstream,
  870. )
  871. }
  872. wroteDownstream = true
  873. }
  874. }
  875. }
  876. if !clientDisconnected && !suppressClientMessage {
  877. markOpenAIWSClientVisibleFailure(c, eventType, upstreamMessage)
  878. }
  879. if upstreamEventErr != nil {
  880. return resultWithUsage(), upstreamEventErr
  881. }
  882. if isOpenAIWSTerminalEvent(eventType) && !bareErrorPending {
  883. if eventType == "response.failed" {
  884. upstreamTerminalEvent = "response.failed"
  885. } else {
  886. upstreamTerminalEvent = s.handleOpenAIWSTerminalTransientFailure(ctx, account, mappedModel, resp.Header, upstreamMessage)
  887. }
  888. terminalEventCount++
  889. firstTokenMsValue := -1
  890. if firstTokenMs != nil {
  891. firstTokenMsValue = *firstTokenMs
  892. }
  893. logOpenAIWSModeInfo(
  894. "ingress_ws_http_bridge_turn_completed account_id=%d turn=%d response_id=%s payload_bytes=%d duration_ms=%d events=%d token_events=%d terminal_events=%d first_event=%s last_event=%s first_token_ms=%d client_disconnected=%v",
  895. account.ID,
  896. turn,
  897. truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen),
  898. payloadBytes,
  899. time.Since(turnStart).Milliseconds(),
  900. eventCount,
  901. tokenEventCount,
  902. terminalEventCount,
  903. truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen),
  904. truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen),
  905. firstTokenMsValue,
  906. clientDisconnected,
  907. )
  908. return resultWithUsage(), nil
  909. }
  910. }
  911. if bareErrorPending {
  912. if finalizeErr := finalizeBareError(); finalizeErr != nil {
  913. return resultWithUsage(), finalizeErr
  914. }
  915. if scanErr := scanner.Err(); scanErr != nil {
  916. return resultWithUsage(), fmt.Errorf("read upstream http bridge stream after error event: %w", scanErr)
  917. }
  918. return resultWithUsage(), errors.New(bareErrorMessage)
  919. }
  920. if err := scanner.Err(); err != nil {
  921. streamErr := fmt.Errorf("read upstream http bridge stream: %w", err)
  922. if turn == 1 && !wroteDownstream {
  923. return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, streamErr, true)
  924. }
  925. return resultWithUsage(), streamErr
  926. }
  927. terminalErr := errors.New("upstream http bridge stream ended before terminal event")
  928. if sawDone {
  929. terminalErr = errors.New("upstream http bridge stream sent [DONE] before terminal event")
  930. }
  931. if turn == 1 && !wroteDownstream {
  932. return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, terminalErr, true)
  933. }
  934. return resultWithUsage(), terminalErr
  935. }
  936. func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, seedPayload, currentPayload []byte, originalModel string) (string, error) {
  937. body, err := prepareOpenAIWSHTTPBridgeBody(account, seedPayload)
  938. if err != nil {
  939. return "", err
  940. }
  941. upstreamModel := resolveGrokWSUpstreamModel(account, currentPayload, originalModel)
  942. body, err = patchGrokResponsesBody(body, upstreamModel)
  943. if err != nil {
  944. return "", err
  945. }
  946. return resolveGrokCacheIdentity(c, body, "", upstreamModel), nil
  947. }
  948. func resolveGrokWSUpstreamModel(account *Account, body []byte, originalModel string) string {
  949. upstreamModel := strings.TrimSpace(gjson.GetBytes(body, "model").String())
  950. originalModel = strings.TrimSpace(originalModel)
  951. // Shared ingress has already applied channel and account mappings when the
  952. // body model differs from the client-facing model. Only resolve from the
  953. // original model when the body still carries that original value.
  954. if account != nil && originalModel != "" && (upstreamModel == "" || upstreamModel == originalModel) {
  955. if mappedModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)); mappedModel != "" {
  956. upstreamModel = mappedModel
  957. }
  958. }
  959. if upstreamModel == "" {
  960. upstreamModel = grokDefaultResponsesModel
  961. }
  962. return upstreamModel
  963. }