| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011 |
- package service
- import (
- "bufio"
- "context"
- "encoding/json"
- "errors"
- "fmt"
- "io"
- "net/http"
- "slices"
- "strings"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
- "github.com/gin-gonic/gin"
- "github.com/tidwall/gjson"
- )
- const (
- openAIWSClientReadLimitBytesDefault int64 = 64 * 1024 * 1024
- openAIWSHTTPBridgeThresholdBytesDefault int64 = 15 * 1024 * 1024
- openAIWSHTTPBridgeErrorBodyLimitBytes = 64 * 1024
- )
- const openAIWSHTTPBridgeToolStateContextKey = "openai_ws_http_bridge_tool_state"
- type openAIWSHTTPBridgeToolState struct {
- ClientMapping apicompat.ResponsesClientToolMapping
- LoweredTools json.RawMessage
- }
- func openAIWSHTTPBridgeToolStateFromContext(c *gin.Context) (openAIWSHTTPBridgeToolState, bool) {
- if c == nil {
- return openAIWSHTTPBridgeToolState{}, false
- }
- value, ok := c.Get(openAIWSHTTPBridgeToolStateContextKey)
- state, typed := value.(openAIWSHTTPBridgeToolState)
- return state, ok && typed
- }
- func setOpenAIWSHTTPBridgeToolState(c *gin.Context, state openAIWSHTTPBridgeToolState) {
- if c == nil {
- return
- }
- state.LoweredTools = append(json.RawMessage(nil), state.LoweredTools...)
- c.Set(openAIWSHTTPBridgeToolStateContextKey, state)
- }
- func decodeOpenAIWSHTTPBridgeLoweredTools(raw json.RawMessage) []any {
- if len(raw) == 0 {
- return nil
- }
- var tools []any
- if err := json.Unmarshal(raw, &tools); err != nil {
- return nil
- }
- return tools
- }
- func openAIWSHTTPBridgeRawField(body []byte, name string) (json.RawMessage, bool) {
- var fields map[string]json.RawMessage
- if err := json.Unmarshal(body, &fields); err != nil {
- return nil, false
- }
- raw, present := fields[name]
- return append(json.RawMessage(nil), raw...), present
- }
- func openAIWSHTTPBridgeToolUpstreamName(account *Account) string {
- if account != nil && account.Platform == PlatformGrok {
- return "Grok WS HTTP bridge"
- }
- return "OpenAI WS HTTP bridge"
- }
- // ResolveOpenAIWSClientFirstMessageTimeout returns the effective client ingress deadline.
- func ResolveOpenAIWSClientFirstMessageTimeout(cfg *config.Config) time.Duration {
- seconds := config.DefaultOpenAIWSClientFirstMessageTimeoutSeconds
- if cfg != nil && cfg.Gateway.OpenAIWS.ClientFirstMessageTimeoutSeconds > 0 {
- seconds = cfg.Gateway.OpenAIWS.ClientFirstMessageTimeoutSeconds
- }
- return time.Duration(seconds) * time.Second
- }
- func ResolveOpenAIWSClientReadLimitBytes(cfg *config.Config) int64 {
- if cfg == nil || cfg.Gateway.OpenAIWS.ClientReadLimitBytes <= 0 {
- return openAIWSClientReadLimitBytesDefault
- }
- return cfg.Gateway.OpenAIWS.ClientReadLimitBytes
- }
- func (s *OpenAIGatewayService) openAIWSHTTPBridgeEnabled() bool {
- return s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.HTTPBridgeEnabled
- }
- func (s *OpenAIGatewayService) openAIWSHTTPBridgeThresholdBytes() int64 {
- if s == nil || s.cfg == nil || s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes <= 0 {
- return openAIWSHTTPBridgeThresholdBytesDefault
- }
- return s.cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes
- }
- func (s *OpenAIGatewayService) shouldBridgeOpenAIWSHTTP(account *Account, payloadBytes int, previousResponseID string) bool {
- if account != nil && account.Platform == PlatformGrok {
- return true
- }
- if !s.openAIWSHTTPBridgeEnabled() {
- return false
- }
- if strings.TrimSpace(previousResponseID) != "" {
- return false
- }
- threshold := s.openAIWSHTTPBridgeThresholdBytes()
- return threshold > 0 && int64(payloadBytes) >= threshold
- }
- func (s *OpenAIGatewayService) shouldBridgeOpenAIWSPassthroughFirstMessage(account *Account, payload []byte) bool {
- if account != nil && account.Platform == PlatformGrok {
- return true
- }
- if !s.openAIWSHTTPBridgeEnabled() || int64(len(payload)) < s.openAIWSHTTPBridgeThresholdBytes() {
- return false
- }
- if !json.Valid(payload) {
- return false
- }
- i := skipOpenAIWSJSONSpace(payload, 0)
- if i >= len(payload) || payload[i] != '{' {
- return false
- }
- i++
- eventType := "response.create"
- previousResponseID := ""
- typeSeen, previousResponseIDSeen := false, false
- for {
- i = skipOpenAIWSJSONSpace(payload, i)
- if payload[i] == '}' {
- break
- }
- keyStart := i
- keyEnd := scanOpenAIWSJSONString(payload, keyStart)
- i = skipOpenAIWSJSONSpace(payload, keyEnd)
- i++ // json.Valid guarantees the colon.
- i = skipOpenAIWSJSONSpace(payload, i)
- valueStart := i
- i = skipOpenAIWSJSONValue(payload, i)
- key := ""
- // A critical key is at most 20 decoded bytes. The generous encoded bound
- // covers escaped spellings without allocating attacker-sized key strings.
- if keyEnd-keyStart <= 128 {
- _ = json.Unmarshal(payload[keyStart:keyEnd], &key)
- }
- switch key {
- case "type":
- if typeSeen {
- return false
- }
- typeSeen = true
- var value *string
- if err := json.Unmarshal(payload[valueStart:i], &value); err != nil {
- return false
- }
- if value == nil || strings.TrimSpace(*value) == "" {
- eventType = "response.create"
- } else {
- eventType = strings.TrimSpace(*value)
- }
- case "previous_response_id":
- if previousResponseIDSeen {
- return false
- }
- previousResponseIDSeen = true
- var value *string
- if err := json.Unmarshal(payload[valueStart:i], &value); err != nil {
- return false
- }
- if value != nil {
- previousResponseID = strings.TrimSpace(*value)
- }
- }
- i = skipOpenAIWSJSONSpace(payload, i)
- if payload[i] == ',' {
- i++
- }
- }
- return eventType == "response.create" && previousResponseID == ""
- }
- func skipOpenAIWSJSONSpace(payload []byte, i int) int {
- for i < len(payload) {
- switch payload[i] {
- case ' ', '\t', '\r', '\n':
- i++
- default:
- return i
- }
- }
- return i
- }
- func scanOpenAIWSJSONString(payload []byte, i int) int {
- for i++; i < len(payload); i++ {
- switch payload[i] {
- case '\\':
- i++
- case '"':
- return i + 1
- }
- }
- return len(payload)
- }
- func skipOpenAIWSJSONValue(payload []byte, i int) int {
- if payload[i] == '"' {
- return scanOpenAIWSJSONString(payload, i)
- }
- if payload[i] != '{' && payload[i] != '[' {
- for i < len(payload) && payload[i] != ',' && payload[i] != '}' {
- i++
- }
- return i
- }
- depth := 0
- for ; i < len(payload); i++ {
- switch payload[i] {
- case '"':
- i = scanOpenAIWSJSONString(payload, i) - 1
- case '{', '[':
- depth++
- case '}', ']':
- depth--
- if depth == 0 {
- return i + 1
- }
- }
- }
- return len(payload)
- }
- func prepareOpenAIWSHTTPBridgeBody(account *Account, payload []byte) ([]byte, error) {
- var body map[string]any
- if err := decodeOpenAIJSONUseNumber(payload, &body); err != nil {
- return nil, err
- }
- if body == nil {
- return nil, errors.New("response.create payload must be a JSON object")
- }
- delete(body, "type")
- delete(body, "generate")
- delete(body, "previous_response_id")
- deleteOpenAIResponsesNoneReasoningEffortFromObject(account, body)
- body["stream"] = true
- return json.Marshal(body)
- }
- type openAIWSToolCallReplayCollector struct {
- items []json.RawMessage
- seen map[string]struct{}
- allItems []json.RawMessage
- allSeen map[string]struct{}
- }
- func (c *openAIWSToolCallReplayCollector) AddEvent(eventType string, message []byte) {
- switch strings.TrimSpace(eventType) {
- case "response.output_item.done":
- item := gjson.GetBytes(message, "item")
- c.addAllItem(item)
- c.addItem(item)
- case "response.completed", "response.done":
- output := gjson.GetBytes(message, "response.output")
- if !output.IsArray() {
- return
- }
- for _, item := range output.Array() {
- c.addAllItem(item)
- c.addItem(item)
- }
- }
- }
- // Items/AllItems 返回浅拷贝头数组;正文由 collector 独立分配且此后不可变,
- // 调用方按 replay 所有权不变式共享持有。
- func (c *openAIWSToolCallReplayCollector) Items() []json.RawMessage {
- return slices.Clone(c.items)
- }
- func (c *openAIWSToolCallReplayCollector) AllItems() []json.RawMessage {
- return slices.Clone(c.allItems)
- }
- func (c *openAIWSToolCallReplayCollector) addAllItem(item gjson.Result) {
- if !item.Exists() || item.Type != gjson.JSON {
- return
- }
- raw := strings.TrimSpace(item.Raw)
- if raw == "" || !strings.HasPrefix(raw, "{") || strings.TrimSpace(item.Get("type").String()) == "" {
- return
- }
- key := strings.TrimSpace(item.Get("id").String())
- if key == "" {
- key = strings.TrimSpace(item.Get("call_id").String())
- }
- if key == "" {
- key = raw
- }
- if c.allSeen == nil {
- c.allSeen = make(map[string]struct{})
- }
- if _, ok := c.allSeen[key]; ok {
- return
- }
- c.allSeen[key] = struct{}{}
- c.allItems = append(c.allItems, json.RawMessage(raw))
- }
- func (c *openAIWSToolCallReplayCollector) addItem(item gjson.Result) {
- if !item.Exists() || item.Type != gjson.JSON {
- return
- }
- raw := strings.TrimSpace(item.Raw)
- if raw == "" || !strings.HasPrefix(raw, "{") {
- return
- }
- if !isCodexToolCallContextItemType(item.Get("type").String()) {
- return
- }
- key := strings.TrimSpace(item.Get("id").String())
- if key == "" {
- key = strings.TrimSpace(item.Get("call_id").String())
- }
- if key == "" {
- key = raw
- }
- if c.seen == nil {
- c.seen = make(map[string]struct{})
- }
- if _, ok := c.seen[key]; ok {
- return
- }
- c.seen[key] = struct{}{}
- c.items = append(c.items, json.RawMessage(raw))
- }
- func buildOpenAIWSHTTPBridgeErrorEvent(statusCode int, message string) []byte {
- message = strings.TrimSpace(message)
- if message == "" {
- message = http.StatusText(statusCode)
- }
- if message == "" {
- message = "upstream request failed"
- }
- event := map[string]any{
- "type": "error",
- "status": statusCode,
- "error": map[string]any{
- "type": "upstream_error",
- "message": message,
- },
- }
- body, err := json.Marshal(event)
- if err != nil {
- return []byte(`{"type":"error","error":{"type":"upstream_error","message":"upstream request failed"}}`)
- }
- return body
- }
- func buildOpenAIWSHTTPBridgeFailedEvent(responseID, model string, source []byte, fallbackMessage string) []byte {
- errorType := strings.TrimSpace(gjson.GetBytes(source, "error.type").String())
- if errorType == "" {
- errorType = strings.TrimSpace(gjson.GetBytes(source, "response.error.type").String())
- }
- code := strings.TrimSpace(gjson.GetBytes(source, "error.code").String())
- if code == "" {
- code = strings.TrimSpace(gjson.GetBytes(source, "response.error.code").String())
- }
- if code == "" {
- code = "upstream_error"
- }
- message := extractOpenAISSEErrorMessage(source)
- if message == "" {
- message = strings.TrimSpace(fallbackMessage)
- }
- if message == "" {
- message = "Upstream response failed"
- }
- errorBody := map[string]any{"code": code, "message": message}
- if errorType != "" {
- errorBody["type"] = errorType
- }
- response := map[string]any{
- "id": responseID, "object": "response", "status": "failed",
- "output": []any{}, "error": errorBody,
- }
- if model = strings.TrimSpace(model); model != "" {
- response["model"] = model
- }
- body, err := json.Marshal(map[string]any{"type": "response.failed", "response": response})
- if err != nil {
- return []byte(`{"type":"response.failed","response":{"status":"failed","output":[],"error":{"code":"upstream_error","message":"Upstream response failed"}}}`)
- }
- return body
- }
- func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
- ctx context.Context,
- c *gin.Context,
- account *Account,
- token string,
- payload []byte,
- payloadBytes int,
- originalModel string,
- imageBillingModel string,
- imageSizeTier string,
- imageInputSize string,
- grokCacheIdentity string,
- turn int,
- writeClientMessage func([]byte) error,
- ) (*OpenAIForwardResult, error) {
- if s == nil {
- return nil, errors.New("service is nil")
- }
- if s.httpUpstream == nil {
- return nil, errors.New("openai http upstream is nil")
- }
- if account == nil {
- return nil, errors.New("account is nil")
- }
- if writeClientMessage == nil {
- return nil, errors.New("client websocket writer is nil")
- }
- responseModelObserver := &upstreamResponseModelObserver{}
- body, err := prepareOpenAIWSHTTPBridgeBody(account, payload)
- if err != nil {
- return nil, fmt.Errorf("prepare http bridge body: %w", err)
- }
- grokIntentSourceBody := append([]byte(nil), body...)
- _, grokExplicitToolsField := openAIWSHTTPBridgeRawField(grokIntentSourceBody, "tools")
- grokExplicitToolIntent := account.Platform == PlatformGrok && hasGrokResponsesToolIntent(grokIntentSourceBody)
- var clientToolMapping apicompat.ResponsesClientToolMapping
- functionToolUpstream := (account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey) || account.Platform == PlatformGrok
- if functionToolUpstream {
- if account.Platform == PlatformGrok {
- body, err = sanitizeGrokResponsesInput(body)
- if err != nil {
- return nil, fmt.Errorf("sanitize Grok WS HTTP bridge input: %w", err)
- }
- }
- inheritedState, _ := openAIWSHTTPBridgeToolStateFromContext(c)
- inheritedLoweredTools := decodeOpenAIWSHTTPBridgeLoweredTools(inheritedState.LoweredTools)
- body, clientToolMapping, err = adaptResponsesClientToolsForFunctionUpstreamWithMapping(
- body,
- openAIWSHTTPBridgeToolUpstreamName(account),
- inheritedState.ClientMapping,
- inheritedLoweredTools,
- )
- if err != nil {
- return nil, fmt.Errorf("adapt %s client tools: %w", openAIWSHTTPBridgeToolUpstreamName(account), err)
- }
- if account.Platform == PlatformGrok && !grokExplicitToolsField && !grokExplicitToolIntent && len(inheritedLoweredTools) > 0 && hasGrokResponsesToolIntent(body) {
- // This continuation omitted tools, so the pre-adapter source cannot
- // represent the effective inherited declarations. Cache routing must
- // see the rehydrated tool intent or it will replace client functions
- // with the native-search tool-free route. Explicit current-turn tool
- // intent still uses the original pre-sanitization source above.
- grokIntentSourceBody = append(grokIntentSourceBody[:0], body...)
- }
- loweredTools := inheritedState.LoweredTools
- if currentTools, present := openAIWSHTTPBridgeRawField(body, "tools"); present {
- loweredTools = currentTools
- }
- setOpenAIWSHTTPBridgeToolState(c, openAIWSHTTPBridgeToolState{
- ClientMapping: clientToolMapping,
- LoweredTools: loweredTools,
- })
- }
- if account.Platform != PlatformGrok && isOpenAIResponsesLiteWebSocketPayload(payload) {
- liteBody, liteChanged, liteErr := normalizeOpenAIResponsesLitePayloadForAccount(body, account)
- if liteErr != nil {
- return nil, fmt.Errorf("normalize responses Lite payload: %w", liteErr)
- }
- if liteChanged {
- body = liteBody
- }
- }
- buildUpstreamRequest := func(requestBody []byte) (*http.Request, error) {
- upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
- defer releaseUpstreamCtx()
- var upstreamReq *http.Request
- var buildErr error
- if account.Platform == PlatformGrok {
- upstreamReq, buildErr = buildGrokResponsesRequest(upstreamCtx, c, account, requestBody, token, grokCacheIdentity, s.cfg, s.settingService)
- } else {
- upstreamReq, buildErr = s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, requestBody, token)
- }
- if buildErr != nil {
- return nil, buildErr
- }
- if account.Platform != PlatformGrok && isOpenAIResponsesLiteWebSocketPayload(payload) {
- upstreamReq.Header.Set(responsesLiteHeader, "true")
- }
- return upstreamReq, nil
- }
- if account.Platform == PlatformGrok {
- upstreamModel := resolveGrokWSUpstreamModel(account, body, originalModel)
- body, err = patchGrokResponsesBody(body, upstreamModel)
- if err != nil {
- return nil, err
- }
- grokMixedCacheIntentBody := append([]byte(nil), body...)
- body, err = applyGrokResponsesCacheIdentity(body, grokIntentSourceBody, grokCacheIdentity, account.IsGrokOAuth())
- if err != nil {
- return nil, fmt.Errorf("apply grok prompt cache identity: %w", err)
- }
- body, err = applyGrokFreeRequestToolCacheRoute(c, body, grokMixedCacheIntentBody, account, grokCacheIdentity)
- if err != nil {
- return nil, fmt.Errorf("apply grok Free function-tool cache route: %w", err)
- }
- }
- actualModel := strings.TrimSpace(gjson.GetBytes(body, "model").String())
- if actualModel == "" {
- actualModel = canonicalOpenAIAccountSchedulingModel(account, originalModel)
- }
- SetOpsUpstreamModel(c, actualModel)
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- if c != nil {
- c.Set("openai_passthrough", true)
- c.Set("openai_ws_http_bridge", true)
- }
- turnStart := time.Now()
- rejectedFieldRetryState := newOpenAIResponsesRejectedFieldRetryState(body)
- var resp *http.Response
- for {
- upstreamReq, buildErr := buildUpstreamRequest(body)
- if buildErr != nil {
- return nil, buildErr
- }
- resp, err = s.doOpenAIUpstream(upstreamReq, proxyURL, account)
- if err != nil {
- if turn == 1 {
- return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true)
- }
- safeErr := sanitizeUpstreamErrorMessage(err.Error())
- clientError := buildOpenAIWSHTTPBridgeErrorEvent(http.StatusBadGateway, "Upstream request failed")
- if writeErr := writeClientMessage(clientError); writeErr == nil {
- markOpenAIWSClientVisibleFailure(c, "error", clientError)
- }
- return nil, fmt.Errorf("upstream http bridge request failed: %s", safeErr)
- }
- if resp.StatusCode < 400 {
- break
- }
- respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes))
- _ = resp.Body.Close()
- markOpenAICyberPolicyEvent(c, respBody, resp.StatusCode, nil)
- if resp.StatusCode == http.StatusBadRequest &&
- extractUpstreamErrorCode(respBody) == openAIWSFallbackReasonInvalidEncryptedContent {
- s.markOpenAIWSInvalidEncryptedContentLineageFromPayload(
- c, body, "ingress_ws_http_bridge_invalid_encrypted_lineage_mark", account.ID, turn,
- )
- }
- retryBody, retryReason, changed, retryErr := normalizeOpenAIResponsesRejectedFieldRetryBody(resp.StatusCode, body, respBody)
- if retryErr != nil {
- return nil, fmt.Errorf("normalize websocket http bridge rejected field retry: %w", retryErr)
- }
- if changed && rejectedFieldRetryState.Allow(retryBody) {
- logOpenAIWSModeInfo(
- "ingress_ws_http_bridge_rejected_field_retry account_id=%d turn=%d reason=%s",
- account.ID,
- turn,
- truncateOpenAIWSLogValue(retryReason, openAIWSLogValueMaxLen),
- )
- body = retryBody
- payloadBytes = len(body)
- continue
- }
- upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
- if upstreamMsg == "" {
- upstreamMsg = http.StatusText(resp.StatusCode)
- }
- shouldFailover := s.shouldFailoverOpenAIUpstreamResponse(account, resp.StatusCode, upstreamMsg, respBody)
- if account.Platform == PlatformGrok {
- shouldFailover = s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody)
- s.handleGrokAccountUpstreamError(withGrokTeamRateLimitModel(ctx, resolveGrokWSUpstreamModel(account, body, originalModel)), account, resp.StatusCode, resp.Header, respBody)
- if shouldFailover && (turn == 1 || resp.StatusCode == http.StatusTooManyRequests) {
- return nil, newOpenAIUpstreamFailoverError(resp.StatusCode, resp.Header, respBody, upstreamMsg, false)
- }
- } else if shouldFailover && (turn == 1 || resp.StatusCode == http.StatusTooManyRequests) {
- return nil, s.handleFailoverErrorResponsePassthrough(ctx, resp, c, account, body, respBody)
- }
- if account.Platform != PlatformGrok && (shouldFailover || shouldCooldownOpenAITransientUpstreamError(resp.StatusCode, respBody)) {
- s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, actualModel)
- }
- clientError := buildOpenAIWSHTTPBridgeErrorEvent(resp.StatusCode, upstreamMsg)
- if writeErr := writeClientMessage(clientError); writeErr == nil {
- markOpenAIWSClientVisibleFailure(c, "error", clientError)
- }
- return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg)
- }
- defer func() { _ = resp.Body.Close() }()
- stopCancelBody := context.AfterFunc(ctx, func() { _ = resp.Body.Close() })
- defer stopCancelBody()
- if account.Platform == PlatformGrok {
- s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, resolveGrokWSUpstreamModel(account, body, originalModel)), account, resp.Header, resp.StatusCode)
- }
- responseID := ""
- usage := OpenAIUsage{}
- imageCounter := newOpenAIImageOutputCounter()
- var firstTokenMs *int
- reqStream := openAIWSPayloadBoolFromRaw(body, "stream", true)
- eventCount := 0
- tokenEventCount := 0
- terminalEventCount := 0
- replayCollector := &openAIWSToolCallReplayCollector{}
- firstEventType := ""
- lastEventType := ""
- upstreamTerminalEvent := ""
- sawDone := false
- wroteDownstream := false
- pendingClientMessages := make([][]byte, 0, 4)
- pendingClientMessageBytes := int64(0)
- capacityFailoverSuppressedLogged := false
- clientDisconnected := false
- officialOpenAIResponses := account != nil && account.Platform == PlatformOpenAI
- bareErrorPending := false
- var bareErrorPayload []byte
- bareErrorMessage := ""
- failureAccountSideEffectsApplied := false
- mappedModel := actualModel
- needModelReplace := false
- var mappedModelBytes []byte
- if originalModel != "" {
- needModelReplace = mappedModel != "" && mappedModel != originalModel
- if needModelReplace {
- mappedModelBytes = []byte(mappedModel)
- }
- }
- resultWithUsage := func() *OpenAIForwardResult {
- imageCount := imageCounter.Count()
- result := &OpenAIForwardResult{
- RequestID: responseID,
- Usage: usage,
- Model: originalModel,
- UpstreamModel: mappedModel,
- UpstreamResponseModel: responseModelObserver.Model(),
- UpstreamResponseModelConflict: responseModelObserver.Conflict(),
- UpstreamResponseServiceTier: responseModelObserver.ServiceTier(),
- ServiceTier: resolvedOpenAIUpstreamServiceTierFromObserver(responseModelObserver, extractOpenAIServiceTierFromBody(body)),
- ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, mappedModel, originalModel), body, mappedModel),
- RequestedReasoningEffort: CanonicalRequestedReasoningEffort(body, originalModel, mappedModel),
- Stream: reqStream,
- OpenAIWSMode: true,
- UpstreamTerminalEvent: upstreamTerminalEvent,
- ResponseHeaders: cloneHeader(resp.Header),
- Duration: time.Since(turnStart),
- FirstTokenMs: firstTokenMs,
- }
- if replayInput := replayCollector.Items(); len(replayInput) > 0 {
- result.wsReplayInput = replayInput
- result.wsReplayInputExists = true
- }
- result.wsAccountFailoverReplayInput = replayCollector.AllItems()
- if imageCount > 0 {
- result.ImageCount = imageCount
- result.ImageSize = imageSizeTier
- result.ImageInputSize = imageInputSize
- result.ImageOutputSizes = imageCounter.Sizes()
- result.BillingModel = imageBillingModel
- }
- return result
- }
- maxLineSize := defaultMaxLineSize
- if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
- maxLineSize = s.cfg.Gateway.MaxLineSize
- }
- if hasResponsesClientToolMapping(clientToolMapping) {
- resp.Body = newResponsesClientToolStreamBody(resp.Body, clientToolMapping, maxLineSize)
- }
- scanner := bufio.NewScanner(resp.Body)
- scanBuf := getSSEScannerBuf64K()
- scanner.Buffer(scanBuf[:0], maxLineSize)
- defer putSSEScannerBuf64K(scanBuf)
- pendingSSEEventType := ""
- finalizeBareError := func() error {
- if !bareErrorPending {
- return nil
- }
- if !failureAccountSideEffectsApplied {
- failureAccountSideEffectsApplied = s.handleOpenAIWSFailureAccountSideEffects(ctx, account, mappedModel, resp.Header, bareErrorPayload)
- }
- upstreamTerminalEvent = "response.failed"
- if clientDisconnected {
- return nil
- }
- clientMessage := buildOpenAIWSHTTPBridgeFailedEvent(responseID, originalModel, bareErrorPayload, bareErrorMessage)
- if rewritten, changed := sanitizeOpenAICapacityShedErrorCodeForClient(clientMessage); changed {
- clientMessage = rewritten
- }
- messages := append(pendingClientMessages, clientMessage)
- pendingClientMessages = nil
- pendingClientMessageBytes = 0
- for _, message := range messages {
- if err := writeClientMessage(message); err != nil {
- if isOpenAIWSClientDisconnectError(err) {
- clientDisconnected = true
- return nil
- }
- return fmt.Errorf("write synthesized websocket response.failed: %w", err)
- }
- wroteDownstream = true
- }
- markOpenAIWSClientVisibleFailure(c, "response.failed", clientMessage)
- return nil
- }
- for scanner.Scan() {
- line := scanner.Text()
- if eventType, ok := extractOpenAISSEEventLine(line); ok {
- pendingSSEEventType = eventType
- continue
- }
- if strings.TrimSpace(line) == "" {
- pendingSSEEventType = ""
- continue
- }
- data, ok := extractOpenAISSEDataLine(line)
- if !ok {
- continue
- }
- trimmedData := strings.TrimSpace(data)
- if trimmedData == "" {
- continue
- }
- if trimmedData == "[DONE]" {
- sawDone = true
- continue
- }
- upstreamMessage := []byte(openAICompatPayloadWithEventType(trimmedData, pendingSSEEventType))
- if normalized, changed := normalizeCompletedImageGenerationStatus(upstreamMessage); changed {
- upstreamMessage = normalized
- }
- eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage)
- responseModelObserver.ObserveOpenAI(upstreamMessage, eventType)
- if responseID == "" && eventResponseID != "" {
- responseID = eventResponseID
- }
- if eventType != "" {
- eventCount++
- if firstEventType == "" {
- firstEventType = eventType
- }
- lastEventType = eventType
- }
- if isOpenAIWSTokenEvent(eventType) {
- tokenEventCount++
- if firstTokenMs == nil {
- ms := int(time.Since(turnStart).Milliseconds())
- firstTokenMs = &ms
- }
- }
- if openAIWSMessageShouldParseUsage(eventType, upstreamMessage) {
- parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage)
- }
- if eventType == "error" || eventType == "response.failed" {
- markOpenAICyberPolicyEvent(c, upstreamMessage, http.StatusOK, &usage)
- }
- imageCounter.AddSSEData(upstreamMessage)
- if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && strings.Contains(trimmedData, mappedModel) {
- upstreamMessage = replaceOpenAIWSMessageModel(upstreamMessage, mappedModel, originalModel)
- }
- if s.toolCorrector != nil && openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(upstreamMessage) {
- if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(upstreamMessage); changed {
- upstreamMessage = corrected
- }
- }
- replayCollector.AddEvent(eventType, upstreamMessage)
- var upstreamEventErr error
- if officialOpenAIResponses && bareErrorPending && (eventType == "response.completed" || eventType == "response.done") {
- // Some upstreams emit a recoverable bare error before the authoritative
- // successful terminal. Do not replace that terminal with a synthetic
- // failure or retain side effects from the superseded error.
- bareErrorPending = false
- bareErrorPayload = nil
- bareErrorMessage = ""
- }
- suppressClientMessage := officialOpenAIResponses && bareErrorPending && eventType != "response.failed"
- if eventType == "error" || eventType == "response.failed" {
- errMessage := extractOpenAISSEErrorMessage(upstreamMessage)
- if errMessage == "" {
- errMessage = "upstream error event"
- }
- statusCode := openAIStreamFailureStatus(upstreamMessage, errMessage)
- shouldFailover := openAIStreamFailedEventShouldFailover(upstreamMessage, errMessage)
- if eventType == "error" {
- errCodeRaw, errTypeRaw, _ := parseOpenAIWSErrorEventFields(upstreamMessage)
- shouldFailover = openAIStreamErrorEventShouldFailover(upstreamMessage, errMessage)
- if account.Platform == PlatformGrok {
- statusCode = openAIWSErrorHTTPStatusFromRaw(errCodeRaw, errTypeRaw)
- }
- if reason, _ := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMessage); reason == openAIWSFallbackReasonInvalidEncryptedContent {
- s.markOpenAIWSInvalidEncryptedContentLineageFromPayload(
- c, body, "ingress_ws_http_bridge_invalid_encrypted_lineage_mark", account.ID, turn,
- )
- }
- }
- requestScopedCapacity := isOpenAIUpstreamCapacityShedEvent(upstreamMessage)
- if account.Platform == PlatformGrok && eventType == "error" {
- // SSE error events do not carry an HTTP status. The local status
- // mapper therefore defaults unknown xAI codes (for example
- // new_sensitive) to 502; classify the body as a request-scoped
- // 403 before applying status-based failover or account state.
- if isGrokContentPolicyRejection(http.StatusForbidden, upstreamMessage) {
- shouldFailover = false
- } else {
- shouldFailover = s.shouldFailoverGrokUpstreamError(statusCode, upstreamMessage)
- s.handleGrokAccountUpstreamError(ctx, account, statusCode, resp.Header, upstreamMessage)
- }
- }
- if !wroteDownstream && shouldFailover && (turn == 1 || statusCode == http.StatusTooManyRequests) {
- if account.Platform == PlatformGrok {
- return nil, newOpenAIUpstreamFailoverError(statusCode, resp.Header, upstreamMessage, errMessage, false)
- }
- return nil, s.newOpenAIStreamFailoverErrorWithModel(c, account, true, resp.Header.Get("x-request-id"), upstreamMessage, errMessage, mappedModel, resp.Header)
- }
- if account.Platform != PlatformGrok && !failureAccountSideEffectsApplied {
- if eventType == "response.failed" || (!officialOpenAIResponses && shouldFailover && !requestScopedCapacity) {
- failureAccountSideEffectsApplied = s.handleOpenAIWSFailureAccountSideEffects(ctx, account, mappedModel, resp.Header, upstreamMessage)
- }
- }
- if wroteDownstream && requestScopedCapacity && !capacityFailoverSuppressedLogged {
- logOpenAICapacityFailoverSuppressed(ctx, account, "ws_http_bridge", resp.Header.Get("x-request-id"), eventType)
- capacityFailoverSuppressedLogged = true
- }
- if eventType == "error" && !officialOpenAIResponses {
- upstreamEventErr = errors.New(errMessage)
- } else if eventType == "error" {
- bareErrorPending = true
- bareErrorPayload = append(bareErrorPayload[:0], upstreamMessage...)
- bareErrorMessage = errMessage
- suppressClientMessage = true
- } else {
- bareErrorPending = false
- }
- }
- // 客户端写出副本改写容量降载码:Codex 对 error/response.failed 中的
- // server_is_overloaded / slow_down 判致命并终止会话,改写后走客户端内置
- // 重试。账号状态与终止事件判定(下方 handleOpenAIWSTerminalTransientFailure)
- // 仍使用未改写的 upstreamMessage。
- clientMessage := upstreamMessage
- if eventType == "error" || eventType == "response.failed" {
- if rewritten, changed := sanitizeOpenAICapacityShedErrorCodeForClient(clientMessage); changed {
- clientMessage = rewritten
- }
- }
- if !clientDisconnected && !suppressClientMessage {
- stageBeforeSemanticOutput := turn == 1 && account.Platform == PlatformOpenAI && !wroteDownstream
- commitStagedMessages := !stageBeforeSemanticOutput ||
- openAIStreamDataStartsClientOutput(string(clientMessage), eventType) ||
- isOpenAIWSTerminalEvent(eventType)
- if stageBeforeSemanticOutput && !commitStagedMessages {
- if pendingClientMessageBytes+int64(len(clientMessage)) > openAIFirstOutputStageMaxBytes {
- return nil, s.newOpenAIStreamFailoverError(
- c,
- account,
- true,
- resp.Header.Get("x-request-id"),
- nil,
- "OpenAI WS HTTP bridge first-output staging limit exceeded",
- resp.Header,
- )
- }
- pendingClientMessages = append(pendingClientMessages, append([]byte(nil), clientMessage...))
- pendingClientMessageBytes += int64(len(clientMessage))
- } else {
- messages := append(pendingClientMessages, clientMessage)
- pendingClientMessages = nil
- pendingClientMessageBytes = 0
- for _, message := range messages {
- if err := writeClientMessage(message); err != nil {
- if isOpenAIWSClientDisconnectError(err) {
- clientDisconnected = true
- closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err)
- logOpenAIWSModeInfo(
- "ingress_ws_http_bridge_client_disconnected_drain account_id=%d turn=%d close_status=%s close_reason=%s",
- account.ID,
- turn,
- closeStatus,
- truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen),
- )
- break
- }
- return nil, wrapOpenAIWSIngressTurnError(
- "write_client",
- fmt.Errorf("write client websocket event: %w", err),
- wroteDownstream,
- )
- }
- wroteDownstream = true
- }
- }
- }
- if !clientDisconnected && !suppressClientMessage {
- markOpenAIWSClientVisibleFailure(c, eventType, upstreamMessage)
- }
- if upstreamEventErr != nil {
- return resultWithUsage(), upstreamEventErr
- }
- if isOpenAIWSTerminalEvent(eventType) && !bareErrorPending {
- if eventType == "response.failed" {
- upstreamTerminalEvent = "response.failed"
- } else {
- upstreamTerminalEvent = s.handleOpenAIWSTerminalTransientFailure(ctx, account, mappedModel, resp.Header, upstreamMessage)
- }
- terminalEventCount++
- firstTokenMsValue := -1
- if firstTokenMs != nil {
- firstTokenMsValue = *firstTokenMs
- }
- logOpenAIWSModeInfo(
- "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",
- account.ID,
- turn,
- truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen),
- payloadBytes,
- time.Since(turnStart).Milliseconds(),
- eventCount,
- tokenEventCount,
- terminalEventCount,
- truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen),
- truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen),
- firstTokenMsValue,
- clientDisconnected,
- )
- return resultWithUsage(), nil
- }
- }
- if bareErrorPending {
- if finalizeErr := finalizeBareError(); finalizeErr != nil {
- return resultWithUsage(), finalizeErr
- }
- if scanErr := scanner.Err(); scanErr != nil {
- return resultWithUsage(), fmt.Errorf("read upstream http bridge stream after error event: %w", scanErr)
- }
- return resultWithUsage(), errors.New(bareErrorMessage)
- }
- if err := scanner.Err(); err != nil {
- streamErr := fmt.Errorf("read upstream http bridge stream: %w", err)
- if turn == 1 && !wroteDownstream {
- return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, streamErr, true)
- }
- return resultWithUsage(), streamErr
- }
- terminalErr := errors.New("upstream http bridge stream ended before terminal event")
- if sawDone {
- terminalErr = errors.New("upstream http bridge stream sent [DONE] before terminal event")
- }
- if turn == 1 && !wroteDownstream {
- return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, terminalErr, true)
- }
- return resultWithUsage(), terminalErr
- }
- func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, seedPayload, currentPayload []byte, originalModel string) (string, error) {
- body, err := prepareOpenAIWSHTTPBridgeBody(account, seedPayload)
- if err != nil {
- return "", err
- }
- upstreamModel := resolveGrokWSUpstreamModel(account, currentPayload, originalModel)
- body, err = patchGrokResponsesBody(body, upstreamModel)
- if err != nil {
- return "", err
- }
- return resolveGrokCacheIdentity(c, body, "", upstreamModel), nil
- }
- func resolveGrokWSUpstreamModel(account *Account, body []byte, originalModel string) string {
- upstreamModel := strings.TrimSpace(gjson.GetBytes(body, "model").String())
- originalModel = strings.TrimSpace(originalModel)
- // Shared ingress has already applied channel and account mappings when the
- // body model differs from the client-facing model. Only resolve from the
- // original model when the body still carries that original value.
- if account != nil && originalModel != "" && (upstreamModel == "" || upstreamModel == originalModel) {
- if mappedModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)); mappedModel != "" {
- upstreamModel = mappedModel
- }
- }
- if upstreamModel == "" {
- upstreamModel = grokDefaultResponsesModel
- }
- return upstreamModel
- }
|