openai_ws_forwarder_support.go 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "net/http"
  8. "strings"
  9. "time"
  10. "github.com/gin-gonic/gin"
  11. "github.com/tidwall/gjson"
  12. "github.com/tidwall/sjson"
  13. )
  14. func (s *OpenAIGatewayService) isOpenAIWSGeneratePrewarmEnabled() bool {
  15. return s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.PrewarmGenerateEnabled
  16. }
  17. // performOpenAIWSGeneratePrewarm 在 WSv2 下执行可选的 generate=false 预热。
  18. // 预热默认关闭,仅在配置开启后生效;失败时按可恢复错误回退到 HTTP。
  19. func (s *OpenAIGatewayService) performOpenAIWSGeneratePrewarm(
  20. ctx context.Context,
  21. lease *openAIWSConnLease,
  22. decision OpenAIWSProtocolDecision,
  23. payload map[string]any,
  24. previousResponseID string,
  25. reqBody map[string]any,
  26. account *Account,
  27. stateStore OpenAIWSStateStore,
  28. groupID int64,
  29. ) error {
  30. if s == nil {
  31. return nil
  32. }
  33. if lease == nil || account == nil {
  34. logOpenAIWSModeInfo("prewarm_skip reason=invalid_state has_lease=%v has_account=%v", lease != nil, account != nil)
  35. return nil
  36. }
  37. connID := strings.TrimSpace(lease.ConnID())
  38. if !s.isOpenAIWSGeneratePrewarmEnabled() {
  39. return nil
  40. }
  41. if decision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 {
  42. logOpenAIWSModeInfo(
  43. "prewarm_skip account_id=%d conn_id=%s reason=transport_not_v2 transport=%s",
  44. account.ID,
  45. connID,
  46. normalizeOpenAIWSLogValue(string(decision.Transport)),
  47. )
  48. return nil
  49. }
  50. if strings.TrimSpace(previousResponseID) != "" {
  51. logOpenAIWSModeInfo(
  52. "prewarm_skip account_id=%d conn_id=%s reason=has_previous_response_id previous_response_id=%s",
  53. account.ID,
  54. connID,
  55. truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen),
  56. )
  57. return nil
  58. }
  59. if lease.IsPrewarmed() {
  60. logOpenAIWSModeInfo("prewarm_skip account_id=%d conn_id=%s reason=already_prewarmed", account.ID, connID)
  61. return nil
  62. }
  63. if NeedsToolContinuation(reqBody) {
  64. logOpenAIWSModeInfo("prewarm_skip account_id=%d conn_id=%s reason=tool_continuation", account.ID, connID)
  65. return nil
  66. }
  67. prewarmStart := time.Now()
  68. logOpenAIWSModeInfo("prewarm_start account_id=%d conn_id=%s", account.ID, connID)
  69. prewarmPayload := make(map[string]any, len(payload)+1)
  70. for k, v := range payload {
  71. prewarmPayload[k] = v
  72. }
  73. prewarmPayload["generate"] = false
  74. prewarmPayloadJSON := payloadAsJSONBytes(prewarmPayload)
  75. if err := lease.WriteJSONWithContextTimeout(ctx, prewarmPayload, s.openAIWSWriteTimeout()); err != nil {
  76. lease.MarkBroken()
  77. logOpenAIWSModeInfo(
  78. "prewarm_write_fail account_id=%d conn_id=%s cause=%s",
  79. account.ID,
  80. connID,
  81. truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen),
  82. )
  83. return wrapOpenAIWSFallback("prewarm_write", err)
  84. }
  85. logOpenAIWSModeInfo("prewarm_write_sent account_id=%d conn_id=%s payload_bytes=%d", account.ID, connID, len(prewarmPayloadJSON))
  86. prewarmResponseID := ""
  87. prewarmEventCount := 0
  88. prewarmTerminalCount := 0
  89. for {
  90. message, readErr := lease.ReadMessageWithContextTimeout(ctx, s.openAIWSReadTimeout())
  91. if readErr != nil {
  92. lease.MarkBroken()
  93. closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr)
  94. logOpenAIWSModeInfo(
  95. "prewarm_read_fail account_id=%d conn_id=%s close_status=%s close_reason=%s cause=%s events=%d",
  96. account.ID,
  97. connID,
  98. closeStatus,
  99. closeReason,
  100. truncateOpenAIWSLogValue(readErr.Error(), openAIWSLogValueMaxLen),
  101. prewarmEventCount,
  102. )
  103. return wrapOpenAIWSFallback("prewarm_"+classifyOpenAIWSReadFallbackReason(readErr), readErr)
  104. }
  105. eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(message)
  106. if eventType == "" {
  107. continue
  108. }
  109. prewarmEventCount++
  110. if prewarmResponseID == "" && eventResponseID != "" {
  111. prewarmResponseID = eventResponseID
  112. }
  113. if prewarmEventCount <= openAIWSPrewarmEventLogHead || eventType == "error" || isOpenAIWSTerminalEvent(eventType) {
  114. logOpenAIWSModeInfo(
  115. "prewarm_event account_id=%d conn_id=%s idx=%d type=%s bytes=%d",
  116. account.ID,
  117. connID,
  118. prewarmEventCount,
  119. truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen),
  120. len(message),
  121. )
  122. }
  123. if eventType == "error" {
  124. errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message)
  125. prewarmModel, _ := reqBody["model"].(string)
  126. s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw, prewarmModel)
  127. errMsg := strings.TrimSpace(errMsgRaw)
  128. if errMsg == "" {
  129. errMsg = "OpenAI websocket prewarm error"
  130. }
  131. fallbackReason, canFallback := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw)
  132. errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw)
  133. logOpenAIWSModeInfo(
  134. "prewarm_error_event account_id=%d conn_id=%s idx=%d fallback_reason=%s can_fallback=%v err_code=%s err_type=%s err_message=%s",
  135. account.ID,
  136. connID,
  137. prewarmEventCount,
  138. truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen),
  139. canFallback,
  140. errCode,
  141. errType,
  142. errMessage,
  143. )
  144. lease.MarkBroken()
  145. if canFallback {
  146. return wrapOpenAIWSFallback("prewarm_"+fallbackReason, errors.New(errMsg))
  147. }
  148. return wrapOpenAIWSFallback("prewarm_error_event", errors.New(errMsg))
  149. }
  150. if isOpenAIWSTerminalEvent(eventType) {
  151. prewarmTerminalCount++
  152. break
  153. }
  154. }
  155. lease.MarkPrewarmed()
  156. if prewarmResponseID != "" && stateStore != nil {
  157. ttl := s.openAIWSResponseStickyTTL()
  158. logOpenAIWSBindResponseAccountWarn(groupID, account.ID, prewarmResponseID, stateStore.BindResponseAccount(ctx, groupID, prewarmResponseID, account.ID, ttl))
  159. stateStore.BindResponseConn(prewarmResponseID, lease.ConnID(), ttl)
  160. }
  161. logOpenAIWSModeInfo(
  162. "prewarm_done account_id=%d conn_id=%s response_id=%s events=%d terminal_events=%d duration_ms=%d",
  163. account.ID,
  164. connID,
  165. truncateOpenAIWSLogValue(prewarmResponseID, openAIWSIDValueMaxLen),
  166. prewarmEventCount,
  167. prewarmTerminalCount,
  168. time.Since(prewarmStart).Milliseconds(),
  169. )
  170. return nil
  171. }
  172. func payloadAsJSON(payload map[string]any) string {
  173. return string(payloadAsJSONBytes(payload))
  174. }
  175. func payloadAsJSONBytes(payload map[string]any) []byte {
  176. if len(payload) == 0 {
  177. return []byte("{}")
  178. }
  179. body, err := json.Marshal(payload)
  180. if err != nil {
  181. return []byte("{}")
  182. }
  183. return body
  184. }
  185. func isOpenAIWSTerminalEvent(eventType string) bool {
  186. switch strings.TrimSpace(eventType) {
  187. case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled":
  188. return true
  189. default:
  190. return false
  191. }
  192. }
  193. func normalizeOpenAIWSTerminalEvent(eventType string) string {
  194. switch strings.TrimSpace(eventType) {
  195. case "response.completed":
  196. return "response.completed"
  197. case "response.done":
  198. return "response.done"
  199. case "response.failed":
  200. return "response.failed"
  201. case "response.incomplete":
  202. return "response.incomplete"
  203. case "response.cancelled", "response.canceled":
  204. return "response.cancelled"
  205. default:
  206. return ""
  207. }
  208. }
  209. // markOpenAIWSClientVisibleFailure records only terminal/error protocol events
  210. // that were delivered to the client. Callers invoke it only after any hidden
  211. // failover/recovery decision and a successful downstream write.
  212. func markOpenAIWSClientVisibleFailure(c *gin.Context, eventType string, payload []byte) {
  213. eventType = strings.TrimSpace(eventType)
  214. if eventType != "error" && eventType != "response.failed" {
  215. return
  216. }
  217. prefix := "error"
  218. if eventType == "response.failed" {
  219. prefix = "response.error"
  220. }
  221. code := strings.TrimSpace(gjson.GetBytes(payload, prefix+".code").String())
  222. errType := strings.TrimSpace(gjson.GetBytes(payload, prefix+".type").String())
  223. message := strings.TrimSpace(gjson.GetBytes(payload, prefix+".message").String())
  224. if eventType == "response.failed" && code == "" && errType == "" && message == "" {
  225. prefix = "error"
  226. code = strings.TrimSpace(gjson.GetBytes(payload, prefix+".code").String())
  227. errType = strings.TrimSpace(gjson.GetBytes(payload, prefix+".type").String())
  228. message = strings.TrimSpace(gjson.GetBytes(payload, prefix+".message").String())
  229. }
  230. status := int(gjson.GetBytes(payload, prefix+".status_code").Int())
  231. if status == 0 {
  232. status = int(gjson.GetBytes(payload, prefix+".status").Int())
  233. }
  234. if status == 0 && eventType == "error" {
  235. status = int(gjson.GetBytes(payload, "status").Int())
  236. }
  237. if status == 0 {
  238. status = openAIWSErrorHTTPStatusFromRaw(code, errType)
  239. }
  240. if errType == "" {
  241. errType = "upstream_error"
  242. }
  243. if code == "" {
  244. code = strings.ReplaceAll(eventType, ".", "_")
  245. }
  246. if message == "" {
  247. message = "upstream websocket request failed"
  248. }
  249. MarkOpsStreamFailure(c, errType, code, message, status)
  250. }
  251. func openAIWSPayloadTransientStatus(payload []byte) int {
  252. if len(payload) == 0 {
  253. return 0
  254. }
  255. status := int(gjson.GetBytes(payload, "response.error.status_code").Int())
  256. if status == 0 {
  257. status = int(gjson.GetBytes(payload, "response.error.status").Int())
  258. }
  259. if status == 0 {
  260. status = int(gjson.GetBytes(payload, "error.status_code").Int())
  261. }
  262. if status == 0 {
  263. status = int(gjson.GetBytes(payload, "error.status").Int())
  264. }
  265. if shouldCooldownOpenAITransientUpstreamError(status, payload) {
  266. return status
  267. }
  268. if status != 0 {
  269. return 0
  270. }
  271. code := strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "response.error.code").String()))
  272. errType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "response.error.type").String()))
  273. if code == "" {
  274. code = strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "error.code").String()))
  275. }
  276. if errType == "" {
  277. errType = strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "error.type").String()))
  278. }
  279. switch {
  280. case code == "server_is_overloaded", code == "slow_down":
  281. return http.StatusServiceUnavailable
  282. case strings.Contains(code, "server_error"),
  283. strings.Contains(code, "internal_error"),
  284. strings.Contains(code, "upstream_error"),
  285. strings.Contains(errType, "server_error"),
  286. strings.Contains(errType, "internal_error"),
  287. strings.Contains(errType, "upstream_error"):
  288. return http.StatusInternalServerError
  289. default:
  290. return 0
  291. }
  292. }
  293. func (s *OpenAIGatewayService) handleOpenAIWSTerminalTransientFailure(ctx context.Context, account *Account, canonicalModel string, headers http.Header, payload []byte) string {
  294. eventType, _, _ := parseOpenAIWSEventEnvelope(payload)
  295. terminalEvent := normalizeOpenAIWSTerminalEvent(eventType)
  296. if terminalEvent != "response.failed" {
  297. return terminalEvent
  298. }
  299. s.handleOpenAIWSFailureAccountSideEffects(ctx, account, canonicalModel, headers, payload)
  300. return terminalEvent
  301. }
  302. func (s *OpenAIGatewayService) handleOpenAIWSErrorEventTransientFailure(ctx context.Context, account *Account, canonicalModel string, headers http.Header, payload []byte) {
  303. eventType, _, _ := parseOpenAIWSEventEnvelope(payload)
  304. if eventType != "error" {
  305. return
  306. }
  307. status := openAIWSPayloadTransientStatus(payload)
  308. if status != 0 {
  309. if status == http.StatusTooManyRequests {
  310. headers = openAIWSSemantic429Headers(account, canonicalModel, headers)
  311. }
  312. s.handleOpenAIAccountUpstreamError(ctx, account, status, headers, payload, canonicalModel)
  313. }
  314. }
  315. // handleOpenAIWSFailureAccountSideEffects applies both structured credential
  316. // failures and transient failures. Its return value lets stream callers avoid
  317. // applying the same transition twice for an error/response.failed pair.
  318. func (s *OpenAIGatewayService) handleOpenAIWSFailureAccountSideEffects(ctx context.Context, account *Account, canonicalModel string, headers http.Header, payload []byte) bool {
  319. message := extractOpenAISSEErrorMessage(payload)
  320. status := openAIStreamFailureStatus(payload, message)
  321. switch status {
  322. case http.StatusUnauthorized, http.StatusTooManyRequests, 529:
  323. s.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, message, headers, canonicalModel)
  324. return true
  325. case http.StatusForbidden:
  326. if !openAIStream403AccountFailure(payload, message) {
  327. return false
  328. }
  329. s.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, message, headers, canonicalModel)
  330. return true
  331. }
  332. status = openAIWSPayloadTransientStatus(payload)
  333. if status == 0 {
  334. return false
  335. }
  336. s.handleOpenAIAccountUpstreamError(ctx, account, status, headers, payload, canonicalModel)
  337. return true
  338. }
  339. func (s *OpenAIGatewayService) handleOpenAIWSDialTransientFailure(ctx context.Context, account *Account, canonicalModel string, err error) {
  340. var dialErr *openAIWSDialError
  341. if !errors.As(err, &dialErr) || dialErr == nil || !shouldCooldownOpenAITransientUpstreamError(dialErr.StatusCode, dialErr.ResponseBody) {
  342. return
  343. }
  344. s.handleOpenAIAccountUpstreamError(ctx, account, dialErr.StatusCode, dialErr.ResponseHeaders, dialErr.ResponseBody, canonicalModel)
  345. }
  346. func isOpenAIWSTokenEvent(eventType string) bool {
  347. eventType = strings.TrimSpace(eventType)
  348. if eventType == "" {
  349. return false
  350. }
  351. switch eventType {
  352. case "response.created", "response.in_progress", "response.output_item.added", "response.output_item.done":
  353. return false
  354. }
  355. if strings.Contains(eventType, ".delta") {
  356. return true
  357. }
  358. if strings.HasPrefix(eventType, "response.output_text") {
  359. return true
  360. }
  361. if strings.HasPrefix(eventType, "response.output") {
  362. return true
  363. }
  364. // 终止事件(response.completed/done/failed/...)由 isOpenAIWSTerminalEvent 单独处理。
  365. // 不能把它们当作 token event,否则当上游没有可识别的 delta 时,
  366. // firstTokenMs 会被填到终止时刻,等于把"总耗时"误报为"首 token 延迟"。
  367. return false
  368. }
  369. func replaceOpenAIWSMessageModel(message []byte, fromModel, toModel string) []byte {
  370. if len(message) == 0 {
  371. return message
  372. }
  373. if strings.TrimSpace(fromModel) == "" || strings.TrimSpace(toModel) == "" || fromModel == toModel {
  374. return message
  375. }
  376. if !bytes.Contains(message, []byte(`"model"`)) || !bytes.Contains(message, []byte(fromModel)) {
  377. return message
  378. }
  379. modelValues := gjson.GetManyBytes(message, "model", "response.model")
  380. replaceModel := modelValues[0].Exists() && modelValues[0].Str == fromModel
  381. replaceResponseModel := modelValues[1].Exists() && modelValues[1].Str == fromModel
  382. if !replaceModel && !replaceResponseModel {
  383. return message
  384. }
  385. updated := message
  386. if replaceModel {
  387. if next, err := sjson.SetBytes(updated, "model", toModel); err == nil {
  388. updated = next
  389. }
  390. }
  391. if replaceResponseModel {
  392. if next, err := sjson.SetBytes(updated, "response.model", toModel); err == nil {
  393. updated = next
  394. }
  395. }
  396. return updated
  397. }
  398. func populateOpenAIUsageFromResponseJSON(body []byte, usage *OpenAIUsage) {
  399. if usage == nil || len(body) == 0 {
  400. return
  401. }
  402. if parsed, ok := extractOpenAIUsageFromJSONBytes(body); ok {
  403. *usage = parsed
  404. }
  405. }
  406. func getOpenAIGroupIDFromContext(c *gin.Context) int64 {
  407. if c == nil {
  408. return 0
  409. }
  410. value, exists := c.Get("api_key")
  411. if !exists {
  412. return 0
  413. }
  414. apiKey, ok := value.(*APIKey)
  415. if !ok || apiKey == nil || apiKey.GroupID == nil {
  416. return 0
  417. }
  418. return *apiKey.GroupID
  419. }
  420. // SelectAccountByPreviousResponseID 按 previous_response_id 命中账号粘连。
  421. // 未命中或账号不可用时返回 (nil, nil),由调用方继续走常规调度。
  422. func (s *OpenAIGatewayService) SelectAccountByPreviousResponseID(
  423. ctx context.Context,
  424. groupID *int64,
  425. previousResponseID string,
  426. requestedModel string,
  427. excludedIDs map[int64]struct{},
  428. requireCompact bool,
  429. ) (*AccountSelectionResult, error) {
  430. // 分组利润控制:公共入口装门,保证不经 selectAccountWithScheduler
  431. // 的调用方也无法绕过利润准入(scheduler 内部路径已在唯一调度入口装门)。
  432. ctx = s.withOpenAIProfitControlGate(ctx, groupID)
  433. return s.selectAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, "", requireCompact)
  434. }
  435. func (s *OpenAIGatewayService) selectAccountByPreviousResponseIDForCapability(
  436. ctx context.Context,
  437. groupID *int64,
  438. previousResponseID string,
  439. requestedModel string,
  440. excludedIDs map[int64]struct{},
  441. requiredCapability OpenAIEndpointCapability,
  442. requireCompact bool,
  443. ) (*AccountSelectionResult, error) {
  444. if s == nil {
  445. return nil, nil
  446. }
  447. accountID, account, responseID, store := s.resolveAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact)
  448. if accountID <= 0 || account == nil || store == nil {
  449. return nil, nil
  450. }
  451. result, acquireErr := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
  452. if acquireErr == nil && result.Acquired {
  453. logOpenAIWSBindResponseAccountWarn(
  454. derefGroupID(groupID),
  455. accountID,
  456. responseID,
  457. store.BindResponseAccount(ctx, derefGroupID(groupID), responseID, accountID, s.openAIWSResponseStickyTTL()),
  458. )
  459. return attachSelectionProfitGate(ctx, &AccountSelectionResult{
  460. Account: account,
  461. Acquired: true,
  462. ReleaseFunc: result.ReleaseFunc,
  463. }), nil
  464. }
  465. cfg := s.schedulingConfig()
  466. if s.concurrencyService != nil {
  467. return attachSelectionProfitGate(ctx, &AccountSelectionResult{
  468. Account: account,
  469. WaitPlan: &AccountWaitPlan{
  470. AccountID: accountID,
  471. MaxConcurrency: account.Concurrency,
  472. Timeout: cfg.StickySessionWaitTimeout,
  473. MaxWaiting: cfg.StickySessionMaxWaiting,
  474. },
  475. }), nil
  476. }
  477. return nil, nil
  478. }
  479. func (s *OpenAIGatewayService) ResolveAccountIDByPreviousResponseIDForScheduler(
  480. ctx context.Context,
  481. groupID *int64,
  482. previousResponseID string,
  483. requestedModel string,
  484. excludedIDs map[int64]struct{},
  485. requiredCapability OpenAIEndpointCapability,
  486. requireCompact bool,
  487. ) int64 {
  488. accountID, _, _, _ := s.resolveAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact)
  489. return accountID
  490. }
  491. func (s *OpenAIGatewayService) resolveAccountByPreviousResponseIDForCapability(
  492. ctx context.Context,
  493. groupID *int64,
  494. previousResponseID string,
  495. requestedModel string,
  496. excludedIDs map[int64]struct{},
  497. requiredCapability OpenAIEndpointCapability,
  498. requireCompact bool,
  499. ) (int64, *Account, string, OpenAIWSStateStore) {
  500. if s == nil {
  501. return 0, nil, "", nil
  502. }
  503. responseID := strings.TrimSpace(previousResponseID)
  504. if responseID == "" {
  505. return 0, nil, "", nil
  506. }
  507. store := s.getOpenAIWSStateStore()
  508. if store == nil {
  509. return 0, nil, "", nil
  510. }
  511. accountID, err := store.GetResponseAccount(ctx, derefGroupID(groupID), responseID)
  512. if err != nil || accountID <= 0 {
  513. return 0, nil, "", nil
  514. }
  515. if excludedIDs != nil {
  516. if _, excluded := excludedIDs[accountID]; excluded {
  517. return 0, nil, "", nil
  518. }
  519. }
  520. account, err := s.getSchedulableAccount(ctx, accountID)
  521. if err != nil || account == nil {
  522. _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
  523. return 0, nil, "", nil
  524. }
  525. // OAuth/SetupToken continuation state lives on the WSv2 session and cannot
  526. // survive an HTTP fallback. Official API-key Responses HTTP requests are
  527. // different: previous_response_id is supported by the provider and scoped to
  528. // the selected key/project, so the response-id binding must retain that key.
  529. if !account.IsOpenAIApiKey() && s.getOpenAIWSProtocolResolver().Resolve(account).Transport != OpenAIUpstreamTransportResponsesWebsocketV2 {
  530. return 0, nil, "", nil
  531. }
  532. if shouldClearStickySession(account, requestedModel) || !account.IsOpenAI() || !account.IsSchedulable() {
  533. _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
  534. return 0, nil, "", nil
  535. }
  536. if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) {
  537. _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
  538. return 0, nil, "", nil
  539. }
  540. if requestedModel != "" && !account.IsModelSupported(requestedModel) {
  541. return 0, nil, "", nil
  542. }
  543. if !account.SupportsOpenAIEndpointCapability(requiredCapability) {
  544. return 0, nil, "", nil
  545. }
  546. // Quota auto-pause must also gate the previous_response_id sticky path; otherwise an
  547. // account over its 5h/7d threshold keeps serving the same response chain even though
  548. // normal scheduling skips it. Pause is transient, so fall through to normal scheduling
  549. // without deleting the binding (the window may reset before the next turn).
  550. if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused {
  551. return 0, nil, "", nil
  552. }
  553. // 分组利润控制:与 quota auto-pause 同语义——利润不合格是暂时
  554. // 状态(上游倍率/高峰随时间变化),只跳过本次复用、落回普通调度,不删除
  555. // 绑定(倍率恢复后可继续按 previous_response_id 粘连)。
  556. if vetoed, _ := openAIProfitControlVetoReason(ctx, account); vetoed {
  557. return 0, nil, "", nil
  558. }
  559. if s.schedulerSnapshot != nil && s.accountRepo != nil {
  560. latest, latestErr := s.accountRepo.GetByID(ctx, account.ID)
  561. if latestErr != nil || latest == nil {
  562. _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
  563. return 0, nil, "", nil
  564. }
  565. if shouldClearStickySession(latest, requestedModel) || !latest.IsOpenAI() || !latest.IsSchedulable() {
  566. _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
  567. return 0, nil, "", nil
  568. }
  569. if !s.openAIAccountMatchesSchedulingGroup(latest, groupID) {
  570. return 0, nil, "", nil
  571. }
  572. if s.openAIGroupRequiresPrivacySet(ctx, groupID) && !latest.IsPrivacySet() {
  573. return 0, nil, "", nil
  574. }
  575. if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) {
  576. _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
  577. return 0, nil, "", nil
  578. }
  579. if requestedModel != "" && !latest.IsModelSupported(requestedModel) {
  580. return 0, nil, "", nil
  581. }
  582. if !latest.SupportsOpenAIEndpointCapability(requiredCapability) {
  583. return 0, nil, "", nil
  584. }
  585. if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, latest); paused {
  586. return 0, nil, "", nil
  587. }
  588. // 利润门对最新账号状态复检一次,语义同上:跳过复用、不删绑定。
  589. if vetoed, _ := openAIProfitControlVetoReason(ctx, latest); vetoed {
  590. return 0, nil, "", nil
  591. }
  592. if s.isOpenAIAccountRequestRuntimeBlocked(latest, requestedModel) {
  593. _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
  594. return 0, nil, "", nil
  595. }
  596. account = latest
  597. }
  598. if requireCompact && openAICompactSupportTier(account) == 0 {
  599. _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
  600. return 0, nil, "", nil
  601. }
  602. return accountID, account, responseID, store
  603. }
  604. func classifyOpenAIWSAcquireError(err error) string {
  605. if err == nil {
  606. return "acquire_conn"
  607. }
  608. var dialErr *openAIWSDialError
  609. if errors.As(err, &dialErr) {
  610. switch dialErr.StatusCode {
  611. case 426:
  612. return "upgrade_required"
  613. case 401, 403:
  614. return "auth_failed"
  615. case 429:
  616. return "upstream_rate_limited"
  617. }
  618. if dialErr.StatusCode >= 500 {
  619. return "upstream_5xx"
  620. }
  621. return "dial_failed"
  622. }
  623. if errors.Is(err, errOpenAIWSConnQueueFull) {
  624. return "conn_queue_full"
  625. }
  626. if errors.Is(err, errOpenAIWSPreferredConnUnavailable) {
  627. return "preferred_conn_unavailable"
  628. }
  629. if errors.Is(err, context.DeadlineExceeded) {
  630. return "acquire_timeout"
  631. }
  632. return "acquire_conn"
  633. }
  634. func isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw string) bool {
  635. code := strings.ToLower(strings.TrimSpace(codeRaw))
  636. errType := strings.ToLower(strings.TrimSpace(errTypeRaw))
  637. msg := strings.ToLower(strings.TrimSpace(msgRaw))
  638. if strings.Contains(errType, "rate_limit") || strings.Contains(errType, "usage_limit") {
  639. return true
  640. }
  641. if strings.Contains(code, "rate_limit") || strings.Contains(code, "usage_limit") || strings.Contains(code, "insufficient_quota") {
  642. return true
  643. }
  644. if strings.Contains(msg, "usage limit") && strings.Contains(msg, "reached") {
  645. return true
  646. }
  647. if strings.Contains(msg, "rate limit") && (strings.Contains(msg, "reached") || strings.Contains(msg, "exceeded")) {
  648. return true
  649. }
  650. return false
  651. }
  652. func (s *OpenAIGatewayService) persistOpenAIWSRateLimitSignal(ctx context.Context, account *Account, headers http.Header, responseBody []byte, codeRaw, errTypeRaw, msgRaw string, canonicalModel ...string) {
  653. if s == nil || s.rateLimitService == nil || account == nil || account.Platform != PlatformOpenAI {
  654. return
  655. }
  656. if !isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw) {
  657. return
  658. }
  659. model := firstNonEmpty(canonicalModel...)
  660. if model == "" {
  661. model = firstNonEmpty(gjson.GetBytes(responseBody, "model").String(), gjson.GetBytes(responseBody, "response.model").String())
  662. }
  663. // 非空 responseBody 表示已建立连接后收到的语义错误事件;握手响应头
  664. // 可能只是成功连接时的全局快照,不能用于普通模型的 429 账号级限流。
  665. // 实际拨号 HTTP 429 使用 nil responseBody,必须保留响应头。
  666. if len(responseBody) > 0 {
  667. headers = openAIWSSemantic429Headers(account, model, headers)
  668. }
  669. s.handleOpenAIAccountUpstreamError(ctx, account, http.StatusTooManyRequests, headers, responseBody, model)
  670. }
  671. func openAIWSSemantic429Headers(account *Account, model string, headers http.Header) http.Header {
  672. if isCodexSparkModel(model) && isOpenAIOAuthAccount(account) {
  673. return headers
  674. }
  675. return nil
  676. }
  677. func (s *OpenAIGatewayService) newOpenAIWSRateLimitFailoverError(account *Account, headers http.Header, responseBody []byte, message string) *UpstreamFailoverError {
  678. return s.newOpenAIAccountFailoverError(
  679. account,
  680. http.StatusTooManyRequests,
  681. headers,
  682. responseBody,
  683. strings.TrimSpace(message),
  684. false,
  685. false,
  686. )
  687. }
  688. func classifyOpenAIWSErrorEventFromRaw(codeRaw, errTypeRaw, msgRaw string) (string, bool) {
  689. code := strings.ToLower(strings.TrimSpace(codeRaw))
  690. errType := strings.ToLower(strings.TrimSpace(errTypeRaw))
  691. msg := strings.ToLower(strings.TrimSpace(msgRaw))
  692. switch code {
  693. case "upgrade_required":
  694. return "upgrade_required", true
  695. case "websocket_not_supported", "websocket_unsupported":
  696. return "ws_unsupported", true
  697. case "websocket_connection_limit_reached":
  698. return "ws_connection_limit_reached", true
  699. case "invalid_encrypted_content":
  700. return "invalid_encrypted_content", true
  701. case "previous_response_not_found":
  702. return "previous_response_not_found", true
  703. }
  704. if isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw) {
  705. return "upstream_rate_limited", false
  706. }
  707. if strings.Contains(msg, "upgrade required") || strings.Contains(msg, "status 426") {
  708. return "upgrade_required", true
  709. }
  710. if strings.Contains(errType, "upgrade") {
  711. return "upgrade_required", true
  712. }
  713. if strings.Contains(msg, "websocket") && strings.Contains(msg, "unsupported") {
  714. return "ws_unsupported", true
  715. }
  716. if strings.Contains(msg, "connection limit") && strings.Contains(msg, "websocket") {
  717. return "ws_connection_limit_reached", true
  718. }
  719. if strings.Contains(msg, "invalid_encrypted_content") ||
  720. (strings.Contains(msg, "encrypted content") && strings.Contains(msg, "could not be verified")) {
  721. return "invalid_encrypted_content", true
  722. }
  723. if strings.Contains(msg, "previous_response_not_found") ||
  724. (strings.Contains(msg, "previous response") && strings.Contains(msg, "not found")) {
  725. return "previous_response_not_found", true
  726. }
  727. if strings.Contains(errType, "server_error") || strings.Contains(code, "server_error") {
  728. return "upstream_error_event", true
  729. }
  730. return "event_error", false
  731. }
  732. func classifyOpenAIWSErrorEvent(message []byte) (string, bool) {
  733. if len(message) == 0 {
  734. return "event_error", false
  735. }
  736. return classifyOpenAIWSErrorEventFromRaw(parseOpenAIWSErrorEventFields(message))
  737. }
  738. func openAIWSErrorHTTPStatusFromRaw(codeRaw, errTypeRaw string) int {
  739. code := strings.ToLower(strings.TrimSpace(codeRaw))
  740. errType := strings.ToLower(strings.TrimSpace(errTypeRaw))
  741. switch {
  742. case strings.Contains(errType, "invalid_request"),
  743. strings.Contains(code, "invalid_request"),
  744. strings.Contains(code, "bad_request"),
  745. code == "invalid_encrypted_content",
  746. code == "previous_response_not_found":
  747. return http.StatusBadRequest
  748. case strings.Contains(errType, "authentication"),
  749. strings.Contains(code, "invalid_api_key"),
  750. strings.Contains(code, "unauthorized"):
  751. return http.StatusUnauthorized
  752. case strings.Contains(errType, "permission"),
  753. strings.Contains(code, "forbidden"):
  754. return http.StatusForbidden
  755. case isOpenAIWSRateLimitError(codeRaw, errTypeRaw, ""):
  756. return http.StatusTooManyRequests
  757. default:
  758. return http.StatusBadGateway
  759. }
  760. }
  761. func openAIWSErrorHTTPStatus(message []byte) int {
  762. if len(message) == 0 {
  763. return http.StatusBadGateway
  764. }
  765. codeRaw, errTypeRaw, _ := parseOpenAIWSErrorEventFields(message)
  766. return openAIWSErrorHTTPStatusFromRaw(codeRaw, errTypeRaw)
  767. }
  768. func (s *OpenAIGatewayService) openAIWSFallbackCooldown() time.Duration {
  769. if s == nil || s.cfg == nil {
  770. return 30 * time.Second
  771. }
  772. seconds := s.cfg.Gateway.OpenAIWS.FallbackCooldownSeconds
  773. if seconds <= 0 {
  774. return 0
  775. }
  776. return time.Duration(seconds) * time.Second
  777. }
  778. func (s *OpenAIGatewayService) isOpenAIWSFallbackCooling(accountID int64) bool {
  779. if s == nil || accountID <= 0 {
  780. return false
  781. }
  782. cooldown := s.openAIWSFallbackCooldown()
  783. if cooldown <= 0 {
  784. return false
  785. }
  786. rawUntil, ok := s.openaiWSFallbackUntil.Load(accountID)
  787. if !ok || rawUntil == nil {
  788. return false
  789. }
  790. until, ok := rawUntil.(time.Time)
  791. if !ok || until.IsZero() {
  792. s.openaiWSFallbackUntil.Delete(accountID)
  793. return false
  794. }
  795. if time.Now().Before(until) {
  796. return true
  797. }
  798. s.openaiWSFallbackUntil.Delete(accountID)
  799. return false
  800. }
  801. func (s *OpenAIGatewayService) markOpenAIWSFallbackCooling(accountID int64, _ string) {
  802. if s == nil || accountID <= 0 {
  803. return
  804. }
  805. cooldown := s.openAIWSFallbackCooldown()
  806. if cooldown <= 0 {
  807. return
  808. }
  809. s.openaiWSFallbackUntil.Store(accountID, time.Now().Add(cooldown))
  810. }
  811. func (s *OpenAIGatewayService) clearOpenAIWSFallbackCooling(accountID int64) {
  812. if s == nil || accountID <= 0 {
  813. return
  814. }
  815. s.openaiWSFallbackUntil.Delete(accountID)
  816. }