diff --git a/backend/internal/service/openai_ws_forwarder.go b/backend/internal/service/openai_ws_forwarder.go index bbca9776ab..e0bd769c3c 100644 --- a/backend/internal/service/openai_ws_forwarder.go +++ b/backend/internal/service/openai_ws_forwarder.go @@ -1,27 +1,15 @@ package service import ( - "bytes" "context" - "encoding/json" "errors" "fmt" - "io" "math/rand" - "net" - "net/http" - "net/url" - "sort" "strings" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "github.com/Wei-Shaw/sub2api/internal/pkg/openai" - "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" coderws "github.com/coder/websocket" - "github.com/gin-gonic/gin" - "github.com/tidwall/gjson" - "github.com/tidwall/sjson" "go.uber.org/zap" ) @@ -227,676 +215,6 @@ type OpenAIWSIngressHooks struct { AfterTurn func(turn int, result *OpenAIForwardResult, turnErr error) } -func normalizeOpenAIWSLogValue(value string) string { - trimmed := strings.TrimSpace(value) - if trimmed == "" { - return "-" - } - return openAIWSLogValueReplacer.Replace(trimmed) -} - -func truncateOpenAIWSLogValue(value string, maxLen int) string { - normalized := normalizeOpenAIWSLogValue(value) - if normalized == "-" || maxLen <= 0 { - return normalized - } - if len(normalized) <= maxLen { - return normalized - } - return normalized[:maxLen] + "..." -} - -func openAIWSHeaderValueForLog(headers http.Header, key string) string { - if headers == nil { - return "-" - } - return truncateOpenAIWSLogValue(headers.Get(key), openAIWSHeaderValueMaxLen) -} - -func hasOpenAIWSHeader(headers http.Header, key string) bool { - if headers == nil { - return false - } - return strings.TrimSpace(headers.Get(key)) != "" -} - -type openAIWSSessionHeaderResolution struct { - SessionID string - ConversationID string - SessionSource string - ConversationSource string -} - -func resolveOpenAIWSSessionHeaders(c *gin.Context, promptCacheKey string) openAIWSSessionHeaderResolution { - resolution := openAIWSSessionHeaderResolution{ - SessionSource: "none", - ConversationSource: "none", - } - if c != nil && c.Request != nil { - if sessionID := strings.TrimSpace(c.Request.Header.Get("session_id")); sessionID != "" { - resolution.SessionID = sessionID - resolution.SessionSource = "header_session_id" - } - if conversationID := strings.TrimSpace(c.Request.Header.Get("conversation_id")); conversationID != "" { - resolution.ConversationID = conversationID - resolution.ConversationSource = "header_conversation_id" - if resolution.SessionID == "" { - resolution.SessionID = conversationID - resolution.SessionSource = "header_conversation_id" - } - } - } - - cacheKey := strings.TrimSpace(promptCacheKey) - if cacheKey != "" { - if resolution.SessionID == "" { - resolution.SessionID = cacheKey - resolution.SessionSource = "prompt_cache_key" - } - } - return resolution -} - -func shouldLogOpenAIWSEvent(idx int, eventType string) bool { - if idx <= openAIWSEventLogHeadLimit { - return true - } - if openAIWSEventLogEveryN > 0 && idx%openAIWSEventLogEveryN == 0 { - return true - } - if eventType == "error" || isOpenAIWSTerminalEvent(eventType) { - return true - } - return false -} - -func shouldLogOpenAIWSBufferedEvent(idx int) bool { - if idx <= openAIWSBufferLogHeadLimit { - return true - } - if openAIWSBufferLogEveryN > 0 && idx%openAIWSBufferLogEveryN == 0 { - return true - } - return false -} - -func openAIWSEventMayContainModel(eventType string) bool { - switch eventType { - case "response.created", - "response.in_progress", - "response.completed", - "response.done", - "response.failed", - "response.incomplete", - "response.cancelled", - "response.canceled": - return true - default: - trimmed := strings.TrimSpace(eventType) - if trimmed == eventType { - return false - } - switch trimmed { - case "response.created", - "response.in_progress", - "response.completed", - "response.done", - "response.failed", - "response.incomplete", - "response.cancelled", - "response.canceled": - return true - default: - return false - } - } -} - -func openAIWSEventMayContainToolCalls(eventType string) bool { - eventType = strings.TrimSpace(eventType) - if eventType == "" { - return false - } - if strings.Contains(eventType, "function_call") || strings.Contains(eventType, "tool_call") { - return true - } - switch eventType { - case "response.output_item.added", "response.output_item.done", "response.completed", "response.done": - return true - default: - return false - } -} - -func openAIWSEventShouldParseUsage(eventType string) bool { - switch strings.TrimSpace(eventType) { - case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled": - return true - default: - return false - } -} - -func parseOpenAIWSEventEnvelope(message []byte) (eventType string, responseID string, response gjson.Result) { - if len(message) == 0 { - return "", "", gjson.Result{} - } - values := gjson.GetManyBytes(message, "type", "response.id", "id", "response") - eventType = strings.TrimSpace(values[0].String()) - if id := strings.TrimSpace(values[1].String()); id != "" { - responseID = id - } else { - responseID = strings.TrimSpace(values[2].String()) - } - return eventType, responseID, values[3] -} - -func openAIWSMessageLikelyContainsToolCalls(message []byte) bool { - if len(message) == 0 { - return false - } - return bytes.Contains(message, []byte(`"tool_calls"`)) || - bytes.Contains(message, []byte(`"tool_call"`)) || - bytes.Contains(message, []byte(`"function_call"`)) -} - -func parseOpenAIWSResponseUsageFromCompletedEvent(message []byte, usage *OpenAIUsage) { - if usage == nil || len(message) == 0 { - return - } - if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(message); ok { - *usage = parsedUsage - } -} - -func parseOpenAIWSErrorEventFields(message []byte) (code string, errType string, errMessage string) { - if len(message) == 0 { - return "", "", "" - } - values := gjson.GetManyBytes(message, "error.code", "error.type", "error.message") - return strings.TrimSpace(values[0].String()), strings.TrimSpace(values[1].String()), strings.TrimSpace(values[2].String()) -} - -func summarizeOpenAIWSErrorEventFieldsFromRaw(codeRaw, errTypeRaw, errMessageRaw string) (code string, errType string, errMessage string) { - code = truncateOpenAIWSLogValue(codeRaw, openAIWSLogValueMaxLen) - errType = truncateOpenAIWSLogValue(errTypeRaw, openAIWSLogValueMaxLen) - errMessage = truncateOpenAIWSLogValue(errMessageRaw, openAIWSLogValueMaxLen) - return code, errType, errMessage -} - -func summarizeOpenAIWSErrorEventFields(message []byte) (code string, errType string, errMessage string) { - if len(message) == 0 { - return "-", "-", "-" - } - return summarizeOpenAIWSErrorEventFieldsFromRaw(parseOpenAIWSErrorEventFields(message)) -} - -func summarizeOpenAIWSPayloadKeySizes(payload map[string]any, topN int) string { - if len(payload) == 0 { - return "-" - } - type keySize struct { - Key string - Size int - } - sizes := make([]keySize, 0, len(payload)) - for key, value := range payload { - size := estimateOpenAIWSPayloadValueSize(value, openAIWSPayloadSizeEstimateDepth) - sizes = append(sizes, keySize{Key: key, Size: size}) - } - sort.Slice(sizes, func(i, j int) bool { - if sizes[i].Size == sizes[j].Size { - return sizes[i].Key < sizes[j].Key - } - return sizes[i].Size > sizes[j].Size - }) - - if topN <= 0 || topN > len(sizes) { - topN = len(sizes) - } - parts := make([]string, 0, topN) - for idx := 0; idx < topN; idx++ { - item := sizes[idx] - parts = append(parts, fmt.Sprintf("%s:%d", item.Key, item.Size)) - } - return strings.Join(parts, ",") -} - -func estimateOpenAIWSPayloadValueSize(value any, depth int) int { - if depth <= 0 { - return -1 - } - switch v := value.(type) { - case nil: - return 0 - case string: - return len(v) - case []byte: - return len(v) - case bool: - return 1 - case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64: - return 8 - case float32, float64: - return 8 - case map[string]any: - if len(v) == 0 { - return 2 - } - total := 2 - count := 0 - for key, item := range v { - count++ - if count > openAIWSPayloadSizeEstimateMaxItems { - return -1 - } - itemSize := estimateOpenAIWSPayloadValueSize(item, depth-1) - if itemSize < 0 { - return -1 - } - total += len(key) + itemSize + 3 - if total > openAIWSPayloadSizeEstimateMaxBytes { - return -1 - } - } - return total - case []any: - if len(v) == 0 { - return 2 - } - total := 2 - limit := len(v) - if limit > openAIWSPayloadSizeEstimateMaxItems { - return -1 - } - for i := 0; i < limit; i++ { - itemSize := estimateOpenAIWSPayloadValueSize(v[i], depth-1) - if itemSize < 0 { - return -1 - } - total += itemSize + 1 - if total > openAIWSPayloadSizeEstimateMaxBytes { - return -1 - } - } - return total - default: - raw, err := json.Marshal(v) - if err != nil { - return -1 - } - if len(raw) > openAIWSPayloadSizeEstimateMaxBytes { - return -1 - } - return len(raw) - } -} - -func openAIWSPayloadString(payload map[string]any, key string) string { - if len(payload) == 0 { - return "" - } - raw, ok := payload[key] - if !ok { - return "" - } - switch v := raw.(type) { - case nil: - return "" - case string: - return strings.TrimSpace(v) - case []byte: - return strings.TrimSpace(string(v)) - default: - return "" - } -} - -func openAIWSPayloadStringFromRaw(payload []byte, key string) string { - if len(payload) == 0 || strings.TrimSpace(key) == "" { - return "" - } - return strings.TrimSpace(gjson.GetBytes(payload, key).String()) -} - -func openAIWSPayloadBoolFromRaw(payload []byte, key string, defaultValue bool) bool { - if len(payload) == 0 || strings.TrimSpace(key) == "" { - return defaultValue - } - value := gjson.GetBytes(payload, key) - if !value.Exists() { - return defaultValue - } - if value.Type != gjson.True && value.Type != gjson.False { - return defaultValue - } - return value.Bool() -} - -func openAIWSSessionHashesFromID(sessionID string) (string, string) { - return deriveOpenAISessionHashes(sessionID) -} - -func extractOpenAIWSImageURL(value any) string { - switch v := value.(type) { - case string: - return strings.TrimSpace(v) - case map[string]any: - if raw, ok := v["url"].(string); ok { - return strings.TrimSpace(raw) - } - } - return "" -} - -func summarizeOpenAIWSInput(input any) string { - items, ok := input.([]any) - if !ok || len(items) == 0 { - return "-" - } - - itemCount := len(items) - textChars := 0 - imageDataURLs := 0 - imageDataURLChars := 0 - imageRemoteURLs := 0 - - handleContentItem := func(contentItem map[string]any) { - contentType, _ := contentItem["type"].(string) - switch strings.TrimSpace(contentType) { - case "input_text", "output_text", "text": - if text, ok := contentItem["text"].(string); ok { - textChars += len(text) - } - case "input_image": - imageURL := extractOpenAIWSImageURL(contentItem["image_url"]) - if imageURL == "" { - return - } - if strings.HasPrefix(strings.ToLower(imageURL), "data:image/") { - imageDataURLs++ - imageDataURLChars += len(imageURL) - return - } - imageRemoteURLs++ - } - } - - handleInputItem := func(inputItem map[string]any) { - if content, ok := inputItem["content"].([]any); ok { - for _, rawContent := range content { - contentItem, ok := rawContent.(map[string]any) - if !ok { - continue - } - handleContentItem(contentItem) - } - return - } - - itemType, _ := inputItem["type"].(string) - switch strings.TrimSpace(itemType) { - case "input_text", "output_text", "text": - if text, ok := inputItem["text"].(string); ok { - textChars += len(text) - } - case "input_image": - imageURL := extractOpenAIWSImageURL(inputItem["image_url"]) - if imageURL == "" { - return - } - if strings.HasPrefix(strings.ToLower(imageURL), "data:image/") { - imageDataURLs++ - imageDataURLChars += len(imageURL) - return - } - imageRemoteURLs++ - } - } - - for _, rawItem := range items { - inputItem, ok := rawItem.(map[string]any) - if !ok { - continue - } - handleInputItem(inputItem) - } - - return fmt.Sprintf( - "items=%d,text_chars=%d,image_data_urls=%d,image_data_url_chars=%d,image_remote_urls=%d", - itemCount, - textChars, - imageDataURLs, - imageDataURLChars, - imageRemoteURLs, - ) -} - -func dropOpenAIWSPayloadKey(payload map[string]any, key string, removed *[]string) { - if len(payload) == 0 || strings.TrimSpace(key) == "" { - return - } - if _, exists := payload[key]; !exists { - return - } - delete(payload, key) - *removed = append(*removed, key) -} - -// applyOpenAIWSRetryPayloadStrategy 在 WS 连续失败时仅移除无语义字段， -// 避免重试成功却改变原始请求语义。 -// 注意：prompt_cache_key 不应在重试中移除；它常用于会话稳定标识（session_id 兜底）。 -func applyOpenAIWSRetryPayloadStrategy(payload map[string]any, attempt int) (strategy string, removedKeys []string) { - if len(payload) == 0 { - return "empty", nil - } - if attempt <= 1 { - return "full", nil - } - - removed := make([]string, 0, 2) - if attempt >= 2 { - dropOpenAIWSPayloadKey(payload, "include", &removed) - } - - if len(removed) == 0 { - return "full", nil - } - sort.Strings(removed) - return "trim_optional_fields", removed -} - -func logOpenAIWSModeInfo(format string, args ...any) { - logger.LegacyPrintf("service.openai_gateway", "[OpenAI WS Mode][openai_ws_mode=true] "+format, args...) -} - -func isOpenAIWSModeDebugEnabled() bool { - return logger.L().Core().Enabled(zap.DebugLevel) -} - -func logOpenAIWSModeDebug(format string, args ...any) { - if !isOpenAIWSModeDebugEnabled() { - return - } - logger.LegacyPrintf("service.openai_gateway", "[debug] [OpenAI WS Mode][openai_ws_mode=true] "+format, args...) -} - -func logOpenAIWSBindResponseAccountWarn(groupID, accountID int64, responseID string, err error) { - if err == nil { - return - } - logger.L().Warn( - "openai.ws_bind_response_account_failed", - zap.Int64("group_id", groupID), - zap.Int64("account_id", accountID), - zap.String("response_id", truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen)), - zap.Error(err), - ) -} - -func summarizeOpenAIWSReadCloseError(err error) (status string, reason string) { - if err == nil { - return "-", "-" - } - statusCode := coderws.CloseStatus(err) - if statusCode == -1 { - return "-", "-" - } - closeStatus := fmt.Sprintf("%d(%s)", int(statusCode), statusCode.String()) - closeReason := "-" - var closeErr coderws.CloseError - if errors.As(err, &closeErr) { - reasonText := strings.TrimSpace(closeErr.Reason) - if reasonText != "" { - closeReason = normalizeOpenAIWSLogValue(reasonText) - } - } - return normalizeOpenAIWSLogValue(closeStatus), closeReason -} - -func unwrapOpenAIWSDialBaseError(err error) error { - if err == nil { - return nil - } - var dialErr *openAIWSDialError - if errors.As(err, &dialErr) && dialErr != nil && dialErr.Err != nil { - return dialErr.Err - } - return err -} - -func openAIWSDialRespHeaderForLog(err error, key string) string { - var dialErr *openAIWSDialError - if !errors.As(err, &dialErr) || dialErr == nil || dialErr.ResponseHeaders == nil { - return "-" - } - return truncateOpenAIWSLogValue(dialErr.ResponseHeaders.Get(key), openAIWSHeaderValueMaxLen) -} - -func classifyOpenAIWSDialError(err error) string { - if err == nil { - return "-" - } - baseErr := unwrapOpenAIWSDialBaseError(err) - if baseErr == nil { - return "-" - } - if errors.Is(baseErr, context.DeadlineExceeded) { - return "ctx_deadline_exceeded" - } - if errors.Is(baseErr, context.Canceled) { - return "ctx_canceled" - } - var netErr net.Error - if errors.As(baseErr, &netErr) && netErr.Timeout() { - return "net_timeout" - } - if status := coderws.CloseStatus(baseErr); status != -1 { - return normalizeOpenAIWSLogValue(fmt.Sprintf("ws_close_%d", int(status))) - } - message := strings.ToLower(strings.TrimSpace(baseErr.Error())) - switch { - case strings.Contains(message, "handshake not finished"): - return "handshake_not_finished" - case strings.Contains(message, "bad handshake"): - return "bad_handshake" - case strings.Contains(message, "connection refused"): - return "connection_refused" - case strings.Contains(message, "no such host"): - return "dns_not_found" - case strings.Contains(message, "tls"): - return "tls_error" - case strings.Contains(message, "i/o timeout"): - return "io_timeout" - case strings.Contains(message, "context deadline exceeded"): - return "ctx_deadline_exceeded" - default: - return "dial_error" - } -} - -func summarizeOpenAIWSDialError(err error) ( - statusCode int, - dialClass string, - closeStatus string, - closeReason string, - respServer string, - respVia string, - respCFRay string, - respRequestID string, -) { - dialClass = "-" - closeStatus = "-" - closeReason = "-" - respServer = "-" - respVia = "-" - respCFRay = "-" - respRequestID = "-" - if err == nil { - return - } - var dialErr *openAIWSDialError - if errors.As(err, &dialErr) && dialErr != nil { - statusCode = dialErr.StatusCode - respServer = openAIWSDialRespHeaderForLog(err, "server") - respVia = openAIWSDialRespHeaderForLog(err, "via") - respCFRay = openAIWSDialRespHeaderForLog(err, "cf-ray") - respRequestID = openAIWSDialRespHeaderForLog(err, "x-request-id") - } - dialClass = normalizeOpenAIWSLogValue(classifyOpenAIWSDialError(err)) - closeStatus, closeReason = summarizeOpenAIWSReadCloseError(unwrapOpenAIWSDialBaseError(err)) - return -} - -func isOpenAIWSClientDisconnectError(err error) bool { - if err == nil { - return false - } - if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || errors.Is(err, context.Canceled) { - return true - } - switch coderws.CloseStatus(err) { - case coderws.StatusNormalClosure, coderws.StatusGoingAway, coderws.StatusNoStatusRcvd, coderws.StatusAbnormalClosure: - return true - } - message := strings.ToLower(strings.TrimSpace(err.Error())) - if message == "" { - return false - } - return strings.Contains(message, "failed to read frame header: eof") || - strings.Contains(message, "unexpected eof") || - strings.Contains(message, "use of closed network connection") || - strings.Contains(message, "connection reset by peer") || - strings.Contains(message, "broken pipe") || - strings.Contains(message, "an established connection was aborted") -} - -func classifyOpenAIWSReadFallbackReason(err error) string { - if err == nil { - return "read_event" - } - switch coderws.CloseStatus(err) { - case coderws.StatusPolicyViolation: - return "policy_violation" - case coderws.StatusMessageTooBig: - return "message_too_big" - default: - return "read_event" - } -} - -func sortedKeys(m map[string]any) []string { - if len(m) == 0 { - return nil - } - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - sort.Strings(keys) - return keys -} - func (s *OpenAIGatewayService) getOpenAIWSConnPool() *openAIWSConnPool { if s == nil { return nil @@ -1067,3609 +385,3 @@ func (s *OpenAIGatewayService) openAIWSAcquireTimeout() time.Duration { } return dial + 2*time.Second } - -func (s *OpenAIGatewayService) buildOpenAIResponsesWSURL(account *Account) (string, error) { - if account == nil { - return "", errors.New("account is nil") - } - var targetURL string - switch account.Type { - case AccountTypeOAuth: - targetURL = chatgptCodexURL - case AccountTypeAPIKey: - baseURL := account.GetOpenAIBaseURL() - if baseURL == "" { - targetURL = openaiPlatformAPIURL - } else { - validatedURL, err := s.validateUpstreamBaseURL(baseURL) - if err != nil { - return "", err - } - targetURL = buildOpenAIResponsesURL(validatedURL) - } - default: - targetURL = openaiPlatformAPIURL - } - - parsed, err := url.Parse(strings.TrimSpace(targetURL)) - if err != nil { - return "", fmt.Errorf("invalid target url: %w", err) - } - switch strings.ToLower(parsed.Scheme) { - case "https": - parsed.Scheme = "wss" - case "http": - parsed.Scheme = "ws" - case "wss", "ws": - // 保持不变 - default: - return "", fmt.Errorf("unsupported scheme for ws: %s", parsed.Scheme) - } - return parsed.String(), nil -} - -func (s *OpenAIGatewayService) buildOpenAIWSHeaders( - ctx context.Context, - c *gin.Context, - account *Account, - token string, - decision OpenAIWSProtocolDecision, - isCodexCLI bool, - turnState string, - turnMetadata string, - promptCacheKey string, -) (http.Header, openAIWSSessionHeaderResolution, error) { - headers := make(http.Header) - headers.Set("authorization", "Bearer "+token) - - sessionResolution := resolveOpenAIWSSessionHeaders(c, promptCacheKey) - if c != nil && c.Request != nil { - if v := strings.TrimSpace(c.Request.Header.Get("accept-language")); v != "" { - headers.Set("accept-language", v) - } - } - // OAuth 账号：将 apiKeyID 混入 session 标识符，防止跨用户会话碰撞。 - if account != nil && account.Type == AccountTypeOAuth { - apiKeyID := getAPIKeyIDFromContext(c) - if sessionResolution.SessionID != "" { - headers.Set("session_id", isolateOpenAISessionID(apiKeyID, sessionResolution.SessionID)) - } - if sessionResolution.ConversationID != "" { - headers.Set("conversation_id", isolateOpenAISessionID(apiKeyID, sessionResolution.ConversationID)) - } - } else { - if sessionResolution.SessionID != "" { - headers.Set("session_id", sessionResolution.SessionID) - } - if sessionResolution.ConversationID != "" { - headers.Set("conversation_id", sessionResolution.ConversationID) - } - } - if state := strings.TrimSpace(turnState); state != "" { - headers.Set(openAIWSTurnStateHeader, state) - } - if metadata := strings.TrimSpace(turnMetadata); metadata != "" { - headers.Set(openAIWSTurnMetadataHeader, metadata) - } - - if account != nil && account.Type == AccountTypeOAuth { - if err := resolveAndSetOpenAIChatGPTAccountHeaders(ctx, s.accountRepo, headers, account); err != nil { - return nil, sessionResolution, fmt.Errorf("resolve chatgpt account headers: %w", err) - } - headers.Set("originator", resolveOpenAIUpstreamOriginator(c, isCodexCLI)) - } - - betaValue := openAIWSBetaV2Value - if decision.Transport == OpenAIUpstreamTransportResponsesWebsocket { - betaValue = openAIWSBetaV1Value - } - headers.Set("OpenAI-Beta", betaValue) - - customUA := "" - if account != nil { - customUA = account.GetOpenAIUserAgent() - } - if strings.TrimSpace(customUA) != "" { - headers.Set("user-agent", customUA) - } else if c != nil { - if ua := strings.TrimSpace(c.GetHeader("User-Agent")); ua != "" { - headers.Set("user-agent", ua) - } - } - if s != nil && s.cfg != nil && s.cfg.Gateway.ForceCodexCLI { - headers.Set("user-agent", codexCLIUserAgent) - } - if account != nil && account.Type == AccountTypeOAuth && !openai.IsCodexCLIRequest(headers.Get("user-agent")) { - headers.Set("user-agent", codexCLIUserAgent) - } - - // 账号级请求头覆写（仅 openai api_key 账号启用时生效；OAuth 路径 no-op）。 - // 覆盖所有 WS 模式（ctx_pool/dedicated/passthrough）的握手头。 - account.ApplyHeaderOverrides(headers) - - return headers, sessionResolution, nil -} - -func (s *OpenAIGatewayService) buildOpenAIWSCreatePayload(reqBody map[string]any, account *Account) map[string]any { - // OpenAI WS Mode 协议：response.create 字段与 HTTP /responses 基本一致。 - // 保留 stream 字段（与 Codex CLI 一致），仅移除 background。 - payload := make(map[string]any, len(reqBody)+1) - for k, v := range reqBody { - payload[k] = v - } - - delete(payload, "background") - if _, exists := payload["stream"]; !exists { - payload["stream"] = true - } - payload["type"] = "response.create" - - // OAuth 默认保持 store=false，避免误依赖服务端历史。 - if account != nil && account.Type == AccountTypeOAuth && !s.isOpenAIWSStoreRecoveryAllowed(account) { - payload["store"] = false - } - return payload -} - -func setOpenAIWSTurnMetadata(payload map[string]any, turnMetadata string) { - if len(payload) == 0 { - return - } - metadata := strings.TrimSpace(turnMetadata) - if metadata == "" { - return - } - - switch existing := payload["client_metadata"].(type) { - case map[string]any: - existing[openAIWSTurnMetadataHeader] = metadata - payload["client_metadata"] = existing - case map[string]string: - next := make(map[string]any, len(existing)+1) - for k, v := range existing { - next[k] = v - } - next[openAIWSTurnMetadataHeader] = metadata - payload["client_metadata"] = next - default: - payload["client_metadata"] = map[string]any{ - openAIWSTurnMetadataHeader: metadata, - } - } -} - -func (s *OpenAIGatewayService) isOpenAIWSStoreRecoveryAllowed(account *Account) bool { - if account != nil && account.IsOpenAIWSAllowStoreRecoveryEnabled() { - return true - } - if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.AllowStoreRecovery { - return true - } - return false -} - -func (s *OpenAIGatewayService) isOpenAIWSStoreDisabledInRequest(reqBody map[string]any, account *Account) bool { - if account != nil && account.Type == AccountTypeOAuth && !s.isOpenAIWSStoreRecoveryAllowed(account) { - return true - } - if len(reqBody) == 0 { - return false - } - rawStore, ok := reqBody["store"] - if !ok { - return false - } - storeEnabled, ok := rawStore.(bool) - if !ok { - return false - } - return !storeEnabled -} - -func (s *OpenAIGatewayService) isOpenAIWSStoreDisabledInRequestRaw(reqBody []byte, account *Account) bool { - if account != nil && account.Type == AccountTypeOAuth && !s.isOpenAIWSStoreRecoveryAllowed(account) { - return true - } - if len(reqBody) == 0 { - return false - } - storeValue := gjson.GetBytes(reqBody, "store") - if !storeValue.Exists() { - return false - } - if storeValue.Type != gjson.True && storeValue.Type != gjson.False { - return false - } - return !storeValue.Bool() -} - -func (s *OpenAIGatewayService) openAIWSStoreDisabledConnMode() string { - if s == nil || s.cfg == nil { - return openAIWSStoreDisabledConnModeStrict - } - mode := strings.ToLower(strings.TrimSpace(s.cfg.Gateway.OpenAIWS.StoreDisabledConnMode)) - switch mode { - case openAIWSStoreDisabledConnModeStrict, openAIWSStoreDisabledConnModeAdaptive, openAIWSStoreDisabledConnModeOff: - return mode - case "": - // 兼容旧配置：仅配置了布尔开关时按旧语义推导。 - if s.cfg.Gateway.OpenAIWS.StoreDisabledForceNewConn { - return openAIWSStoreDisabledConnModeStrict - } - return openAIWSStoreDisabledConnModeOff - default: - return openAIWSStoreDisabledConnModeStrict - } -} - -func shouldForceNewConnOnStoreDisabled(mode, lastFailureReason string) bool { - switch mode { - case openAIWSStoreDisabledConnModeOff: - return false - case openAIWSStoreDisabledConnModeAdaptive: - reason := strings.TrimPrefix(strings.TrimSpace(lastFailureReason), "prewarm_") - switch reason { - case "policy_violation", "message_too_big", "auth_failed", "write_request", "write": - return true - default: - return false - } - default: - return true - } -} - -func dropPreviousResponseIDFromRawPayload(payload []byte) ([]byte, bool, error) { - return dropPreviousResponseIDFromRawPayloadWithDeleteFn(payload, sjson.DeleteBytes) -} - -func dropPreviousResponseIDFromRawPayloadWithDeleteFn( - payload []byte, - deleteFn func([]byte, string) ([]byte, error), -) ([]byte, bool, error) { - if len(payload) == 0 { - return payload, false, nil - } - if !gjson.GetBytes(payload, "previous_response_id").Exists() { - return payload, false, nil - } - if deleteFn == nil { - deleteFn = sjson.DeleteBytes - } - - updated := payload - for i := 0; i < openAIWSMaxPrevResponseIDDeletePasses && - gjson.GetBytes(updated, "previous_response_id").Exists(); i++ { - next, err := deleteFn(updated, "previous_response_id") - if err != nil { - return payload, false, err - } - updated = next - } - return updated, !gjson.GetBytes(updated, "previous_response_id").Exists(), nil -} - -func setPreviousResponseIDToRawPayload(payload []byte, previousResponseID string) ([]byte, error) { - normalizedPrevID := strings.TrimSpace(previousResponseID) - if len(payload) == 0 || normalizedPrevID == "" { - return payload, nil - } - updated, err := sjson.SetBytes(payload, "previous_response_id", normalizedPrevID) - if err == nil { - return updated, nil - } - - var reqBody map[string]any - if unmarshalErr := json.Unmarshal(payload, &reqBody); unmarshalErr != nil { - return nil, err - } - reqBody["previous_response_id"] = normalizedPrevID - rebuilt, marshalErr := json.Marshal(reqBody) - if marshalErr != nil { - return nil, marshalErr - } - return rebuilt, nil -} - -func shouldInferIngressFunctionCallOutputPreviousResponseID( - storeDisabled bool, - turn int, - signals ToolContinuationSignals, - currentPreviousResponseID string, - expectedPreviousResponseID string, -) bool { - if !storeDisabled || turn <= 1 || !signals.HasFunctionCallOutput { - return false - } - if strings.TrimSpace(currentPreviousResponseID) != "" { - return false - } - if signals.HasFunctionCallOutputMissingCallID { - return false - } - // If the client already sent the actual tool-call context, treat this as - // a full replay / self-contained continuation payload rather than - // downgrading it into an inferred delta continuation. item_reference alone - // is not enough on the store=false WS path: it still needs a valid prior - // response anchor so upstream can resolve the referenced function_call. - if signals.HasToolCallContext { - return false - } - return strings.TrimSpace(expectedPreviousResponseID) != "" -} - -func alignStoreDisabledPreviousResponseID( - payload []byte, - expectedPreviousResponseID string, -) ([]byte, bool, error) { - if len(payload) == 0 { - return payload, false, nil - } - expected := strings.TrimSpace(expectedPreviousResponseID) - if expected == "" { - return payload, false, nil - } - current := openAIWSPayloadStringFromRaw(payload, "previous_response_id") - if current == "" || current == expected { - return payload, false, nil - } - - withoutPrev, removed, dropErr := dropPreviousResponseIDFromRawPayload(payload) - if dropErr != nil { - return payload, false, dropErr - } - if !removed { - return payload, false, nil - } - updated, setErr := setPreviousResponseIDToRawPayload(withoutPrev, expected) - if setErr != nil { - return payload, false, setErr - } - return updated, true, nil -} - -func cloneOpenAIWSPayloadBytes(payload []byte) []byte { - if len(payload) == 0 { - return nil - } - cloned := make([]byte, len(payload)) - copy(cloned, payload) - return cloned -} - -func cloneOpenAIWSRawMessages(items []json.RawMessage) []json.RawMessage { - if items == nil { - return nil - } - cloned := make([]json.RawMessage, 0, len(items)) - for idx := range items { - cloned = append(cloned, json.RawMessage(cloneOpenAIWSPayloadBytes(items[idx]))) - } - return cloned -} - -func normalizeOpenAIWSJSONForCompare(raw []byte) ([]byte, error) { - trimmed := bytes.TrimSpace(raw) - if len(trimmed) == 0 { - return nil, errors.New("json is empty") - } - var decoded any - if err := json.Unmarshal(trimmed, &decoded); err != nil { - return nil, err - } - return json.Marshal(decoded) -} - -func normalizeOpenAIWSJSONForCompareOrRaw(raw []byte) []byte { - normalized, err := normalizeOpenAIWSJSONForCompare(raw) - if err != nil { - return bytes.TrimSpace(raw) - } - return normalized -} - -func normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(payload []byte) ([]byte, error) { - if len(payload) == 0 { - return nil, errors.New("payload is empty") - } - var decoded map[string]any - if err := json.Unmarshal(payload, &decoded); err != nil { - return nil, err - } - delete(decoded, "input") - delete(decoded, "previous_response_id") - return json.Marshal(decoded) -} - -func openAIWSExtractNormalizedInputSequence(payload []byte) ([]json.RawMessage, bool, error) { - if len(payload) == 0 { - return nil, false, nil - } - inputValue := gjson.GetBytes(payload, "input") - if !inputValue.Exists() { - return nil, false, nil - } - if inputValue.Type == gjson.JSON { - raw := strings.TrimSpace(inputValue.Raw) - if strings.HasPrefix(raw, "[") { - var items []json.RawMessage - if err := json.Unmarshal([]byte(raw), &items); err != nil { - return nil, true, err - } - return items, true, nil - } - return []json.RawMessage{json.RawMessage(raw)}, true, nil - } - if inputValue.Type == gjson.String { - encoded, _ := json.Marshal(inputValue.String()) - return []json.RawMessage{encoded}, true, nil - } - return []json.RawMessage{json.RawMessage(inputValue.Raw)}, true, nil -} - -func openAIWSInputIsPrefixExtended(previousPayload, currentPayload []byte) (bool, error) { - previousItems, previousExists, prevErr := openAIWSExtractNormalizedInputSequence(previousPayload) - if prevErr != nil { - return false, prevErr - } - currentItems, currentExists, currentErr := openAIWSExtractNormalizedInputSequence(currentPayload) - if currentErr != nil { - return false, currentErr - } - if !previousExists && !currentExists { - return true, nil - } - if !previousExists { - return len(currentItems) == 0, nil - } - if !currentExists { - return len(previousItems) == 0, nil - } - if len(currentItems) < len(previousItems) { - return false, nil - } - - for idx := range previousItems { - previousNormalized := normalizeOpenAIWSJSONForCompareOrRaw(previousItems[idx]) - currentNormalized := normalizeOpenAIWSJSONForCompareOrRaw(currentItems[idx]) - if !bytes.Equal(previousNormalized, currentNormalized) { - return false, nil - } - } - return true, nil -} - -func openAIWSRawItemsHasPrefix(items []json.RawMessage, prefix []json.RawMessage) bool { - if len(prefix) == 0 { - return true - } - if len(items) < len(prefix) { - return false - } - for idx := range prefix { - previousNormalized := normalizeOpenAIWSJSONForCompareOrRaw(prefix[idx]) - currentNormalized := normalizeOpenAIWSJSONForCompareOrRaw(items[idx]) - if !bytes.Equal(previousNormalized, currentNormalized) { - return false - } - } - return true -} - -func openAIWSRawItemsHasFunctionCallOutput(items []json.RawMessage) bool { - for _, item := range items { - if isCodexToolCallOutputItemType(gjson.GetBytes(item, "type").String()) { - return true - } - } - return false -} - -func openAIWSRawItemsHaveToolCallContextForOutputs(items []json.RawMessage) bool { - if len(items) == 0 { - return false - } - contextCallIDs := make(map[string]struct{}) - outputCallIDs := make(map[string]struct{}) - for _, item := range items { - itemType := gjson.GetBytes(item, "type").String() - callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String()) - switch { - case isCodexToolCallContextItemType(itemType): - if callID != "" { - contextCallIDs[callID] = struct{}{} - } - case isCodexToolCallOutputItemType(itemType): - if callID == "" { - return false - } - outputCallIDs[callID] = struct{}{} - } - } - if len(outputCallIDs) == 0 || len(contextCallIDs) == 0 { - return false - } - for callID := range outputCallIDs { - if _, ok := contextCallIDs[callID]; !ok { - return false - } - } - return true -} - -func openAIWSRawPayloadHasToolCallOutput(payload []byte) bool { - if len(payload) == 0 { - return false - } - input := gjson.GetBytes(payload, "input") - if !input.Exists() { - return false - } - if input.IsArray() { - for _, item := range input.Array() { - if isCodexToolCallOutputItemType(item.Get("type").String()) { - return true - } - } - return false - } - if input.Type == gjson.JSON { - return isCodexToolCallOutputItemType(input.Get("type").String()) - } - return false -} - -func buildOpenAIWSReplayInputSequence( - previousFullInput []json.RawMessage, - previousFullInputExists bool, - currentPayload []byte, - hasPreviousResponseID bool, -) ([]json.RawMessage, bool, error) { - currentItems, currentExists, currentErr := openAIWSExtractNormalizedInputSequence(currentPayload) - if currentErr != nil { - return nil, false, currentErr - } - if !hasPreviousResponseID { - return cloneOpenAIWSRawMessages(currentItems), currentExists, nil - } - if !previousFullInputExists { - return cloneOpenAIWSRawMessages(currentItems), currentExists, nil - } - if !currentExists || len(currentItems) == 0 { - return cloneOpenAIWSRawMessages(previousFullInput), true, nil - } - if openAIWSRawItemsHasPrefix(currentItems, previousFullInput) { - return cloneOpenAIWSRawMessages(currentItems), true, nil - } - merged := make([]json.RawMessage, 0, len(previousFullInput)+len(currentItems)) - merged = append(merged, cloneOpenAIWSRawMessages(previousFullInput)...) - merged = append(merged, cloneOpenAIWSRawMessages(currentItems)...) - return merged, true, nil -} - -func setOpenAIWSPayloadInputSequence( - payload []byte, - fullInput []json.RawMessage, - fullInputExists bool, -) ([]byte, error) { - if !fullInputExists { - return payload, nil - } - // Preserve [] vs null semantics when input exists but is empty. - inputForMarshal := fullInput - if inputForMarshal == nil { - inputForMarshal = []json.RawMessage{} - } - inputRaw, marshalErr := json.Marshal(inputForMarshal) - if marshalErr != nil { - return nil, marshalErr - } - return sjson.SetRawBytes(payload, "input", inputRaw) -} - -func shouldKeepIngressPreviousResponseID( - previousPayload []byte, - currentPayload []byte, - lastTurnResponseID string, - hasFunctionCallOutput bool, -) (bool, string, error) { - if hasFunctionCallOutput { - return true, "has_function_call_output", nil - } - currentPreviousResponseID := strings.TrimSpace(openAIWSPayloadStringFromRaw(currentPayload, "previous_response_id")) - if currentPreviousResponseID == "" { - return false, "missing_previous_response_id", nil - } - expectedPreviousResponseID := strings.TrimSpace(lastTurnResponseID) - if expectedPreviousResponseID == "" { - return false, "missing_last_turn_response_id", nil - } - if currentPreviousResponseID != expectedPreviousResponseID { - return false, "previous_response_id_mismatch", nil - } - if len(previousPayload) == 0 { - return false, "missing_previous_turn_payload", nil - } - - previousComparable, previousComparableErr := normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(previousPayload) - if previousComparableErr != nil { - return false, "non_input_compare_error", previousComparableErr - } - currentComparable, currentComparableErr := normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(currentPayload) - if currentComparableErr != nil { - return false, "non_input_compare_error", currentComparableErr - } - if !bytes.Equal(previousComparable, currentComparable) { - return false, "non_input_changed", nil - } - return true, "strict_incremental_ok", nil -} - -type openAIWSIngressPreviousTurnStrictState struct { - nonInputComparable []byte -} - -func buildOpenAIWSIngressPreviousTurnStrictState(payload []byte) (*openAIWSIngressPreviousTurnStrictState, error) { - if len(payload) == 0 { - return nil, nil - } - nonInputComparable, nonInputErr := normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(payload) - if nonInputErr != nil { - return nil, nonInputErr - } - return &openAIWSIngressPreviousTurnStrictState{ - nonInputComparable: nonInputComparable, - }, nil -} - -func shouldKeepIngressPreviousResponseIDWithStrictState( - previousState *openAIWSIngressPreviousTurnStrictState, - currentPayload []byte, - lastTurnResponseID string, - hasFunctionCallOutput bool, -) (bool, string, error) { - if hasFunctionCallOutput { - return true, "has_function_call_output", nil - } - currentPreviousResponseID := strings.TrimSpace(openAIWSPayloadStringFromRaw(currentPayload, "previous_response_id")) - if currentPreviousResponseID == "" { - return false, "missing_previous_response_id", nil - } - expectedPreviousResponseID := strings.TrimSpace(lastTurnResponseID) - if expectedPreviousResponseID == "" { - return false, "missing_last_turn_response_id", nil - } - if currentPreviousResponseID != expectedPreviousResponseID { - return false, "previous_response_id_mismatch", nil - } - if previousState == nil { - return false, "missing_previous_turn_payload", nil - } - - currentComparable, currentComparableErr := normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(currentPayload) - if currentComparableErr != nil { - return false, "non_input_compare_error", currentComparableErr - } - if !bytes.Equal(previousState.nonInputComparable, currentComparable) { - return false, "non_input_changed", nil - } - return true, "strict_incremental_ok", nil -} - -func (s *OpenAIGatewayService) forwardOpenAIWSV2( - ctx context.Context, - c *gin.Context, - account *Account, - reqBody map[string]any, - token string, - decision OpenAIWSProtocolDecision, - isCodexCLI bool, - reqStream bool, - originalModel string, - mappedModel string, - startTime time.Time, - attempt int, - lastFailureReason string, -) (*OpenAIForwardResult, error) { - if s == nil || account == nil { - return nil, wrapOpenAIWSFallback("invalid_state", errors.New("service or account is nil")) - } - - wsURL, err := s.buildOpenAIResponsesWSURL(account) - if err != nil { - return nil, wrapOpenAIWSFallback("build_ws_url", err) - } - wsHost := "-" - wsPath := "-" - if parsed, parseErr := url.Parse(wsURL); parseErr == nil && parsed != nil { - if h := strings.TrimSpace(parsed.Host); h != "" { - wsHost = normalizeOpenAIWSLogValue(h) - } - if p := strings.TrimSpace(parsed.Path); p != "" { - wsPath = normalizeOpenAIWSLogValue(p) - } - } - logOpenAIWSModeDebug( - "dial_target account_id=%d account_type=%s ws_host=%s ws_path=%s", - account.ID, - account.Type, - wsHost, - wsPath, - ) - - payload := s.buildOpenAIWSCreatePayload(reqBody, account) - payloadStrategy, removedKeys := applyOpenAIWSRetryPayloadStrategy(payload, attempt) - previousResponseID := openAIWSPayloadString(payload, "previous_response_id") - previousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(previousResponseID) - promptCacheKey := openAIWSPayloadString(payload, "prompt_cache_key") - _, hasTools := payload["tools"] - debugEnabled := isOpenAIWSModeDebugEnabled() - payloadBytes := -1 - resolvePayloadBytes := func() int { - if payloadBytes >= 0 { - return payloadBytes - } - payloadBytes = len(payloadAsJSONBytes(payload)) - return payloadBytes - } - streamValue := "-" - if raw, ok := payload["stream"]; ok { - streamValue = normalizeOpenAIWSLogValue(strings.TrimSpace(fmt.Sprintf("%v", raw))) - } - turnState := "" - turnMetadata := "" - if c != nil && c.Request != nil { - turnState = strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader)) - turnMetadata = strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)) - } - setOpenAIWSTurnMetadata(payload, turnMetadata) - payloadEventType := openAIWSPayloadString(payload, "type") - if payloadEventType == "" { - payloadEventType = "response.create" - } - if s.shouldEmitOpenAIWSPayloadSchema(attempt) { - logOpenAIWSModeInfo( - "[debug] payload_schema account_id=%d attempt=%d event=%s payload_keys=%s payload_bytes=%d payload_key_sizes=%s input_summary=%s stream=%s payload_strategy=%s removed_keys=%s has_previous_response_id=%v has_prompt_cache_key=%v has_tools=%v", - account.ID, - attempt, - payloadEventType, - normalizeOpenAIWSLogValue(strings.Join(sortedKeys(payload), ",")), - resolvePayloadBytes(), - normalizeOpenAIWSLogValue(summarizeOpenAIWSPayloadKeySizes(payload, openAIWSPayloadKeySizeTopN)), - normalizeOpenAIWSLogValue(summarizeOpenAIWSInput(payload["input"])), - streamValue, - normalizeOpenAIWSLogValue(payloadStrategy), - normalizeOpenAIWSLogValue(strings.Join(removedKeys, ",")), - previousResponseID != "", - promptCacheKey != "", - hasTools, - ) - } - - stateStore := s.getOpenAIWSStateStore() - groupID := getOpenAIGroupIDFromContext(c) - sessionHash := s.GenerateSessionHash(c, nil) - if sessionHash == "" { - var legacySessionHash string - sessionHash, legacySessionHash = openAIWSSessionHashesFromID(promptCacheKey) - attachOpenAILegacySessionHashToGin(c, legacySessionHash) - } - if turnState == "" && stateStore != nil && sessionHash != "" { - if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { - turnState = savedTurnState - } - } - preferredConnID := "" - if stateStore != nil && previousResponseID != "" { - if connID, ok := stateStore.GetResponseConn(previousResponseID); ok { - preferredConnID = connID - } - } - storeDisabled := s.isOpenAIWSStoreDisabledInRequest(reqBody, account) - if stateStore != nil && storeDisabled && previousResponseID == "" && sessionHash != "" { - if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok { - preferredConnID = connID - } - } - storeDisabledConnMode := s.openAIWSStoreDisabledConnMode() - forceNewConnByPolicy := shouldForceNewConnOnStoreDisabled(storeDisabledConnMode, lastFailureReason) - forceNewConn := forceNewConnByPolicy && storeDisabled && previousResponseID == "" && sessionHash != "" && preferredConnID == "" - wsHeaders, sessionResolution, buildHdrErr := s.buildOpenAIWSHeaders(ctx, c, account, token, decision, isCodexCLI, turnState, turnMetadata, promptCacheKey) - if buildHdrErr != nil { - return nil, fmt.Errorf("build ws headers: %w", buildHdrErr) - } - logOpenAIWSModeDebug( - "acquire_start account_id=%d account_type=%s transport=%s preferred_conn_id=%s has_previous_response_id=%v session_hash=%s has_turn_state=%v turn_state_len=%d has_turn_metadata=%v turn_metadata_len=%d store_disabled=%v store_disabled_conn_mode=%s retry_last_reason=%s force_new_conn=%v header_user_agent=%s header_openai_beta=%s header_originator=%s header_accept_language=%s header_session_id=%s header_conversation_id=%s session_id_source=%s conversation_id_source=%s has_prompt_cache_key=%v has_chatgpt_account_id=%v has_authorization=%v has_session_id=%v has_conversation_id=%v proxy_enabled=%v", - account.ID, - account.Type, - normalizeOpenAIWSLogValue(string(decision.Transport)), - truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), - previousResponseID != "", - truncateOpenAIWSLogValue(sessionHash, 12), - turnState != "", - len(turnState), - turnMetadata != "", - len(turnMetadata), - storeDisabled, - normalizeOpenAIWSLogValue(storeDisabledConnMode), - truncateOpenAIWSLogValue(lastFailureReason, openAIWSLogValueMaxLen), - forceNewConn, - openAIWSHeaderValueForLog(wsHeaders, "user-agent"), - openAIWSHeaderValueForLog(wsHeaders, "openai-beta"), - openAIWSHeaderValueForLog(wsHeaders, "originator"), - openAIWSHeaderValueForLog(wsHeaders, "accept-language"), - openAIWSHeaderValueForLog(wsHeaders, "session_id"), - openAIWSHeaderValueForLog(wsHeaders, "conversation_id"), - normalizeOpenAIWSLogValue(sessionResolution.SessionSource), - normalizeOpenAIWSLogValue(sessionResolution.ConversationSource), - promptCacheKey != "", - hasOpenAIWSHeader(wsHeaders, "chatgpt-account-id"), - hasOpenAIWSHeader(wsHeaders, "authorization"), - hasOpenAIWSHeader(wsHeaders, "session_id"), - hasOpenAIWSHeader(wsHeaders, "conversation_id"), - account.ProxyID != nil && account.Proxy != nil, - ) - - acquireCtx, acquireCancel := context.WithTimeout(ctx, s.openAIWSAcquireTimeout()) - defer acquireCancel() - - lease, err := s.getOpenAIWSConnPool().Acquire(acquireCtx, openAIWSAcquireRequest{ - Account: account, - WSURL: wsURL, - Headers: wsHeaders, - PreferredConnID: preferredConnID, - ForceNewConn: forceNewConn, - ProxyURL: func() string { - if account.ProxyID != nil && account.Proxy != nil { - return account.Proxy.URL() - } - return "" - }(), - }) - if err != nil { - dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(err) - logOpenAIWSModeInfo( - "acquire_fail account_id=%d account_type=%s transport=%s reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_new_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v", - account.ID, - account.Type, - normalizeOpenAIWSLogValue(string(decision.Transport)), - normalizeOpenAIWSLogValue(classifyOpenAIWSAcquireError(err)), - dialStatus, - dialClass, - dialCloseStatus, - truncateOpenAIWSLogValue(dialCloseReason, openAIWSHeaderValueMaxLen), - dialRespServer, - dialRespVia, - dialRespCFRay, - dialRespReqID, - truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), - truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), - forceNewConn, - wsHost, - wsPath, - account.ProxyID != nil && account.Proxy != nil, - ) - var dialErr *openAIWSDialError - if errors.As(err, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusTooManyRequests { - s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error())) - } - return nil, wrapOpenAIWSFallback(classifyOpenAIWSAcquireError(err), err) - } - // cleanExit 标记正常终端事件退出，此时上游不会再发送帧，连接可安全归还复用。 - // 所有异常路径（读写错误、error 事件等）已在各自分支中提前调用 MarkBroken， - // 因此 defer 中只需处理正常退出时不 MarkBroken 即可。 - cleanExit := false - defer func() { - if !cleanExit { - lease.MarkBroken() - } - lease.Release() - }() - connID := strings.TrimSpace(lease.ConnID()) - logOpenAIWSModeDebug( - "connected account_id=%d account_type=%s transport=%s conn_id=%s conn_reused=%v conn_pick_ms=%d queue_wait_ms=%d has_previous_response_id=%v", - account.ID, - account.Type, - normalizeOpenAIWSLogValue(string(decision.Transport)), - connID, - lease.Reused(), - lease.ConnPickDuration().Milliseconds(), - lease.QueueWaitDuration().Milliseconds(), - previousResponseID != "", - ) - if previousResponseID != "" { - logOpenAIWSModeInfo( - "continuation_probe account_id=%d account_type=%s conn_id=%s previous_response_id=%s previous_response_id_kind=%s preferred_conn_id=%s conn_reused=%v store_disabled=%v session_hash=%s header_session_id=%s header_conversation_id=%s session_id_source=%s conversation_id_source=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v", - account.ID, - account.Type, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(previousResponseIDKind), - truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), - lease.Reused(), - storeDisabled, - truncateOpenAIWSLogValue(sessionHash, 12), - openAIWSHeaderValueForLog(wsHeaders, "session_id"), - openAIWSHeaderValueForLog(wsHeaders, "conversation_id"), - normalizeOpenAIWSLogValue(sessionResolution.SessionSource), - normalizeOpenAIWSLogValue(sessionResolution.ConversationSource), - turnState != "", - len(turnState), - promptCacheKey != "", - ) - } - if c != nil { - SetOpsLatencyMs(c, OpsOpenAIWSConnPickMsKey, lease.ConnPickDuration().Milliseconds()) - SetOpsLatencyMs(c, OpsOpenAIWSQueueWaitMsKey, lease.QueueWaitDuration().Milliseconds()) - c.Set(OpsOpenAIWSConnReusedKey, lease.Reused()) - if connID != "" { - c.Set(OpsOpenAIWSConnIDKey, connID) - } - } - - handshakeTurnState := strings.TrimSpace(lease.HandshakeHeader(openAIWSTurnStateHeader)) - logOpenAIWSModeDebug( - "handshake account_id=%d conn_id=%s has_turn_state=%v turn_state_len=%d", - account.ID, - connID, - handshakeTurnState != "", - len(handshakeTurnState), - ) - if handshakeTurnState != "" { - if stateStore != nil && sessionHash != "" { - stateStore.BindSessionTurnState(groupID, sessionHash, handshakeTurnState, s.openAIWSSessionStickyTTL()) - } - if c != nil { - c.Header(http.CanonicalHeaderKey(openAIWSTurnStateHeader), handshakeTurnState) - } - } - - if err := s.performOpenAIWSGeneratePrewarm( - ctx, - lease, - decision, - payload, - previousResponseID, - reqBody, - account, - stateStore, - groupID, - ); err != nil { - return nil, err - } - - if err := lease.WriteJSONWithContextTimeout(ctx, payload, s.openAIWSWriteTimeout()); err != nil { - lease.MarkBroken() - logOpenAIWSModeInfo( - "write_request_fail account_id=%d conn_id=%s cause=%s payload_bytes=%d", - account.ID, - connID, - truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), - resolvePayloadBytes(), - ) - return nil, wrapOpenAIWSFallback("write_request", err) - } - if debugEnabled { - logOpenAIWSModeDebug( - "write_request_sent account_id=%d conn_id=%s stream=%v payload_bytes=%d previous_response_id=%s", - account.ID, - connID, - reqStream, - resolvePayloadBytes(), - truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen), - ) - } - - usage := &OpenAIUsage{} - imageCounter := newOpenAIImageOutputCounter() - var firstTokenMs *int - responseID := "" - var finalResponse []byte - wroteDownstream := false - needModelReplace := originalModel != mappedModel - var mappedModelBytes []byte - if needModelReplace && mappedModel != "" { - mappedModelBytes = []byte(mappedModel) - } - bufferedStreamEvents := make([][]byte, 0, 4) - eventCount := 0 - tokenEventCount := 0 - terminalEventCount := 0 - bufferedEventCount := 0 - flushedBufferedEventCount := 0 - firstEventType := "" - lastEventType := "" - - var flusher http.Flusher - if reqStream { - if s.responseHeaderFilter != nil { - responseheaders.WriteFilteredHeaders(c.Writer.Header(), http.Header{}, s.responseHeaderFilter) - } - c.Header("Content-Type", "text/event-stream") - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - c.Header("X-Accel-Buffering", "no") - f, ok := c.Writer.(http.Flusher) - if !ok { - lease.MarkBroken() - return nil, wrapOpenAIWSFallback("streaming_not_supported", errors.New("streaming not supported")) - } - flusher = f - } - - clientDisconnected := false - flushBatchSize := s.openAIWSEventFlushBatchSize() - flushInterval := s.openAIWSEventFlushInterval() - pendingFlushEvents := 0 - lastFlushAt := time.Now() - flushStreamWriter := func(force bool) { - if clientDisconnected || flusher == nil || pendingFlushEvents <= 0 { - return - } - if !force && flushBatchSize > 1 && pendingFlushEvents < flushBatchSize { - if flushInterval <= 0 || time.Since(lastFlushAt) < flushInterval { - return - } - } - flusher.Flush() - pendingFlushEvents = 0 - lastFlushAt = time.Now() - } - emitStreamMessage := func(message []byte, forceFlush bool) { - if clientDisconnected { - return - } - frame := make([]byte, 0, len(message)+8) - frame = append(frame, "data: "...) - frame = append(frame, message...) - frame = append(frame, '\n', '\n') - _, wErr := c.Writer.Write(frame) - if wErr == nil { - wroteDownstream = true - pendingFlushEvents++ - flushStreamWriter(forceFlush) - return - } - clientDisconnected = true - logger.LegacyPrintf("service.openai_gateway", "[OpenAI WS Mode] client disconnected, continue draining upstream: account=%d", account.ID) - } - flushBufferedStreamEvents := func(reason string) { - if len(bufferedStreamEvents) == 0 { - return - } - flushed := len(bufferedStreamEvents) - for _, buffered := range bufferedStreamEvents { - emitStreamMessage(buffered, false) - } - bufferedStreamEvents = bufferedStreamEvents[:0] - flushStreamWriter(true) - flushedBufferedEventCount += flushed - if debugEnabled { - logOpenAIWSModeDebug( - "buffer_flush account_id=%d conn_id=%s reason=%s flushed=%d total_flushed=%d client_disconnected=%v", - account.ID, - connID, - truncateOpenAIWSLogValue(reason, openAIWSLogValueMaxLen), - flushed, - flushedBufferedEventCount, - clientDisconnected, - ) - } - } - - readTimeout := s.openAIWSReadTimeout() - - for { - message, readErr := lease.ReadMessageWithContextTimeout(ctx, readTimeout) - if readErr != nil { - lease.MarkBroken() - closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) - logOpenAIWSModeInfo( - "read_fail account_id=%d conn_id=%s wrote_downstream=%v close_status=%s close_reason=%s cause=%s events=%d token_events=%d terminal_events=%d buffered_pending=%d buffered_flushed=%d first_event=%s last_event=%s", - account.ID, - connID, - wroteDownstream, - closeStatus, - closeReason, - truncateOpenAIWSLogValue(readErr.Error(), openAIWSLogValueMaxLen), - eventCount, - tokenEventCount, - terminalEventCount, - len(bufferedStreamEvents), - flushedBufferedEventCount, - truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen), - truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen), - ) - if !wroteDownstream { - return nil, wrapOpenAIWSFallback(classifyOpenAIWSReadFallbackReason(readErr), readErr) - } - if clientDisconnected { - break - } - setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(readErr.Error()), "") - return nil, fmt.Errorf("openai ws read event: %w", readErr) - } - - eventType, eventResponseID, responseField := parseOpenAIWSEventEnvelope(message) - if eventType == "" { - continue - } - eventCount++ - if firstEventType == "" { - firstEventType = eventType - } - lastEventType = eventType - - if responseID == "" && eventResponseID != "" { - responseID = eventResponseID - } - - isTokenEvent := isOpenAIWSTokenEvent(eventType) - if isTokenEvent { - tokenEventCount++ - } - isTerminalEvent := isOpenAIWSTerminalEvent(eventType) - if isTerminalEvent { - terminalEventCount++ - } - if firstTokenMs == nil && isTokenEvent { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - if debugEnabled && shouldLogOpenAIWSEvent(eventCount, eventType) { - logOpenAIWSModeDebug( - "event_received account_id=%d conn_id=%s idx=%d type=%s bytes=%d token=%v terminal=%v buffered_pending=%d", - account.ID, - connID, - eventCount, - truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen), - len(message), - isTokenEvent, - isTerminalEvent, - len(bufferedStreamEvents), - ) - } - - if !clientDisconnected { - if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && bytes.Contains(message, mappedModelBytes) { - message = replaceOpenAIWSMessageModel(message, mappedModel, originalModel) - } - if openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(message) { - if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(message); changed { - message = corrected - } - } - } - if openAIWSEventShouldParseUsage(eventType) { - parseOpenAIWSResponseUsageFromCompletedEvent(message, usage) - } - imageCounter.AddSSEData(message) - - if eventType == "response.failed" { - if hit, code, msg := detectOpenAICyberPolicy(message); hit { - MarkOpsCyberPolicy(c, CyberPolicyMark{ - Code: code, - Message: msg, - Body: truncateString(string(message), 4096), - UpstreamStatus: http.StatusOK, - UpstreamInTok: usage.InputTokens, - UpstreamOutTok: usage.OutputTokens, - }) - } - } - - if eventType == "error" { - errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message) - s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw) - errMsg := strings.TrimSpace(errMsgRaw) - if errMsg == "" { - errMsg = "Upstream websocket error" - } - fallbackReason, canFallback := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) - errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) - logOpenAIWSModeInfo( - "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", - account.ID, - connID, - eventCount, - truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen), - canFallback, - errCode, - errType, - errMessage, - ) - if fallbackReason == "previous_response_not_found" { - logOpenAIWSModeInfo( - "previous_response_not_found_diag account_id=%d account_type=%s conn_id=%s previous_response_id=%s previous_response_id_kind=%s response_id=%s event_idx=%d req_stream=%v store_disabled=%v conn_reused=%v session_hash=%s header_session_id=%s header_conversation_id=%s session_id_source=%s conversation_id_source=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v err_code=%s err_type=%s err_message=%s", - account.ID, - account.Type, - connID, - truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(previousResponseIDKind), - truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), - eventCount, - reqStream, - storeDisabled, - lease.Reused(), - truncateOpenAIWSLogValue(sessionHash, 12), - openAIWSHeaderValueForLog(wsHeaders, "session_id"), - openAIWSHeaderValueForLog(wsHeaders, "conversation_id"), - normalizeOpenAIWSLogValue(sessionResolution.SessionSource), - normalizeOpenAIWSLogValue(sessionResolution.ConversationSource), - turnState != "", - len(turnState), - promptCacheKey != "", - errCode, - errType, - errMessage, - ) - } - // error 事件后连接不再可复用，避免回池后污染下一请求。 - lease.MarkBroken() - if !wroteDownstream && canFallback { - return nil, wrapOpenAIWSFallback(fallbackReason, errors.New(errMsg)) - } - statusCode := openAIWSErrorHTTPStatusFromRaw(errCodeRaw, errTypeRaw) - setOpsUpstreamError(c, statusCode, errMsg, "") - if reqStream && !clientDisconnected { - flushBufferedStreamEvents("error_event") - emitStreamMessage(message, true) - } - if !reqStream { - c.JSON(statusCode, gin.H{ - "error": gin.H{ - "type": "upstream_error", - "message": errMsg, - }, - }) - } - return nil, fmt.Errorf("openai ws error event: %s", errMsg) - } - - if reqStream { - // 在首个 token 前先缓冲事件（如 response.created）， - // 以便上游早期断连时仍可安全回退到 HTTP，不给下游发送半截流。 - shouldBuffer := firstTokenMs == nil && !isTokenEvent && !isTerminalEvent - if shouldBuffer { - buffered := make([]byte, len(message)) - copy(buffered, message) - bufferedStreamEvents = append(bufferedStreamEvents, buffered) - bufferedEventCount++ - if debugEnabled && shouldLogOpenAIWSBufferedEvent(bufferedEventCount) { - logOpenAIWSModeDebug( - "buffer_enqueue account_id=%d conn_id=%s idx=%d event_idx=%d event_type=%s buffer_size=%d", - account.ID, - connID, - bufferedEventCount, - eventCount, - truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen), - len(bufferedStreamEvents), - ) - } - } else { - flushBufferedStreamEvents(eventType) - emitStreamMessage(message, isTerminalEvent) - } - } else { - if responseField.Exists() && responseField.Type == gjson.JSON { - finalResponse = []byte(responseField.Raw) - } - } - - if isTerminalEvent { - cleanExit = true - break - } - } - - if !reqStream { - if len(finalResponse) == 0 { - logOpenAIWSModeInfo( - "missing_final_response account_id=%d conn_id=%s events=%d token_events=%d terminal_events=%d wrote_downstream=%v", - account.ID, - connID, - eventCount, - tokenEventCount, - terminalEventCount, - wroteDownstream, - ) - if !wroteDownstream { - return nil, wrapOpenAIWSFallback("missing_final_response", errors.New("no terminal response payload")) - } - return nil, errors.New("ws finished without final response") - } - - if needModelReplace { - finalResponse = s.replaceModelInResponseBody(finalResponse, mappedModel, originalModel) - } - finalResponse = s.correctToolCallsInResponseBody(finalResponse) - populateOpenAIUsageFromResponseJSON(finalResponse, usage) - if responseID == "" { - responseID = strings.TrimSpace(gjson.GetBytes(finalResponse, "id").String()) - } - - c.Data(http.StatusOK, "application/json", finalResponse) - } else { - flushStreamWriter(true) - } - - if responseID != "" && stateStore != nil { - ttl := s.openAIWSResponseStickyTTL() - logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) - stateStore.BindResponseConn(responseID, lease.ConnID(), ttl) - } - if stateStore != nil && storeDisabled && sessionHash != "" { - stateStore.BindSessionConn(groupID, sessionHash, lease.ConnID(), s.openAIWSSessionStickyTTL()) - } - firstTokenMsValue := -1 - if firstTokenMs != nil { - firstTokenMsValue = *firstTokenMs - } - logOpenAIWSModeDebug( - "completed account_id=%d conn_id=%s response_id=%s stream=%v duration_ms=%d events=%d token_events=%d terminal_events=%d buffered_events=%d buffered_flushed=%d first_event=%s last_event=%s first_token_ms=%d wrote_downstream=%v client_disconnected=%v", - account.ID, - connID, - truncateOpenAIWSLogValue(strings.TrimSpace(responseID), openAIWSIDValueMaxLen), - reqStream, - time.Since(startTime).Milliseconds(), - eventCount, - tokenEventCount, - terminalEventCount, - bufferedEventCount, - flushedBufferedEventCount, - truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen), - truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen), - firstTokenMsValue, - wroteDownstream, - clientDisconnected, - ) - - return &OpenAIForwardResult{ - RequestID: responseID, - Usage: *usage, - Model: originalModel, - UpstreamModel: mappedModel, - ImageCount: imageCounter.Count(), - ImageOutputSizes: imageCounter.Sizes(), - ServiceTier: extractOpenAIServiceTier(reqBody), - ReasoningEffort: extractOpenAIReasoningEffort(reqBody, originalModel), - Stream: reqStream, - OpenAIWSMode: true, - ResponseHeaders: lease.HandshakeHeaders(), - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - }, nil -} - -// ProxyResponsesWebSocketFromClient 处理客户端入站 WebSocket（OpenAI Responses WS Mode）并转发到上游。 -// 当前实现按“单请求 -> 终止事件 -> 下一请求”的顺序代理，适配 Codex CLI 的 turn 模式。 -// stripCodexSparkImageGenerationToolFromRawPayload removes the image_generation -// tool from a raw /responses payload when the upstream model is gpt-5.3-codex-spark. -// Spark rejects that tool upstream with HTTP 400 (invalid_request_error, param=tools); -// Codex clients advertise it by default. Returns the (possibly unchanged) payload, -// whether it changed, and any JSON decode error. -func stripCodexSparkImageGenerationToolFromRawPayload(payload []byte, model string) ([]byte, bool, error) { - if !isCodexSparkModel(model) || !openAIRequestBodyHasImageGenerationTool(payload) { - return payload, false, nil - } - return stripOpenAIImageGenerationToolFromRawPayload(payload) -} - -func stripOpenAIImageGenerationToolFromRawPayload(payload []byte) ([]byte, bool, error) { - payloadMap := make(map[string]any) - if err := json.Unmarshal(payload, &payloadMap); err != nil { - return payload, false, err - } - if !stripOpenAIImageGenerationTools(payloadMap) { - return payload, false, nil - } - rebuilt, err := json.Marshal(payloadMap) - if err != nil { - return payload, false, err - } - return rebuilt, true, nil -} - -func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( - ctx context.Context, - c *gin.Context, - clientConn *coderws.Conn, - account *Account, - token string, - firstClientMessage []byte, - hooks *OpenAIWSIngressHooks, -) error { - if s == nil { - return errors.New("service is nil") - } - if c == nil { - return errors.New("gin context is nil") - } - if clientConn == nil { - return errors.New("client websocket is nil") - } - if account == nil { - return errors.New("account is nil") - } - if strings.TrimSpace(token) == "" { - return errors.New("token is empty") - } - - // 预取一次 OpenAI Fast Policy settings，绑定到 ctx，让该 WS session - // 内所有帧的 evaluateOpenAIFastPolicy 调用复用同一份快照，避免每帧 - // 进入 DB / settingRepo。Trade-off 见 withOpenAIFastPolicyContext 注释。 - if s.settingService != nil { - if settings, err := s.settingService.GetOpenAIFastPolicySettings(ctx); err == nil && settings != nil { - ctx = withOpenAIFastPolicyContext(ctx, settings) - } - } - - wsDecision := s.getOpenAIWSProtocolResolver().Resolve(account) - forceHTTPBridge := account.Platform == PlatformGrok - modeRouterV2Enabled := s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.ModeRouterV2Enabled - ingressMode := OpenAIWSIngressModeCtxPool - if modeRouterV2Enabled && !forceHTTPBridge { - ingressMode = account.ResolveOpenAIResponsesWebSocketV2Mode(s.cfg.Gateway.OpenAIWS.IngressModeDefault) - if ingressMode == OpenAIWSIngressModeOff { - return NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - "websocket mode is disabled for this account", - nil, - ) - } - switch ingressMode { - case OpenAIWSIngressModePassthrough: - if wsDecision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { - return fmt.Errorf("websocket ingress requires ws_v2 transport, got=%s", wsDecision.Transport) - } - return s.proxyResponsesWebSocketV2Passthrough( - ctx, - c, - clientConn, - account, - token, - firstClientMessage, - hooks, - wsDecision, - ) - case OpenAIWSIngressModeHTTPBridge: - forceHTTPBridge = true - case OpenAIWSIngressModeCtxPool, OpenAIWSIngressModeShared, OpenAIWSIngressModeDedicated: - // continue - default: - return NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - "websocket mode only supports ctx_pool/passthrough/http_bridge", - nil, - ) - } - } - if !forceHTTPBridge && wsDecision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { - return fmt.Errorf("websocket ingress requires ws_v2 transport, got=%s", wsDecision.Transport) - } - dedicatedMode := modeRouterV2Enabled && ingressMode == OpenAIWSIngressModeDedicated - - wsURL := "" - wsHost := "-" - wsPath := "-" - if forceHTTPBridge { - wsHost = "xai-http-bridge" - wsPath = "/v1/responses" - } else { - var err error - wsURL, err = s.buildOpenAIResponsesWSURL(account) - if err != nil { - return fmt.Errorf("build ws url: %w", err) - } - if parsedURL, parseErr := url.Parse(wsURL); parseErr == nil && parsedURL != nil { - wsHost = normalizeOpenAIWSLogValue(parsedURL.Host) - wsPath = normalizeOpenAIWSLogValue(parsedURL.Path) - } - } - debugEnabled := isOpenAIWSModeDebugEnabled() - isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator")) || (s.cfg != nil && s.cfg.Gateway.ForceCodexCLI) - - type openAIWSClientPayload struct { - payloadRaw []byte - rawForHash []byte - promptCacheKey string - previousResponseID string - originalModel string - imageBillingModel string - imageSizeTier string - imageInputSize string - payloadBytes int - } - ingressSessionOriginalModel := "" - - applyPayloadMutation := func(current []byte, path string, value any) ([]byte, error) { - next, err := sjson.SetBytes(current, path, value) - if err == nil { - return next, nil - } - - // 仅在确实需要修改 payload 且 sjson 失败时，退回 map 路径确保兼容性。 - payload := make(map[string]any) - if unmarshalErr := json.Unmarshal(current, &payload); unmarshalErr != nil { - return nil, err - } - switch path { - case "type", "model": - payload[path] = value - case "client_metadata." + openAIWSTurnMetadataHeader: - setOpenAIWSTurnMetadata(payload, fmt.Sprintf("%v", value)) - default: - return nil, err - } - rebuilt, marshalErr := json.Marshal(payload) - if marshalErr != nil { - return nil, marshalErr - } - return rebuilt, nil - } - - parseClientPayload := func(raw []byte) (openAIWSClientPayload, error) { - trimmed := bytes.TrimSpace(raw) - if len(trimmed) == 0 { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "empty websocket request payload", nil) - } - if !gjson.ValidBytes(trimmed) { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", errors.New("invalid json")) - } - - values := gjson.GetManyBytes(trimmed, "type", "model", "prompt_cache_key", "previous_response_id") - eventType := strings.TrimSpace(values[0].String()) - normalized := trimmed - switch eventType { - case "": - eventType = "response.create" - next, setErr := applyPayloadMutation(normalized, "type", eventType) - if setErr != nil { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", setErr) - } - normalized = next - case "response.create": - case "response.append": - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - "response.append is not supported in ws v2; use response.create with previous_response_id", - nil, - ) - default: - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - fmt.Sprintf("unsupported websocket request type: %s", eventType), - nil, - ) - } - - originalModel := strings.TrimSpace(values[1].String()) - modelMissing := originalModel == "" - if originalModel == "" { - // 入站 WS 长会话里，部分客户端只在第一轮 response.create 上声明 - // model，后续 turn 复用同一 session-level model。为避免因省略 - // model 直接断开用户连接，这里回落到上一轮已通过校验的客户端模型， - // 并在下方写回上游 payload，保证账号模型映射/fast policy/图片权限 - // 仍按同一模型执行。 - originalModel = ingressSessionOriginalModel - if originalModel == "" { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - "model is required in response.create payload", - nil, - ) - } - } - promptCacheKey := strings.TrimSpace(values[2].String()) - previousResponseID := strings.TrimSpace(values[3].String()) - previousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(previousResponseID) - if previousResponseID != "" && previousResponseIDKind == OpenAIPreviousResponseIDKindMessageID { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - "previous_response_id must be a response.id (resp_*), not a message id", - nil, - ) - } - if turnMetadata := strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)); turnMetadata != "" { - next, setErr := applyPayloadMutation(normalized, "client_metadata."+openAIWSTurnMetadataHeader, turnMetadata) - if setErr != nil { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", setErr) - } - normalized = next - } - apiKey := getAPIKeyFromContext(c) - imageGenerationAllowed := GroupAllowsImageGeneration(apiKeyGroup(apiKey)) - codexImageGenerationExplicitToolPolicy := codexImageGenerationExplicitToolPolicyAllow - if isCodexCLI { - codexImageGenerationExplicitToolPolicy = account.CodexImageGenerationExplicitToolPolicy() - } - codexBridgeEnabled := isCodexCLI && imageGenerationAllowed && codexImageGenerationExplicitToolPolicy != codexImageGenerationExplicitToolPolicyStrip && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) - if codexBridgeEnabled { - payloadMap := make(map[string]any) - if err := json.Unmarshal(normalized, &payloadMap); err != nil { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", err) - } - bridgeModified := false - if ensureOpenAIResponsesImageGenerationTool(payloadMap) { - bridgeModified = true - logOpenAIWSModeInfo("ingress_ws_codex_image_tool_injected account_id=%d", account.ID) - } - if ensureOpenAIResponsesImageGenerationToolChoiceAuto(payloadMap) { - bridgeModified = true - logOpenAIWSModeInfo("ingress_ws_codex_image_tool_choice_auto account_id=%d", account.ID) - } - if normalizeOpenAIResponsesImageGenerationTools(payloadMap) { - bridgeModified = true - } - if applyCodexImageGenerationBridgeInstructions(payloadMap) { - bridgeModified = true - logOpenAIWSModeInfo("ingress_ws_codex_image_bridge_instructions_added account_id=%d", account.ID) - } - if bridgeModified { - rebuilt, marshalErr := json.Marshal(payloadMap) - if marshalErr != nil { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", marshalErr) - } - normalized = rebuilt - } - } - upstreamModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)) - if modelMissing || upstreamModel != originalModel { - next, setErr := applyPayloadMutation(normalized, "model", upstreamModel) - if setErr != nil { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", setErr) - } - normalized = next - } - if isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip { - if stripped, changed, stripErr := stripOpenAIImageGenerationToolFromRawPayload(normalized); stripErr != nil { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", stripErr) - } else if changed { - normalized = stripped - logOpenAIWSModeInfo("ingress_ws_codex_image_tool_stripped_by_policy account_id=%d", account.ID) - } - } - if stripped, changed, stripErr := stripCodexSparkImageGenerationToolFromRawPayload(normalized, upstreamModel); stripErr != nil { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", stripErr) - } else if changed { - normalized = stripped - logOpenAIWSModeInfo("ingress_ws_codex_spark_image_tool_stripped account_id=%d", account.ID) - } - imageIntent := IsImageGenerationIntent(openAIResponsesEndpoint, originalModel, normalized) - if imageIntent && !imageGenerationAllowed { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, ImageGenerationPermissionMessage(), nil) - } - imageBillingModel := "" - imageSizeTier := "" - imageInputSize := "" - if imageIntent { - var imageCfgErr error - imageCfg, imageCfgErr := resolveOpenAIResponsesImageBillingConfigDetailedFromBody(normalized, originalModel) - if imageCfgErr != nil { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, imageCfgErr.Error(), imageCfgErr) - } - imageBillingModel = imageCfg.Model - imageSizeTier = imageCfg.SizeTier - imageInputSize = imageCfg.InputSize - } - - // Apply OpenAI Fast Policy on the response.create frame using the same - // evaluator/normalize/scope rules as the HTTP entrypoints. This is the - // single integration point for all WS ingress turns (first + follow-up - // frames flow through here). - // - // Model fallback: first turn still requires model at the handler layer； - // follow-up response.create frames may omit it and then reuse - // ingressSessionOriginalModel. We always write a concrete upstream model - // before evaluating policy, so whitelist / filter behavior remains stable. - policyApplied, blocked, policyErr := s.applyOpenAIFastPolicyToWSResponseCreate(ctx, account, upstreamModel, normalized) - if policyErr != nil { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", policyErr) - } - if blocked != nil { - MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalPolicyDenied) - // Send a Realtime-style error event to the client first, then - // signal the handler to close the connection with PolicyViolation. - // We intentionally do NOT forward this frame upstream. - // - // coder/websocket@v1.8.14 Conn.Write is synchronous and flushes - // the underlying bufio writer before returning (write.go:42 → - // 307-311), and the subsequent close handshake re-acquires the - // same writeFrameMu, so the error event is guaranteed to reach - // the kernel send buffer before any close frame is queued. - eventBytes := buildOpenAIFastPolicyBlockedWSEvent(blocked) - if eventBytes != nil { - writeCtx, cancel := context.WithTimeout(ctx, s.openAIWSWriteTimeout()) - _ = clientConn.Write(writeCtx, coderws.MessageText, eventBytes) - cancel() - } - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - blocked.Message, - blocked, - ) - } - normalized = policyApplied - ingressSessionOriginalModel = originalModel - - return openAIWSClientPayload{ - payloadRaw: normalized, - rawForHash: trimmed, - promptCacheKey: promptCacheKey, - previousResponseID: previousResponseID, - originalModel: originalModel, - imageBillingModel: imageBillingModel, - imageSizeTier: imageSizeTier, - imageInputSize: imageInputSize, - payloadBytes: len(normalized), - }, nil - } - - writeClientMessage := func(message []byte) error { - writeCtx, cancel := context.WithTimeout(ctx, s.openAIWSWriteTimeout()) - defer cancel() - return clientConn.Write(writeCtx, coderws.MessageText, message) - } - - readClientMessage := func() ([]byte, error) { - msgType, payload, readErr := clientConn.Read(ctx) - if readErr != nil { - return nil, readErr - } - if msgType != coderws.MessageText && msgType != coderws.MessageBinary { - return nil, NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - fmt.Sprintf("unsupported websocket client message type: %s", msgType.String()), - nil, - ) - } - return payload, nil - } - - firstPayload, err := parseClientPayload(firstClientMessage) - if err != nil { - return err - } - - turnState := strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader)) - stateStore := s.getOpenAIWSStateStore() - groupID := getOpenAIGroupIDFromContext(c) - storeDisabledConnMode := s.openAIWSStoreDisabledConnMode() - sessionHash := "" - preferredConnID := "" - storeDisabled := false - refreshIngressRouteState := func(payload openAIWSClientPayload) { - sessionHash = s.GenerateSessionHash(c, payload.rawForHash) - if turnState == "" && stateStore != nil && sessionHash != "" { - if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { - turnState = savedTurnState - } - } - - preferredConnID = "" - if stateStore != nil && payload.previousResponseID != "" { - if connID, ok := stateStore.GetResponseConn(payload.previousResponseID); ok { - preferredConnID = connID - } - } - - storeDisabled = s.isOpenAIWSStoreDisabledInRequestRaw(payload.payloadRaw, account) - if stateStore != nil && storeDisabled && payload.previousResponseID == "" && sessionHash != "" { - if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok { - preferredConnID = connID - } - } - } - refreshIngressRouteState(firstPayload) - - if forceHTTPBridge || s.shouldBridgeOpenAIWSHTTP(account, firstPayload.payloadBytes, firstPayload.previousResponseID) { - logOpenAIWSModeInfo( - "ingress_ws_http_bridge_start account_id=%d account_type=%s payload_bytes=%d threshold_bytes=%d has_session_hash=%v store_disabled=%v", - account.ID, - account.Type, - firstPayload.payloadBytes, - s.openAIWSHTTPBridgeThresholdBytes(), - sessionHash != "", - storeDisabled, - ) - currentBridgePayload := firstPayload - var bridgeReplayInput []json.RawMessage - bridgeReplayInputExists := false - for turn := 1; ; turn++ { - if turn > 1 && hooks != nil && hooks.BeforeRequest != nil { - if err := hooks.BeforeRequest(turn, currentBridgePayload.payloadRaw, currentBridgePayload.originalModel); err != nil { - return err - } - } - if hooks != nil && hooks.BeforeTurn != nil { - if err := hooks.BeforeTurn(turn); err != nil { - return err - } - } - if turnState != "" && c != nil && c.Request != nil { - c.Request.Header.Set(openAIWSTurnStateHeader, turnState) - } - bridgePayloadRaw := currentBridgePayload.payloadRaw - bridgePayloadBytes := currentBridgePayload.payloadBytes - needsBridgeReplay := currentBridgePayload.previousResponseID != "" || openAIWSRawPayloadHasToolCallOutput(currentBridgePayload.payloadRaw) - turnReplayInput, turnReplayInputExists, replayInputErr := buildOpenAIWSReplayInputSequence( - bridgeReplayInput, - bridgeReplayInputExists, - currentBridgePayload.payloadRaw, - needsBridgeReplay, - ) - if replayInputErr != nil { - return fmt.Errorf("build websocket http bridge replay input: %w", replayInputErr) - } - if needsBridgeReplay && turnReplayInputExists { - updatedPayload, setInputErr := setOpenAIWSPayloadInputSequence( - currentBridgePayload.payloadRaw, - turnReplayInput, - true, - ) - if setInputErr != nil { - return fmt.Errorf("set websocket http bridge replay input: %w", setInputErr) - } - bridgePayloadRaw = updatedPayload - bridgePayloadBytes = len(updatedPayload) - logOpenAIWSModeInfo( - "ingress_ws_http_bridge_replay_input account_id=%d turn=%d input_items=%d previous_response_id_present=%v has_tool_output=%v", - account.ID, - turn, - len(turnReplayInput), - currentBridgePayload.previousResponseID != "", - openAIWSRawPayloadHasToolCallOutput(currentBridgePayload.payloadRaw), - ) - } - result, bridgeErr := s.proxyOpenAIWSHTTPBridgeTurn( - ctx, - c, - account, - token, - bridgePayloadRaw, - bridgePayloadBytes, - currentBridgePayload.originalModel, - currentBridgePayload.imageBillingModel, - currentBridgePayload.imageSizeTier, - currentBridgePayload.imageInputSize, - turn, - writeClientMessage, - ) - if hooks != nil && hooks.AfterTurn != nil { - hooks.AfterTurn(turn, result, bridgeErr) - } - if bridgeErr != nil { - return bridgeErr - } - if result == nil { - return errors.New("websocket http bridge turn result is nil") - } - bridgeReplayInput = cloneOpenAIWSRawMessages(turnReplayInput) - bridgeReplayInputExists = turnReplayInputExists - if result.wsReplayInputExists { - bridgeReplayInput = append(bridgeReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...) - bridgeReplayInputExists = true - } - if bridgeTurnState := strings.TrimSpace(result.ResponseHeaders.Get(openAIWSTurnStateHeader)); bridgeTurnState != "" { - turnState = bridgeTurnState - if stateStore != nil && sessionHash != "" { - stateStore.BindSessionTurnState(groupID, sessionHash, bridgeTurnState, s.openAIWSSessionStickyTTL()) - } - } - responseID := strings.TrimSpace(result.RequestID) - if responseID != "" && stateStore != nil { - ttl := s.openAIWSResponseStickyTTL() - logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) - } - nextClientMessage, readErr := readClientMessage() - if readErr != nil { - if isOpenAIWSClientDisconnectError(readErr) { - closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) - logOpenAIWSModeInfo( - "ingress_ws_http_bridge_client_closed account_id=%d close_status=%s close_reason=%s", - account.ID, - closeStatus, - truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), - ) - return nil - } - return fmt.Errorf("read client websocket request: %w", readErr) - } - nextPayload, parseErr := parseClientPayload(nextClientMessage) - if parseErr != nil { - return parseErr - } - currentBridgePayload = nextPayload - } - } - - wsHeaders, _, buildHdrErr := s.buildOpenAIWSHeaders(ctx, c, account, token, wsDecision, isCodexCLI, turnState, strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)), firstPayload.promptCacheKey) - if buildHdrErr != nil { - return fmt.Errorf("build ws headers: %w", buildHdrErr) - } - baseAcquireReq := openAIWSAcquireRequest{ - Account: account, - WSURL: wsURL, - Headers: wsHeaders, - ProxyURL: func() string { - if account.ProxyID != nil && account.Proxy != nil { - return account.Proxy.URL() - } - return "" - }(), - ForceNewConn: false, - } - pool := s.getOpenAIWSConnPool() - if pool == nil { - return errors.New("openai ws conn pool is nil") - } - - logOpenAIWSModeInfo( - "ingress_ws_protocol_confirm account_id=%d account_type=%s transport=%s ws_host=%s ws_path=%s ws_mode=%s store_disabled=%v has_session_hash=%v has_previous_response_id=%v", - account.ID, - account.Type, - normalizeOpenAIWSLogValue(string(wsDecision.Transport)), - wsHost, - wsPath, - normalizeOpenAIWSLogValue(ingressMode), - storeDisabled, - sessionHash != "", - firstPayload.previousResponseID != "", - ) - - if debugEnabled { - logOpenAIWSModeDebug( - "ingress_ws_start account_id=%d account_type=%s transport=%s ws_host=%s preferred_conn_id=%s has_session_hash=%v has_previous_response_id=%v store_disabled=%v", - account.ID, - account.Type, - normalizeOpenAIWSLogValue(string(wsDecision.Transport)), - wsHost, - truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), - sessionHash != "", - firstPayload.previousResponseID != "", - storeDisabled, - ) - } - if firstPayload.previousResponseID != "" { - firstPreviousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(firstPayload.previousResponseID) - logOpenAIWSModeInfo( - "ingress_ws_continuation_probe account_id=%d turn=%d previous_response_id=%s previous_response_id_kind=%s preferred_conn_id=%s session_hash=%s header_session_id=%s header_conversation_id=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v store_disabled=%v", - account.ID, - 1, - truncateOpenAIWSLogValue(firstPayload.previousResponseID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(firstPreviousResponseIDKind), - truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(sessionHash, 12), - openAIWSHeaderValueForLog(baseAcquireReq.Headers, "session_id"), - openAIWSHeaderValueForLog(baseAcquireReq.Headers, "conversation_id"), - turnState != "", - len(turnState), - firstPayload.promptCacheKey != "", - storeDisabled, - ) - } - - acquireTimeout := s.openAIWSAcquireTimeout() - if acquireTimeout <= 0 { - acquireTimeout = 30 * time.Second - } - - acquireTurnLease := func(turn int, preferred string, forcePreferredConn bool) (*openAIWSConnLease, error) { - req := cloneOpenAIWSAcquireRequest(baseAcquireReq) - req.PreferredConnID = strings.TrimSpace(preferred) - req.ForcePreferredConn = forcePreferredConn - // dedicated 模式下每次获取均新建连接，避免跨会话复用残留上下文。 - req.ForceNewConn = dedicatedMode - acquireCtx, acquireCancel := context.WithTimeout(ctx, acquireTimeout) - lease, acquireErr := pool.Acquire(acquireCtx, req) - acquireCancel() - if acquireErr != nil { - dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(acquireErr) - logOpenAIWSModeInfo( - "ingress_ws_upstream_acquire_fail account_id=%d turn=%d reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_preferred_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v", - account.ID, - turn, - normalizeOpenAIWSLogValue(classifyOpenAIWSAcquireError(acquireErr)), - dialStatus, - dialClass, - dialCloseStatus, - truncateOpenAIWSLogValue(dialCloseReason, openAIWSHeaderValueMaxLen), - dialRespServer, - dialRespVia, - dialRespCFRay, - dialRespReqID, - truncateOpenAIWSLogValue(acquireErr.Error(), openAIWSLogValueMaxLen), - truncateOpenAIWSLogValue(preferred, openAIWSIDValueMaxLen), - forcePreferredConn, - wsHost, - wsPath, - account.ProxyID != nil && account.Proxy != nil, - ) - var dialErr *openAIWSDialError - if errors.As(acquireErr, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusTooManyRequests { - s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(acquireErr.Error())) - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusTooManyRequests, - ResponseHeaders: cloneHeader(dialErr.ResponseHeaders), - } - } - if errors.Is(acquireErr, errOpenAIWSPreferredConnUnavailable) { - return nil, NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - "upstream continuation connection is unavailable; please restart the conversation", - acquireErr, - ) - } - if errors.Is(acquireErr, context.DeadlineExceeded) || errors.Is(acquireErr, errOpenAIWSConnQueueFull) { - return nil, NewOpenAIWSClientCloseError( - coderws.StatusTryAgainLater, - "upstream websocket is busy, please retry later", - acquireErr, - ) - } - return nil, acquireErr - } - connID := strings.TrimSpace(lease.ConnID()) - if handshakeTurnState := strings.TrimSpace(lease.HandshakeHeader(openAIWSTurnStateHeader)); handshakeTurnState != "" { - turnState = handshakeTurnState - if stateStore != nil && sessionHash != "" { - stateStore.BindSessionTurnState(groupID, sessionHash, handshakeTurnState, s.openAIWSSessionStickyTTL()) - } - updatedHeaders := cloneHeader(baseAcquireReq.Headers) - if updatedHeaders == nil { - updatedHeaders = make(http.Header) - } - updatedHeaders.Set(openAIWSTurnStateHeader, handshakeTurnState) - baseAcquireReq.Headers = updatedHeaders - } - logOpenAIWSModeInfo( - "ingress_ws_upstream_connected account_id=%d turn=%d conn_id=%s conn_reused=%v conn_pick_ms=%d queue_wait_ms=%d preferred_conn_id=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - lease.Reused(), - lease.ConnPickDuration().Milliseconds(), - lease.QueueWaitDuration().Milliseconds(), - truncateOpenAIWSLogValue(preferred, openAIWSIDValueMaxLen), - ) - return lease, nil - } - - sendAndRelay := func(turn int, lease *openAIWSConnLease, payload []byte, payloadBytes int, originalModel string, imageBillingModel string, imageSizeTier string, imageInputSize string) (*OpenAIForwardResult, error) { - if lease == nil { - return nil, errors.New("upstream websocket lease is nil") - } - turnStart := time.Now() - wroteDownstream := false - if err := lease.WriteJSONWithContextTimeout(ctx, json.RawMessage(payload), s.openAIWSWriteTimeout()); err != nil { - return nil, wrapOpenAIWSIngressTurnError( - "write_upstream", - fmt.Errorf("write upstream websocket request: %w", err), - false, - ) - } - if debugEnabled { - logOpenAIWSModeDebug( - "ingress_ws_turn_request_sent account_id=%d turn=%d conn_id=%s payload_bytes=%d", - account.ID, - turn, - truncateOpenAIWSLogValue(lease.ConnID(), openAIWSIDValueMaxLen), - payloadBytes, - ) - } - - responseID := "" - usage := OpenAIUsage{} - imageCounter := newOpenAIImageOutputCounter() - var firstTokenMs *int - reqStream := openAIWSPayloadBoolFromRaw(payload, "stream", true) - turnPreviousResponseID := openAIWSPayloadStringFromRaw(payload, "previous_response_id") - turnPreviousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(turnPreviousResponseID) - turnPromptCacheKey := openAIWSPayloadStringFromRaw(payload, "prompt_cache_key") - turnStoreDisabled := s.isOpenAIWSStoreDisabledInRequestRaw(payload, account) - turnHasFunctionCallOutput := openAIWSRawPayloadHasToolCallOutput(payload) - eventCount := 0 - tokenEventCount := 0 - terminalEventCount := 0 - replayCollector := &openAIWSToolCallReplayCollector{} - firstEventType := "" - lastEventType := "" - needModelReplace := false - clientDisconnected := false - mappedModel := "" - var mappedModelBytes []byte - if originalModel != "" { - mappedModel = normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)) - needModelReplace = mappedModel != "" && mappedModel != originalModel - if needModelReplace { - mappedModelBytes = []byte(mappedModel) - } - } - for { - upstreamMessage, readErr := lease.ReadMessageWithContextTimeout(ctx, s.openAIWSReadTimeout()) - if readErr != nil { - lease.MarkBroken() - return nil, wrapOpenAIWSIngressTurnError( - "read_upstream", - fmt.Errorf("read upstream websocket event: %w", readErr), - wroteDownstream, - ) - } - - eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage) - if responseID == "" && eventResponseID != "" { - responseID = eventResponseID - } - if eventType != "" { - eventCount++ - if firstEventType == "" { - firstEventType = eventType - } - lastEventType = eventType - } - if eventType == "error" { - errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) - s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw) - fallbackReason, _ := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) - errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) - recoverablePrevNotFound := fallbackReason == openAIWSIngressStagePreviousResponseNotFound && - turnPreviousResponseID != "" && - !turnHasFunctionCallOutput && - s.openAIWSIngressPreviousResponseRecoveryEnabled() && - !wroteDownstream - if recoverablePrevNotFound { - // 可恢复场景使用非 error 关键字日志，避免被 LegacyPrintf 误判为 ERROR 级别。 - logOpenAIWSModeInfo( - "ingress_ws_prev_response_recoverable account_id=%d turn=%d conn_id=%s idx=%d reason=%s code=%s type=%s message=%s previous_response_id=%s previous_response_id_kind=%s response_id=%s store_disabled=%v has_prompt_cache_key=%v", - account.ID, - turn, - truncateOpenAIWSLogValue(lease.ConnID(), openAIWSIDValueMaxLen), - eventCount, - truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen), - errCode, - errType, - errMessage, - truncateOpenAIWSLogValue(turnPreviousResponseID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(turnPreviousResponseIDKind), - truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), - turnStoreDisabled, - turnPromptCacheKey != "", - ) - } else { - logOpenAIWSModeInfo( - "ingress_ws_error_event account_id=%d turn=%d conn_id=%s idx=%d fallback_reason=%s err_code=%s err_type=%s err_message=%s previous_response_id=%s previous_response_id_kind=%s response_id=%s store_disabled=%v has_prompt_cache_key=%v", - account.ID, - turn, - truncateOpenAIWSLogValue(lease.ConnID(), openAIWSIDValueMaxLen), - eventCount, - truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen), - errCode, - errType, - errMessage, - truncateOpenAIWSLogValue(turnPreviousResponseID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(turnPreviousResponseIDKind), - truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), - turnStoreDisabled, - turnPromptCacheKey != "", - ) - } - // previous_response_not_found 在 ingress 模式支持单次恢复重试： - // 不把该 error 直接下发客户端，而是由上层去掉 previous_response_id 后重放当前 turn。 - if recoverablePrevNotFound { - lease.MarkBroken() - errMsg := strings.TrimSpace(errMsgRaw) - if errMsg == "" { - errMsg = "previous response not found" - } - return nil, wrapOpenAIWSIngressTurnError( - openAIWSIngressStagePreviousResponseNotFound, - errors.New(errMsg), - false, - ) - } - if !wroteDownstream && isOpenAIWSRateLimitError(errCodeRaw, errTypeRaw, errMsgRaw) { - lease.MarkBroken() - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusTooManyRequests, - ResponseBody: append([]byte(nil), upstreamMessage...), - ResponseHeaders: cloneHeader(lease.HandshakeHeaders()), - } - } - } - isTokenEvent := isOpenAIWSTokenEvent(eventType) - if isTokenEvent { - tokenEventCount++ - } - isTerminalEvent := isOpenAIWSTerminalEvent(eventType) - if isTerminalEvent { - terminalEventCount++ - } - if firstTokenMs == nil && isTokenEvent { - ms := int(time.Since(turnStart).Milliseconds()) - firstTokenMs = &ms - } - if openAIWSEventShouldParseUsage(eventType) { - parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage) - } - imageCounter.AddSSEData(upstreamMessage) - - if eventType == "response.failed" { - if hit, code, msg := detectOpenAICyberPolicy(upstreamMessage); hit { - MarkOpsCyberPolicy(c, CyberPolicyMark{ - Code: code, - Message: msg, - Body: truncateString(string(upstreamMessage), 4096), - UpstreamStatus: http.StatusOK, - UpstreamInTok: usage.InputTokens, - UpstreamOutTok: usage.OutputTokens, - }) - } - } - - if !clientDisconnected { - if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && bytes.Contains(upstreamMessage, mappedModelBytes) { - upstreamMessage = replaceOpenAIWSMessageModel(upstreamMessage, mappedModel, originalModel) - } - if openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(upstreamMessage) { - if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(upstreamMessage); changed { - upstreamMessage = corrected - } - } - replayCollector.AddEvent(eventType, upstreamMessage) - if err := writeClientMessage(upstreamMessage); err != nil { - if isOpenAIWSClientDisconnectError(err) { - clientDisconnected = true - closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err) - logOpenAIWSModeInfo( - "ingress_ws_client_disconnected_drain account_id=%d turn=%d conn_id=%s close_status=%s close_reason=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(lease.ConnID(), openAIWSIDValueMaxLen), - closeStatus, - truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), - ) - } else { - return nil, wrapOpenAIWSIngressTurnError( - "write_client", - fmt.Errorf("write client websocket event: %w", err), - wroteDownstream, - ) - } - } else { - wroteDownstream = true - } - } - if isTerminalEvent { - // 客户端已断连时，上游连接的 session 状态不可信，标记 broken 避免回池复用。 - if clientDisconnected { - lease.MarkBroken() - } - firstTokenMsValue := -1 - if firstTokenMs != nil { - firstTokenMsValue = *firstTokenMs - } - if debugEnabled { - logOpenAIWSModeDebug( - "ingress_ws_turn_completed account_id=%d turn=%d conn_id=%s response_id=%s 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(lease.ConnID(), openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), - time.Since(turnStart).Milliseconds(), - eventCount, - tokenEventCount, - terminalEventCount, - truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen), - truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen), - firstTokenMsValue, - clientDisconnected, - ) - } - imageCount := imageCounter.Count() - result := &OpenAIForwardResult{ - RequestID: responseID, - Usage: usage, - Model: originalModel, - UpstreamModel: mappedModel, - ServiceTier: extractOpenAIServiceTierFromBody(payload), - ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, originalModel), payload, mappedModel), - Stream: reqStream, - OpenAIWSMode: true, - ResponseHeaders: lease.HandshakeHeaders(), - Duration: time.Since(turnStart), - FirstTokenMs: firstTokenMs, - } - if replayInput := replayCollector.Items(); len(replayInput) > 0 { - result.wsReplayInput = replayInput - result.wsReplayInputExists = true - } - if imageCount > 0 { - result.ImageCount = imageCount - result.ImageSize = imageSizeTier - result.ImageInputSize = imageInputSize - result.ImageOutputSizes = imageCounter.Sizes() - result.BillingModel = imageBillingModel - } - return result, nil - } - } - } - - currentPayload := firstPayload.payloadRaw - currentOriginalModel := firstPayload.originalModel - currentImageBillingModel := firstPayload.imageBillingModel - currentImageSizeTier := firstPayload.imageSizeTier - currentImageInputSize := firstPayload.imageInputSize - currentPayloadBytes := firstPayload.payloadBytes - isStrictAffinityTurn := func(payload []byte) bool { - if !storeDisabled { - return false - } - return strings.TrimSpace(openAIWSPayloadStringFromRaw(payload, "previous_response_id")) != "" - } - var sessionLease *openAIWSConnLease - sessionConnID := "" - pinnedSessionConnID := "" - unpinSessionConn := func(connID string) { - connID = strings.TrimSpace(connID) - if connID == "" || pinnedSessionConnID != connID { - return - } - pool.UnpinConn(account.ID, connID) - pinnedSessionConnID = "" - } - pinSessionConn := func(connID string) { - if !storeDisabled { - return - } - connID = strings.TrimSpace(connID) - if connID == "" || pinnedSessionConnID == connID { - return - } - if pinnedSessionConnID != "" { - pool.UnpinConn(account.ID, pinnedSessionConnID) - pinnedSessionConnID = "" - } - if pool.PinConn(account.ID, connID) { - pinnedSessionConnID = connID - } - } - // lastTurnClean 标记最后一轮 sendAndRelay 是否正常完成（收到终端事件且客户端未断连）。 - // 所有异常路径（读写错误、error 事件、客户端断连）已在各自分支或上层（L3403）中 MarkBroken， - // 因此 releaseSessionLease 中只需在非正常结束时 MarkBroken。 - lastTurnClean := false - releaseSessionLease := func() { - if sessionLease == nil { - return - } - if !lastTurnClean { - sessionLease.MarkBroken() - } - unpinSessionConn(sessionConnID) - sessionLease.Release() - if debugEnabled { - logOpenAIWSModeDebug( - "ingress_ws_upstream_released account_id=%d conn_id=%s", - account.ID, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - ) - } - } - defer releaseSessionLease() - - turn := 1 - turnRetry := 0 - turnPrevRecoveryTried := false - lastTurnFinishedAt := time.Time{} - lastTurnResponseID := "" - lastTurnPayload := []byte(nil) - var lastTurnStrictState *openAIWSIngressPreviousTurnStrictState - lastTurnReplayInput := []json.RawMessage(nil) - lastTurnReplayInputExists := false - currentTurnReplayInput := []json.RawMessage(nil) - currentTurnReplayInputExists := false - skipBeforeTurn := false - hasCurrentOrReplayFunctionCallOutput := func(payload []byte) bool { - if openAIWSRawPayloadHasToolCallOutput(payload) { - return true - } - return currentTurnReplayInputExists && openAIWSRawItemsHasFunctionCallOutput(currentTurnReplayInput) - } - resetSessionLease := func(markBroken bool) { - if sessionLease == nil { - return - } - if markBroken { - sessionLease.MarkBroken() - } - releaseSessionLease() - sessionLease = nil - sessionConnID = "" - preferredConnID = "" - } - recoverIngressPrevResponseNotFound := func(relayErr error, turn int, connID string) bool { - if !isOpenAIWSIngressPreviousResponseNotFound(relayErr) { - return false - } - if turnPrevRecoveryTried || !s.openAIWSIngressPreviousResponseRecoveryEnabled() { - return false - } - // 携带 function_call_output 的请求不能丢弃 previous_response_id： - // 上游 API 需要 response chain 来匹配 tool_result 与之前的 tool_use， - // 丢弃后会导致 "No tool call found for function call output" 400 错误。 - if hasCurrentOrReplayFunctionCallOutput(currentPayload) { - return false - } - if isStrictAffinityTurn(currentPayload) { - // Layer 2：严格亲和链路命中 previous_response_not_found 时，降级为“去掉 previous_response_id 后重放一次”。 - // 该错误说明续链锚点已失效，继续 strict fail-close 只会直接中断本轮请求。 - logOpenAIWSModeInfo( - "ingress_ws_prev_response_recovery_layer2 account_id=%d turn=%d conn_id=%s store_disabled_conn_mode=%s action=drop_previous_response_id_retry", - account.ID, - turn, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(storeDisabledConnMode), - ) - } - turnPrevRecoveryTried = true - updatedPayload, removed, dropErr := dropPreviousResponseIDFromRawPayload(currentPayload) - if dropErr != nil || !removed { - reason := "not_removed" - if dropErr != nil { - reason = "drop_error" - } - logOpenAIWSModeInfo( - "ingress_ws_prev_response_recovery_skip account_id=%d turn=%d conn_id=%s reason=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(reason), - ) - return false - } - updatedWithInput, setInputErr := setOpenAIWSPayloadInputSequence( - updatedPayload, - currentTurnReplayInput, - currentTurnReplayInputExists, - ) - if setInputErr != nil { - logOpenAIWSModeInfo( - "ingress_ws_prev_response_recovery_skip account_id=%d turn=%d conn_id=%s reason=set_full_input_error cause=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(setInputErr.Error(), openAIWSLogValueMaxLen), - ) - return false - } - logOpenAIWSModeInfo( - "ingress_ws_prev_response_recovery account_id=%d turn=%d conn_id=%s action=drop_previous_response_id retry=1", - account.ID, - turn, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - ) - currentPayload = updatedWithInput - currentPayloadBytes = len(updatedWithInput) - resetSessionLease(true) - skipBeforeTurn = true - return true - } - retryIngressTurn := func(relayErr error, turn int, connID string) bool { - if !isOpenAIWSIngressTurnRetryable(relayErr) || turnRetry >= 1 { - return false - } - if isStrictAffinityTurn(currentPayload) { - logOpenAIWSModeInfo( - "ingress_ws_turn_retry_skip account_id=%d turn=%d conn_id=%s reason=strict_affinity", - account.ID, - turn, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - ) - return false - } - turnRetry++ - logOpenAIWSModeInfo( - "ingress_ws_turn_retry account_id=%d turn=%d retry=%d reason=%s conn_id=%s", - account.ID, - turn, - turnRetry, - truncateOpenAIWSLogValue(openAIWSIngressTurnRetryReason(relayErr), openAIWSLogValueMaxLen), - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - ) - resetSessionLease(true) - skipBeforeTurn = true - return true - } - for { - if turn > 1 && !skipBeforeTurn && hooks != nil && hooks.BeforeRequest != nil { - if err := hooks.BeforeRequest(turn, currentPayload, currentOriginalModel); err != nil { - return err - } - } - if !skipBeforeTurn && hooks != nil && hooks.BeforeTurn != nil { - if err := hooks.BeforeTurn(turn); err != nil { - return err - } - } - skipBeforeTurn = false - currentPreviousResponseID := openAIWSPayloadStringFromRaw(currentPayload, "previous_response_id") - expectedPrev := strings.TrimSpace(lastTurnResponseID) - toolSignals := ToolContinuationSignals{ - HasFunctionCallOutput: openAIWSRawPayloadHasToolCallOutput(currentPayload), - } - if toolSignals.HasFunctionCallOutput { - var currentReqBody map[string]any - if err := json.Unmarshal(currentPayload, ¤tReqBody); err == nil { - toolSignals = AnalyzeToolContinuationSignals(currentReqBody) - } - } - hasFunctionCallOutput := toolSignals.HasFunctionCallOutput - // store=false + function_call_output 场景必须有续链锚点。 - // 若客户端未传 previous_response_id，优先回填上一轮响应 ID，避免上游报 call_id 无法关联。 - if shouldInferIngressFunctionCallOutputPreviousResponseID( - storeDisabled, - turn, - toolSignals, - currentPreviousResponseID, - expectedPrev, - ) { - updatedPayload, setPrevErr := setPreviousResponseIDToRawPayload(currentPayload, expectedPrev) - if setPrevErr != nil { - logOpenAIWSModeInfo( - "ingress_ws_function_call_output_prev_infer_skip account_id=%d turn=%d conn_id=%s reason=set_previous_response_id_error cause=%s expected_previous_response_id=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(setPrevErr.Error(), openAIWSLogValueMaxLen), - truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), - ) - } else { - currentPayload = updatedPayload - currentPayloadBytes = len(updatedPayload) - currentPreviousResponseID = expectedPrev - logOpenAIWSModeInfo( - "ingress_ws_function_call_output_prev_infer account_id=%d turn=%d conn_id=%s action=set_previous_response_id previous_response_id=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), - ) - } - } - nextReplayInput, nextReplayInputExists, replayInputErr := buildOpenAIWSReplayInputSequence( - lastTurnReplayInput, - lastTurnReplayInputExists, - currentPayload, - currentPreviousResponseID != "", - ) - if replayInputErr != nil { - logOpenAIWSModeInfo( - "ingress_ws_replay_input_skip account_id=%d turn=%d conn_id=%s reason=build_error cause=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(replayInputErr.Error(), openAIWSLogValueMaxLen), - ) - currentTurnReplayInput = nil - currentTurnReplayInputExists = false - } else { - currentTurnReplayInput = nextReplayInput - currentTurnReplayInputExists = nextReplayInputExists - } - replayHasFunctionCallOutput := currentTurnReplayInputExists && - openAIWSRawItemsHasFunctionCallOutput(currentTurnReplayInput) - hasFunctionCallOutput = hasFunctionCallOutput || replayHasFunctionCallOutput - if storeDisabled && turn > 1 && currentPreviousResponseID != "" { - shouldKeepPreviousResponseID := false - strictReason := "" - var strictErr error - if lastTurnStrictState != nil { - shouldKeepPreviousResponseID, strictReason, strictErr = shouldKeepIngressPreviousResponseIDWithStrictState( - lastTurnStrictState, - currentPayload, - lastTurnResponseID, - hasFunctionCallOutput, - ) - } else { - shouldKeepPreviousResponseID, strictReason, strictErr = shouldKeepIngressPreviousResponseID( - lastTurnPayload, - currentPayload, - lastTurnResponseID, - hasFunctionCallOutput, - ) - } - if strictErr != nil { - logOpenAIWSModeInfo( - "ingress_ws_prev_response_strict_eval account_id=%d turn=%d conn_id=%s action=keep_previous_response_id reason=%s cause=%s previous_response_id=%s expected_previous_response_id=%s has_function_call_output=%v", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(strictReason), - truncateOpenAIWSLogValue(strictErr.Error(), openAIWSLogValueMaxLen), - truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), - hasFunctionCallOutput, - ) - } else if !shouldKeepPreviousResponseID { - updatedPayload, removed, dropErr := dropPreviousResponseIDFromRawPayload(currentPayload) - if dropErr != nil || !removed { - dropReason := "not_removed" - if dropErr != nil { - dropReason = "drop_error" - } - logOpenAIWSModeInfo( - "ingress_ws_prev_response_strict_eval account_id=%d turn=%d conn_id=%s action=keep_previous_response_id reason=%s drop_reason=%s previous_response_id=%s expected_previous_response_id=%s has_function_call_output=%v", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(strictReason), - normalizeOpenAIWSLogValue(dropReason), - truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), - hasFunctionCallOutput, - ) - } else { - updatedWithInput, setInputErr := setOpenAIWSPayloadInputSequence( - updatedPayload, - currentTurnReplayInput, - currentTurnReplayInputExists, - ) - if setInputErr != nil { - logOpenAIWSModeInfo( - "ingress_ws_prev_response_strict_eval account_id=%d turn=%d conn_id=%s action=keep_previous_response_id reason=%s drop_reason=set_full_input_error previous_response_id=%s expected_previous_response_id=%s cause=%s has_function_call_output=%v", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(strictReason), - truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(setInputErr.Error(), openAIWSLogValueMaxLen), - hasFunctionCallOutput, - ) - } else { - currentPayload = updatedWithInput - currentPayloadBytes = len(updatedWithInput) - logOpenAIWSModeInfo( - "ingress_ws_prev_response_strict_eval account_id=%d turn=%d conn_id=%s action=drop_previous_response_id_full_create reason=%s previous_response_id=%s expected_previous_response_id=%s has_function_call_output=%v", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(strictReason), - truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), - hasFunctionCallOutput, - ) - currentPreviousResponseID = "" - } - } - } - } - forcePreferredConn := isStrictAffinityTurn(currentPayload) - if sessionLease == nil { - acquiredLease, acquireErr := acquireTurnLease(turn, preferredConnID, forcePreferredConn) - if acquireErr != nil { - return fmt.Errorf("acquire upstream websocket: %w", acquireErr) - } - sessionLease = acquiredLease - sessionConnID = strings.TrimSpace(sessionLease.ConnID()) - if storeDisabled { - pinSessionConn(sessionConnID) - } else { - unpinSessionConn(sessionConnID) - } - } - shouldPreflightPing := turn > 1 && sessionLease != nil && turnRetry == 0 - if shouldPreflightPing && openAIWSIngressPreflightPingIdle > 0 && !lastTurnFinishedAt.IsZero() { - if time.Since(lastTurnFinishedAt) < openAIWSIngressPreflightPingIdle { - shouldPreflightPing = false - } - } - if shouldPreflightPing { - if pingErr := sessionLease.PingWithTimeout(openAIWSConnHealthCheckTO); pingErr != nil { - logOpenAIWSModeInfo( - "ingress_ws_upstream_preflight_ping_fail account_id=%d turn=%d conn_id=%s cause=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(pingErr.Error(), openAIWSLogValueMaxLen), - ) - if forcePreferredConn { - // 携带 function_call_output 的请求不能丢弃 previous_response_id： - // 上游 API 需要 response chain 来匹配 tool_result 与之前的 tool_use， - // 除非 replay input 已经包含与每个 tool_result 匹配的 tool_use 上下文。 - hasFCOutput := hasFunctionCallOutput - hasReplayToolContext := hasFCOutput && - currentTurnReplayInputExists && - openAIWSRawItemsHaveToolCallContextForOutputs(currentTurnReplayInput) - if !turnPrevRecoveryTried && currentPreviousResponseID != "" && (!hasFCOutput || hasReplayToolContext) { - updatedPayload, removed, dropErr := dropPreviousResponseIDFromRawPayload(currentPayload) - if dropErr != nil || !removed { - reason := "not_removed" - if dropErr != nil { - reason = "drop_error" - } - logOpenAIWSModeInfo( - "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=%s previous_response_id=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(reason), - truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), - ) - } else { - updatedWithInput, setInputErr := setOpenAIWSPayloadInputSequence( - updatedPayload, - currentTurnReplayInput, - currentTurnReplayInputExists, - ) - if setInputErr != nil { - logOpenAIWSModeInfo( - "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=set_full_input_error previous_response_id=%s cause=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(setInputErr.Error(), openAIWSLogValueMaxLen), - ) - } else { - logOpenAIWSModeInfo( - "ingress_ws_preflight_ping_recovery account_id=%d turn=%d conn_id=%s action=drop_previous_response_id_retry previous_response_id=%s has_function_call_output=%v has_replay_tool_context=%v", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), - hasFCOutput, - hasReplayToolContext, - ) - turnPrevRecoveryTried = true - currentPayload = updatedWithInput - currentPayloadBytes = len(updatedWithInput) - resetSessionLease(true) - skipBeforeTurn = true - continue - } - } - } - if hasFCOutput && currentPreviousResponseID != "" { - reason := "function_call_output_missing_replay_context" - if hasReplayToolContext { - reason = "function_call_output_replay_not_applied" - } - logOpenAIWSModeInfo( - "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=%s action=fail_close previous_response_id=%s has_replay_tool_context=%v", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - reason, - truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), - hasReplayToolContext, - ) - } - resetSessionLease(true) - return NewOpenAIWSClientCloseError( - coderws.StatusPolicyViolation, - "upstream continuation connection is unavailable; please restart the conversation", - pingErr, - ) - } - resetSessionLease(true) - - acquiredLease, acquireErr := acquireTurnLease(turn, preferredConnID, forcePreferredConn) - if acquireErr != nil { - return fmt.Errorf("acquire upstream websocket after preflight ping fail: %w", acquireErr) - } - sessionLease = acquiredLease - sessionConnID = strings.TrimSpace(sessionLease.ConnID()) - if storeDisabled { - pinSessionConn(sessionConnID) - } - } - } - connID := sessionConnID - if currentPreviousResponseID != "" { - chainedFromLast := expectedPrev != "" && currentPreviousResponseID == expectedPrev - currentPreviousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(currentPreviousResponseID) - logOpenAIWSModeInfo( - "ingress_ws_turn_chain account_id=%d turn=%d conn_id=%s previous_response_id=%s previous_response_id_kind=%s last_turn_response_id=%s chained_from_last=%v preferred_conn_id=%s header_session_id=%s header_conversation_id=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v store_disabled=%v", - account.ID, - turn, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(currentPreviousResponseIDKind), - truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), - chainedFromLast, - truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), - openAIWSHeaderValueForLog(baseAcquireReq.Headers, "session_id"), - openAIWSHeaderValueForLog(baseAcquireReq.Headers, "conversation_id"), - turnState != "", - len(turnState), - openAIWSPayloadStringFromRaw(currentPayload, "prompt_cache_key") != "", - storeDisabled, - ) - } - - result, relayErr := sendAndRelay(turn, sessionLease, currentPayload, currentPayloadBytes, currentOriginalModel, currentImageBillingModel, currentImageSizeTier, currentImageInputSize) - if relayErr != nil { - lastTurnClean = false - if recoverIngressPrevResponseNotFound(relayErr, turn, connID) { - continue - } - if retryIngressTurn(relayErr, turn, connID) { - continue - } - finalErr := relayErr - if unwrapped := errors.Unwrap(relayErr); unwrapped != nil { - finalErr = unwrapped - } - if hooks != nil && hooks.AfterTurn != nil { - hooks.AfterTurn(turn, nil, finalErr) - } - sessionLease.MarkBroken() - return finalErr - } - turnRetry = 0 - turnPrevRecoveryTried = false - lastTurnFinishedAt = time.Now() - lastTurnClean = true - if hooks != nil && hooks.AfterTurn != nil { - hooks.AfterTurn(turn, result, nil) - } - if result == nil { - return errors.New("websocket turn result is nil") - } - responseID := strings.TrimSpace(result.RequestID) - lastTurnResponseID = responseID - lastTurnPayload = cloneOpenAIWSPayloadBytes(currentPayload) - lastTurnReplayInput = cloneOpenAIWSRawMessages(currentTurnReplayInput) - lastTurnReplayInputExists = currentTurnReplayInputExists - if result.wsReplayInputExists { - lastTurnReplayInput = append(lastTurnReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...) - lastTurnReplayInputExists = true - } - nextStrictState, strictStateErr := buildOpenAIWSIngressPreviousTurnStrictState(currentPayload) - if strictStateErr != nil { - lastTurnStrictState = nil - logOpenAIWSModeInfo( - "ingress_ws_prev_response_strict_state_skip account_id=%d turn=%d conn_id=%s reason=build_error cause=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(strictStateErr.Error(), openAIWSLogValueMaxLen), - ) - } else { - lastTurnStrictState = nextStrictState - } - - if responseID != "" && stateStore != nil { - ttl := s.openAIWSResponseStickyTTL() - logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) - stateStore.BindResponseConn(responseID, connID, ttl) - } - if stateStore != nil && storeDisabled && sessionHash != "" { - stateStore.BindSessionConn(groupID, sessionHash, connID, s.openAIWSSessionStickyTTL()) - } - if connID != "" { - preferredConnID = connID - } - - nextClientMessage, readErr := readClientMessage() - if readErr != nil { - if isOpenAIWSClientDisconnectError(readErr) { - closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) - logOpenAIWSModeInfo( - "ingress_ws_client_closed account_id=%d conn_id=%s close_status=%s close_reason=%s", - account.ID, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - closeStatus, - truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), - ) - return nil - } - return fmt.Errorf("read client websocket request: %w", readErr) - } - - nextPayload, parseErr := parseClientPayload(nextClientMessage) - if parseErr != nil { - return parseErr - } - if nextPayload.promptCacheKey != "" { - // ingress 会话在整个客户端 WS 生命周期内复用同一上游连接； - // prompt_cache_key 对握手头的更新仅在未来需要重新建连时生效。 - updatedHeaders, _, updHdrErr := s.buildOpenAIWSHeaders(ctx, c, account, token, wsDecision, isCodexCLI, turnState, strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)), nextPayload.promptCacheKey) - if updHdrErr != nil { - logOpenAIWSModeInfo("ingress_ws_update_headers_failed account_id=%d err=%v", account.ID, updHdrErr) - } else { - baseAcquireReq.Headers = updatedHeaders - } - } - if nextPayload.previousResponseID != "" { - expectedPrev := strings.TrimSpace(lastTurnResponseID) - chainedFromLast := expectedPrev != "" && nextPayload.previousResponseID == expectedPrev - nextPreviousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(nextPayload.previousResponseID) - logOpenAIWSModeInfo( - "ingress_ws_next_turn_chain account_id=%d turn=%d next_turn=%d conn_id=%s previous_response_id=%s previous_response_id_kind=%s last_turn_response_id=%s chained_from_last=%v has_prompt_cache_key=%v store_disabled=%v", - account.ID, - turn, - turn+1, - truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(nextPayload.previousResponseID, openAIWSIDValueMaxLen), - normalizeOpenAIWSLogValue(nextPreviousResponseIDKind), - truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), - chainedFromLast, - nextPayload.promptCacheKey != "", - storeDisabled, - ) - } - if stateStore != nil && nextPayload.previousResponseID != "" { - if stickyConnID, ok := stateStore.GetResponseConn(nextPayload.previousResponseID); ok { - if sessionConnID != "" && stickyConnID != "" && stickyConnID != sessionConnID { - logOpenAIWSModeInfo( - "ingress_ws_keep_session_conn account_id=%d turn=%d conn_id=%s sticky_conn_id=%s previous_response_id=%s", - account.ID, - turn, - truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(stickyConnID, openAIWSIDValueMaxLen), - truncateOpenAIWSLogValue(nextPayload.previousResponseID, openAIWSIDValueMaxLen), - ) - } else { - preferredConnID = stickyConnID - } - } - } - currentPayload = nextPayload.payloadRaw - currentOriginalModel = nextPayload.originalModel - currentImageBillingModel = nextPayload.imageBillingModel - currentImageSizeTier = nextPayload.imageSizeTier - currentImageInputSize = nextPayload.imageInputSize - currentPayloadBytes = nextPayload.payloadBytes - storeDisabled = s.isOpenAIWSStoreDisabledInRequestRaw(currentPayload, account) - if !storeDisabled { - unpinSessionConn(sessionConnID) - } - turn++ - } -} - -func (s *OpenAIGatewayService) isOpenAIWSGeneratePrewarmEnabled() bool { - return s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.PrewarmGenerateEnabled -} - -// performOpenAIWSGeneratePrewarm 在 WSv2 下执行可选的 generate=false 预热。 -// 预热默认关闭，仅在配置开启后生效；失败时按可恢复错误回退到 HTTP。 -func (s *OpenAIGatewayService) performOpenAIWSGeneratePrewarm( - ctx context.Context, - lease *openAIWSConnLease, - decision OpenAIWSProtocolDecision, - payload map[string]any, - previousResponseID string, - reqBody map[string]any, - account *Account, - stateStore OpenAIWSStateStore, - groupID int64, -) error { - if s == nil { - return nil - } - if lease == nil || account == nil { - logOpenAIWSModeInfo("prewarm_skip reason=invalid_state has_lease=%v has_account=%v", lease != nil, account != nil) - return nil - } - connID := strings.TrimSpace(lease.ConnID()) - if !s.isOpenAIWSGeneratePrewarmEnabled() { - return nil - } - if decision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { - logOpenAIWSModeInfo( - "prewarm_skip account_id=%d conn_id=%s reason=transport_not_v2 transport=%s", - account.ID, - connID, - normalizeOpenAIWSLogValue(string(decision.Transport)), - ) - return nil - } - if strings.TrimSpace(previousResponseID) != "" { - logOpenAIWSModeInfo( - "prewarm_skip account_id=%d conn_id=%s reason=has_previous_response_id previous_response_id=%s", - account.ID, - connID, - truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen), - ) - return nil - } - if lease.IsPrewarmed() { - logOpenAIWSModeInfo("prewarm_skip account_id=%d conn_id=%s reason=already_prewarmed", account.ID, connID) - return nil - } - if NeedsToolContinuation(reqBody) { - logOpenAIWSModeInfo("prewarm_skip account_id=%d conn_id=%s reason=tool_continuation", account.ID, connID) - return nil - } - prewarmStart := time.Now() - logOpenAIWSModeInfo("prewarm_start account_id=%d conn_id=%s", account.ID, connID) - - prewarmPayload := make(map[string]any, len(payload)+1) - for k, v := range payload { - prewarmPayload[k] = v - } - prewarmPayload["generate"] = false - prewarmPayloadJSON := payloadAsJSONBytes(prewarmPayload) - - if err := lease.WriteJSONWithContextTimeout(ctx, prewarmPayload, s.openAIWSWriteTimeout()); err != nil { - lease.MarkBroken() - logOpenAIWSModeInfo( - "prewarm_write_fail account_id=%d conn_id=%s cause=%s", - account.ID, - connID, - truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), - ) - return wrapOpenAIWSFallback("prewarm_write", err) - } - logOpenAIWSModeInfo("prewarm_write_sent account_id=%d conn_id=%s payload_bytes=%d", account.ID, connID, len(prewarmPayloadJSON)) - - prewarmResponseID := "" - prewarmEventCount := 0 - prewarmTerminalCount := 0 - for { - message, readErr := lease.ReadMessageWithContextTimeout(ctx, s.openAIWSReadTimeout()) - if readErr != nil { - lease.MarkBroken() - closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) - logOpenAIWSModeInfo( - "prewarm_read_fail account_id=%d conn_id=%s close_status=%s close_reason=%s cause=%s events=%d", - account.ID, - connID, - closeStatus, - closeReason, - truncateOpenAIWSLogValue(readErr.Error(), openAIWSLogValueMaxLen), - prewarmEventCount, - ) - return wrapOpenAIWSFallback("prewarm_"+classifyOpenAIWSReadFallbackReason(readErr), readErr) - } - - eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(message) - if eventType == "" { - continue - } - prewarmEventCount++ - if prewarmResponseID == "" && eventResponseID != "" { - prewarmResponseID = eventResponseID - } - if prewarmEventCount <= openAIWSPrewarmEventLogHead || eventType == "error" || isOpenAIWSTerminalEvent(eventType) { - logOpenAIWSModeInfo( - "prewarm_event account_id=%d conn_id=%s idx=%d type=%s bytes=%d", - account.ID, - connID, - prewarmEventCount, - truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen), - len(message), - ) - } - - if eventType == "error" { - errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message) - s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw) - errMsg := strings.TrimSpace(errMsgRaw) - if errMsg == "" { - errMsg = "OpenAI websocket prewarm error" - } - fallbackReason, canFallback := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) - errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) - logOpenAIWSModeInfo( - "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", - account.ID, - connID, - prewarmEventCount, - truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen), - canFallback, - errCode, - errType, - errMessage, - ) - lease.MarkBroken() - if canFallback { - return wrapOpenAIWSFallback("prewarm_"+fallbackReason, errors.New(errMsg)) - } - return wrapOpenAIWSFallback("prewarm_error_event", errors.New(errMsg)) - } - - if isOpenAIWSTerminalEvent(eventType) { - prewarmTerminalCount++ - break - } - } - - lease.MarkPrewarmed() - if prewarmResponseID != "" && stateStore != nil { - ttl := s.openAIWSResponseStickyTTL() - logOpenAIWSBindResponseAccountWarn(groupID, account.ID, prewarmResponseID, stateStore.BindResponseAccount(ctx, groupID, prewarmResponseID, account.ID, ttl)) - stateStore.BindResponseConn(prewarmResponseID, lease.ConnID(), ttl) - } - logOpenAIWSModeInfo( - "prewarm_done account_id=%d conn_id=%s response_id=%s events=%d terminal_events=%d duration_ms=%d", - account.ID, - connID, - truncateOpenAIWSLogValue(prewarmResponseID, openAIWSIDValueMaxLen), - prewarmEventCount, - prewarmTerminalCount, - time.Since(prewarmStart).Milliseconds(), - ) - return nil -} - -func payloadAsJSON(payload map[string]any) string { - return string(payloadAsJSONBytes(payload)) -} - -func payloadAsJSONBytes(payload map[string]any) []byte { - if len(payload) == 0 { - return []byte("{}") - } - body, err := json.Marshal(payload) - if err != nil { - return []byte("{}") - } - return body -} - -func isOpenAIWSTerminalEvent(eventType string) bool { - switch strings.TrimSpace(eventType) { - case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled": - return true - default: - return false - } -} - -func isOpenAIWSTokenEvent(eventType string) bool { - eventType = strings.TrimSpace(eventType) - if eventType == "" { - return false - } - switch eventType { - case "response.created", "response.in_progress", "response.output_item.added", "response.output_item.done": - return false - } - if strings.Contains(eventType, ".delta") { - return true - } - if strings.HasPrefix(eventType, "response.output_text") { - return true - } - if strings.HasPrefix(eventType, "response.output") { - return true - } - // 终止事件（response.completed/done/failed/...）由 isOpenAIWSTerminalEvent 单独处理。 - // 不能把它们当作 token event，否则当上游没有可识别的 delta 时， - // firstTokenMs 会被填到终止时刻，等于把"总耗时"误报为"首 token 延迟"。 - return false -} - -func replaceOpenAIWSMessageModel(message []byte, fromModel, toModel string) []byte { - if len(message) == 0 { - return message - } - if strings.TrimSpace(fromModel) == "" || strings.TrimSpace(toModel) == "" || fromModel == toModel { - return message - } - if !bytes.Contains(message, []byte(`"model"`)) || !bytes.Contains(message, []byte(fromModel)) { - return message - } - modelValues := gjson.GetManyBytes(message, "model", "response.model") - replaceModel := modelValues[0].Exists() && modelValues[0].Str == fromModel - replaceResponseModel := modelValues[1].Exists() && modelValues[1].Str == fromModel - if !replaceModel && !replaceResponseModel { - return message - } - updated := message - if replaceModel { - if next, err := sjson.SetBytes(updated, "model", toModel); err == nil { - updated = next - } - } - if replaceResponseModel { - if next, err := sjson.SetBytes(updated, "response.model", toModel); err == nil { - updated = next - } - } - return updated -} - -func populateOpenAIUsageFromResponseJSON(body []byte, usage *OpenAIUsage) { - if usage == nil || len(body) == 0 { - return - } - values := gjson.GetManyBytes( - body, - "usage.input_tokens", - "usage.output_tokens", - "usage.input_tokens_details.cached_tokens", - ) - usage.InputTokens = int(values[0].Int()) - usage.OutputTokens = int(values[1].Int()) - usage.CacheReadInputTokens = int(values[2].Int()) -} - -func getOpenAIGroupIDFromContext(c *gin.Context) int64 { - if c == nil { - return 0 - } - value, exists := c.Get("api_key") - if !exists { - return 0 - } - apiKey, ok := value.(*APIKey) - if !ok || apiKey == nil || apiKey.GroupID == nil { - return 0 - } - return *apiKey.GroupID -} - -// SelectAccountByPreviousResponseID 按 previous_response_id 命中账号粘连。 -// 未命中或账号不可用时返回 (nil, nil)，由调用方继续走常规调度。 -func (s *OpenAIGatewayService) SelectAccountByPreviousResponseID( - ctx context.Context, - groupID *int64, - previousResponseID string, - requestedModel string, - excludedIDs map[int64]struct{}, - requireCompact bool, -) (*AccountSelectionResult, error) { - return s.selectAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, "", requireCompact) -} - -func (s *OpenAIGatewayService) selectAccountByPreviousResponseIDForCapability( - ctx context.Context, - groupID *int64, - previousResponseID string, - requestedModel string, - excludedIDs map[int64]struct{}, - requiredCapability OpenAIEndpointCapability, - requireCompact bool, -) (*AccountSelectionResult, error) { - if s == nil { - return nil, nil - } - accountID, account, responseID, store := s.resolveAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact) - if accountID <= 0 || account == nil || store == nil { - return nil, nil - } - - result, acquireErr := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency) - if acquireErr == nil && result.Acquired { - logOpenAIWSBindResponseAccountWarn( - derefGroupID(groupID), - accountID, - responseID, - store.BindResponseAccount(ctx, derefGroupID(groupID), responseID, accountID, s.openAIWSResponseStickyTTL()), - ) - return &AccountSelectionResult{ - Account: account, - Acquired: true, - ReleaseFunc: result.ReleaseFunc, - }, nil - } - - cfg := s.schedulingConfig() - if s.concurrencyService != nil { - return &AccountSelectionResult{ - Account: account, - WaitPlan: &AccountWaitPlan{ - AccountID: accountID, - MaxConcurrency: account.Concurrency, - Timeout: cfg.StickySessionWaitTimeout, - MaxWaiting: cfg.StickySessionMaxWaiting, - }, - }, nil - } - return nil, nil -} - -func (s *OpenAIGatewayService) ResolveAccountIDByPreviousResponseIDForScheduler( - ctx context.Context, - groupID *int64, - previousResponseID string, - requestedModel string, - excludedIDs map[int64]struct{}, - requiredCapability OpenAIEndpointCapability, - requireCompact bool, -) int64 { - accountID, _, _, _ := s.resolveAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact) - return accountID -} - -func (s *OpenAIGatewayService) resolveAccountByPreviousResponseIDForCapability( - ctx context.Context, - groupID *int64, - previousResponseID string, - requestedModel string, - excludedIDs map[int64]struct{}, - requiredCapability OpenAIEndpointCapability, - requireCompact bool, -) (int64, *Account, string, OpenAIWSStateStore) { - if s == nil { - return 0, nil, "", nil - } - responseID := strings.TrimSpace(previousResponseID) - if responseID == "" { - return 0, nil, "", nil - } - store := s.getOpenAIWSStateStore() - if store == nil { - return 0, nil, "", nil - } - - accountID, err := store.GetResponseAccount(ctx, derefGroupID(groupID), responseID) - if err != nil || accountID <= 0 { - return 0, nil, "", nil - } - if excludedIDs != nil { - if _, excluded := excludedIDs[accountID]; excluded { - return 0, nil, "", nil - } - } - - account, err := s.getSchedulableAccount(ctx, accountID) - if err != nil || account == nil { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return 0, nil, "", nil - } - // 非 WSv2 场景（如 force_http/全局关闭）不应使用 previous_response_id 粘连， - // 以保持“回滚到 HTTP”后的历史行为一致性。 - if s.getOpenAIWSProtocolResolver().Resolve(account).Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { - return 0, nil, "", nil - } - if shouldClearStickySession(account, requestedModel) || !account.IsOpenAI() || !account.IsSchedulable() { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return 0, nil, "", nil - } - if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return 0, nil, "", nil - } - if requestedModel != "" && !account.IsModelSupported(requestedModel) { - return 0, nil, "", nil - } - if !account.SupportsOpenAIEndpointCapability(requiredCapability) { - return 0, nil, "", nil - } - // Quota auto-pause must also gate the previous_response_id sticky path; otherwise an - // account over its 5h/7d threshold keeps serving the same response chain even though - // normal scheduling skips it. Pause is transient, so fall through to normal scheduling - // without deleting the binding (the window may reset before the next turn). - if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { - return 0, nil, "", nil - } - if s.schedulerSnapshot != nil && s.accountRepo != nil { - latest, latestErr := s.accountRepo.GetByID(ctx, account.ID) - if latestErr != nil || latest == nil { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return 0, nil, "", nil - } - if shouldClearStickySession(latest, requestedModel) || !latest.IsOpenAI() || !latest.IsSchedulable() { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return 0, nil, "", nil - } - if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return 0, nil, "", nil - } - if requestedModel != "" && !latest.IsModelSupported(requestedModel) { - return 0, nil, "", nil - } - if !latest.SupportsOpenAIEndpointCapability(requiredCapability) { - return 0, nil, "", nil - } - if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, latest); paused { - return 0, nil, "", nil - } - if s.isOpenAIAccountRuntimeBlocked(latest) { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return 0, nil, "", nil - } - account = latest - } - if requireCompact && openAICompactSupportTier(account) == 0 { - _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) - return 0, nil, "", nil - } - return accountID, account, responseID, store -} - -func classifyOpenAIWSAcquireError(err error) string { - if err == nil { - return "acquire_conn" - } - var dialErr *openAIWSDialError - if errors.As(err, &dialErr) { - switch dialErr.StatusCode { - case 426: - return "upgrade_required" - case 401, 403: - return "auth_failed" - case 429: - return "upstream_rate_limited" - } - if dialErr.StatusCode >= 500 { - return "upstream_5xx" - } - return "dial_failed" - } - if errors.Is(err, errOpenAIWSConnQueueFull) { - return "conn_queue_full" - } - if errors.Is(err, errOpenAIWSPreferredConnUnavailable) { - return "preferred_conn_unavailable" - } - if errors.Is(err, context.DeadlineExceeded) { - return "acquire_timeout" - } - return "acquire_conn" -} - -func isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw string) bool { - code := strings.ToLower(strings.TrimSpace(codeRaw)) - errType := strings.ToLower(strings.TrimSpace(errTypeRaw)) - msg := strings.ToLower(strings.TrimSpace(msgRaw)) - - if strings.Contains(errType, "rate_limit") || strings.Contains(errType, "usage_limit") { - return true - } - if strings.Contains(code, "rate_limit") || strings.Contains(code, "usage_limit") || strings.Contains(code, "insufficient_quota") { - return true - } - if strings.Contains(msg, "usage limit") && strings.Contains(msg, "reached") { - return true - } - if strings.Contains(msg, "rate limit") && (strings.Contains(msg, "reached") || strings.Contains(msg, "exceeded")) { - return true - } - return false -} - -func (s *OpenAIGatewayService) persistOpenAIWSRateLimitSignal(ctx context.Context, account *Account, headers http.Header, responseBody []byte, codeRaw, errTypeRaw, msgRaw string) { - if s == nil || s.rateLimitService == nil || account == nil || account.Platform != PlatformOpenAI { - return - } - if !isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw) { - return - } - s.handleOpenAIAccountUpstreamError(ctx, account, http.StatusTooManyRequests, headers, responseBody) -} - -func classifyOpenAIWSErrorEventFromRaw(codeRaw, errTypeRaw, msgRaw string) (string, bool) { - code := strings.ToLower(strings.TrimSpace(codeRaw)) - errType := strings.ToLower(strings.TrimSpace(errTypeRaw)) - msg := strings.ToLower(strings.TrimSpace(msgRaw)) - - switch code { - case "upgrade_required": - return "upgrade_required", true - case "websocket_not_supported", "websocket_unsupported": - return "ws_unsupported", true - case "websocket_connection_limit_reached": - return "ws_connection_limit_reached", true - case "invalid_encrypted_content": - return "invalid_encrypted_content", true - case "previous_response_not_found": - return "previous_response_not_found", true - } - if isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw) { - return "upstream_rate_limited", false - } - if strings.Contains(msg, "upgrade required") || strings.Contains(msg, "status 426") { - return "upgrade_required", true - } - if strings.Contains(errType, "upgrade") { - return "upgrade_required", true - } - if strings.Contains(msg, "websocket") && strings.Contains(msg, "unsupported") { - return "ws_unsupported", true - } - if strings.Contains(msg, "connection limit") && strings.Contains(msg, "websocket") { - return "ws_connection_limit_reached", true - } - if strings.Contains(msg, "invalid_encrypted_content") || - (strings.Contains(msg, "encrypted content") && strings.Contains(msg, "could not be verified")) { - return "invalid_encrypted_content", true - } - if strings.Contains(msg, "previous_response_not_found") || - (strings.Contains(msg, "previous response") && strings.Contains(msg, "not found")) { - return "previous_response_not_found", true - } - if strings.Contains(errType, "server_error") || strings.Contains(code, "server_error") { - return "upstream_error_event", true - } - return "event_error", false -} - -func classifyOpenAIWSErrorEvent(message []byte) (string, bool) { - if len(message) == 0 { - return "event_error", false - } - return classifyOpenAIWSErrorEventFromRaw(parseOpenAIWSErrorEventFields(message)) -} - -func openAIWSErrorHTTPStatusFromRaw(codeRaw, errTypeRaw string) int { - code := strings.ToLower(strings.TrimSpace(codeRaw)) - errType := strings.ToLower(strings.TrimSpace(errTypeRaw)) - switch { - case strings.Contains(errType, "invalid_request"), - strings.Contains(code, "invalid_request"), - strings.Contains(code, "bad_request"), - code == "invalid_encrypted_content", - code == "previous_response_not_found": - return http.StatusBadRequest - case strings.Contains(errType, "authentication"), - strings.Contains(code, "invalid_api_key"), - strings.Contains(code, "unauthorized"): - return http.StatusUnauthorized - case strings.Contains(errType, "permission"), - strings.Contains(code, "forbidden"): - return http.StatusForbidden - case isOpenAIWSRateLimitError(codeRaw, errTypeRaw, ""): - return http.StatusTooManyRequests - default: - return http.StatusBadGateway - } -} - -func openAIWSErrorHTTPStatus(message []byte) int { - if len(message) == 0 { - return http.StatusBadGateway - } - codeRaw, errTypeRaw, _ := parseOpenAIWSErrorEventFields(message) - return openAIWSErrorHTTPStatusFromRaw(codeRaw, errTypeRaw) -} - -func (s *OpenAIGatewayService) openAIWSFallbackCooldown() time.Duration { - if s == nil || s.cfg == nil { - return 30 * time.Second - } - seconds := s.cfg.Gateway.OpenAIWS.FallbackCooldownSeconds - if seconds <= 0 { - return 0 - } - return time.Duration(seconds) * time.Second -} - -func (s *OpenAIGatewayService) isOpenAIWSFallbackCooling(accountID int64) bool { - if s == nil || accountID <= 0 { - return false - } - cooldown := s.openAIWSFallbackCooldown() - if cooldown <= 0 { - return false - } - rawUntil, ok := s.openaiWSFallbackUntil.Load(accountID) - if !ok || rawUntil == nil { - return false - } - until, ok := rawUntil.(time.Time) - if !ok || until.IsZero() { - s.openaiWSFallbackUntil.Delete(accountID) - return false - } - if time.Now().Before(until) { - return true - } - s.openaiWSFallbackUntil.Delete(accountID) - return false -} - -func (s *OpenAIGatewayService) markOpenAIWSFallbackCooling(accountID int64, _ string) { - if s == nil || accountID <= 0 { - return - } - cooldown := s.openAIWSFallbackCooldown() - if cooldown <= 0 { - return - } - s.openaiWSFallbackUntil.Store(accountID, time.Now().Add(cooldown)) -} - -func (s *OpenAIGatewayService) clearOpenAIWSFallbackCooling(accountID int64) { - if s == nil || accountID <= 0 { - return - } - s.openaiWSFallbackUntil.Delete(accountID) -} diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go new file mode 100644 index 0000000000..5fb9395461 --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -0,0 +1,1579 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/url" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/openai" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( + ctx context.Context, + c *gin.Context, + clientConn *coderws.Conn, + account *Account, + token string, + firstClientMessage []byte, + hooks *OpenAIWSIngressHooks, +) error { + if s == nil { + return errors.New("service is nil") + } + if c == nil { + return errors.New("gin context is nil") + } + if clientConn == nil { + return errors.New("client websocket is nil") + } + if account == nil { + return errors.New("account is nil") + } + if strings.TrimSpace(token) == "" { + return errors.New("token is empty") + } + + // 预取一次 OpenAI Fast Policy settings，绑定到 ctx，让该 WS session + // 内所有帧的 evaluateOpenAIFastPolicy 调用复用同一份快照，避免每帧 + // 进入 DB / settingRepo。Trade-off 见 withOpenAIFastPolicyContext 注释。 + if s.settingService != nil { + if settings, err := s.settingService.GetOpenAIFastPolicySettings(ctx); err == nil && settings != nil { + ctx = withOpenAIFastPolicyContext(ctx, settings) + } + } + + wsDecision := s.getOpenAIWSProtocolResolver().Resolve(account) + forceHTTPBridge := account.Platform == PlatformGrok + modeRouterV2Enabled := s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.ModeRouterV2Enabled + ingressMode := OpenAIWSIngressModeCtxPool + if modeRouterV2Enabled && !forceHTTPBridge { + ingressMode = account.ResolveOpenAIResponsesWebSocketV2Mode(s.cfg.Gateway.OpenAIWS.IngressModeDefault) + if ingressMode == OpenAIWSIngressModeOff { + return NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + "websocket mode is disabled for this account", + nil, + ) + } + switch ingressMode { + case OpenAIWSIngressModePassthrough: + if wsDecision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { + return fmt.Errorf("websocket ingress requires ws_v2 transport, got=%s", wsDecision.Transport) + } + return s.proxyResponsesWebSocketV2Passthrough( + ctx, + c, + clientConn, + account, + token, + firstClientMessage, + hooks, + wsDecision, + ) + case OpenAIWSIngressModeHTTPBridge: + forceHTTPBridge = true + case OpenAIWSIngressModeCtxPool, OpenAIWSIngressModeShared, OpenAIWSIngressModeDedicated: + // continue + default: + return NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + "websocket mode only supports ctx_pool/passthrough/http_bridge", + nil, + ) + } + } + if !forceHTTPBridge && wsDecision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { + return fmt.Errorf("websocket ingress requires ws_v2 transport, got=%s", wsDecision.Transport) + } + dedicatedMode := modeRouterV2Enabled && ingressMode == OpenAIWSIngressModeDedicated + + wsURL := "" + wsHost := "-" + wsPath := "-" + if forceHTTPBridge { + wsHost = "xai-http-bridge" + wsPath = "/v1/responses" + } else { + var err error + wsURL, err = s.buildOpenAIResponsesWSURL(account) + if err != nil { + return fmt.Errorf("build ws url: %w", err) + } + if parsedURL, parseErr := url.Parse(wsURL); parseErr == nil && parsedURL != nil { + wsHost = normalizeOpenAIWSLogValue(parsedURL.Host) + wsPath = normalizeOpenAIWSLogValue(parsedURL.Path) + } + } + debugEnabled := isOpenAIWSModeDebugEnabled() + isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator")) || (s.cfg != nil && s.cfg.Gateway.ForceCodexCLI) + + type openAIWSClientPayload struct { + payloadRaw []byte + rawForHash []byte + promptCacheKey string + previousResponseID string + originalModel string + imageBillingModel string + imageSizeTier string + imageInputSize string + payloadBytes int + } + ingressSessionOriginalModel := "" + + applyPayloadMutation := func(current []byte, path string, value any) ([]byte, error) { + next, err := sjson.SetBytes(current, path, value) + if err == nil { + return next, nil + } + + // 仅在确实需要修改 payload 且 sjson 失败时，退回 map 路径确保兼容性。 + payload := make(map[string]any) + if unmarshalErr := json.Unmarshal(current, &payload); unmarshalErr != nil { + return nil, err + } + switch path { + case "type", "model": + payload[path] = value + case "client_metadata." + openAIWSTurnMetadataHeader: + setOpenAIWSTurnMetadata(payload, fmt.Sprintf("%v", value)) + default: + return nil, err + } + rebuilt, marshalErr := json.Marshal(payload) + if marshalErr != nil { + return nil, marshalErr + } + return rebuilt, nil + } + + parseClientPayload := func(raw []byte) (openAIWSClientPayload, error) { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "empty websocket request payload", nil) + } + if !gjson.ValidBytes(trimmed) { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", errors.New("invalid json")) + } + + values := gjson.GetManyBytes(trimmed, "type", "model", "prompt_cache_key", "previous_response_id") + eventType := strings.TrimSpace(values[0].String()) + normalized := trimmed + switch eventType { + case "": + eventType = "response.create" + next, setErr := applyPayloadMutation(normalized, "type", eventType) + if setErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", setErr) + } + normalized = next + case "response.create": + case "response.append": + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + "response.append is not supported in ws v2; use response.create with previous_response_id", + nil, + ) + default: + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + fmt.Sprintf("unsupported websocket request type: %s", eventType), + nil, + ) + } + + originalModel := strings.TrimSpace(values[1].String()) + modelMissing := originalModel == "" + if originalModel == "" { + // 入站 WS 长会话里，部分客户端只在第一轮 response.create 上声明 + // model，后续 turn 复用同一 session-level model。为避免因省略 + // model 直接断开用户连接，这里回落到上一轮已通过校验的客户端模型， + // 并在下方写回上游 payload，保证账号模型映射/fast policy/图片权限 + // 仍按同一模型执行。 + originalModel = ingressSessionOriginalModel + if originalModel == "" { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + "model is required in response.create payload", + nil, + ) + } + } + promptCacheKey := strings.TrimSpace(values[2].String()) + previousResponseID := strings.TrimSpace(values[3].String()) + previousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(previousResponseID) + if previousResponseID != "" && previousResponseIDKind == OpenAIPreviousResponseIDKindMessageID { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + "previous_response_id must be a response.id (resp_*), not a message id", + nil, + ) + } + if turnMetadata := strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)); turnMetadata != "" { + next, setErr := applyPayloadMutation(normalized, "client_metadata."+openAIWSTurnMetadataHeader, turnMetadata) + if setErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", setErr) + } + normalized = next + } + apiKey := getAPIKeyFromContext(c) + imageGenerationAllowed := GroupAllowsImageGeneration(apiKeyGroup(apiKey)) + codexImageGenerationExplicitToolPolicy := codexImageGenerationExplicitToolPolicyAllow + if isCodexCLI { + codexImageGenerationExplicitToolPolicy = account.CodexImageGenerationExplicitToolPolicy() + } + codexBridgeEnabled := isCodexCLI && imageGenerationAllowed && codexImageGenerationExplicitToolPolicy != codexImageGenerationExplicitToolPolicyStrip && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey) + if codexBridgeEnabled { + payloadMap := make(map[string]any) + if err := json.Unmarshal(normalized, &payloadMap); err != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", err) + } + bridgeModified := false + if ensureOpenAIResponsesImageGenerationTool(payloadMap) { + bridgeModified = true + logOpenAIWSModeInfo("ingress_ws_codex_image_tool_injected account_id=%d", account.ID) + } + if ensureOpenAIResponsesImageGenerationToolChoiceAuto(payloadMap) { + bridgeModified = true + logOpenAIWSModeInfo("ingress_ws_codex_image_tool_choice_auto account_id=%d", account.ID) + } + if normalizeOpenAIResponsesImageGenerationTools(payloadMap) { + bridgeModified = true + } + if applyCodexImageGenerationBridgeInstructions(payloadMap) { + bridgeModified = true + logOpenAIWSModeInfo("ingress_ws_codex_image_bridge_instructions_added account_id=%d", account.ID) + } + if bridgeModified { + rebuilt, marshalErr := json.Marshal(payloadMap) + if marshalErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", marshalErr) + } + normalized = rebuilt + } + } + upstreamModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)) + if modelMissing || upstreamModel != originalModel { + next, setErr := applyPayloadMutation(normalized, "model", upstreamModel) + if setErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", setErr) + } + normalized = next + } + if isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip { + if stripped, changed, stripErr := stripOpenAIImageGenerationToolFromRawPayload(normalized); stripErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", stripErr) + } else if changed { + normalized = stripped + logOpenAIWSModeInfo("ingress_ws_codex_image_tool_stripped_by_policy account_id=%d", account.ID) + } + } + if stripped, changed, stripErr := stripCodexSparkImageGenerationToolFromRawPayload(normalized, upstreamModel); stripErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", stripErr) + } else if changed { + normalized = stripped + logOpenAIWSModeInfo("ingress_ws_codex_spark_image_tool_stripped account_id=%d", account.ID) + } + imageIntent := IsImageGenerationIntent(openAIResponsesEndpoint, originalModel, normalized) + if imageIntent && !imageGenerationAllowed { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, ImageGenerationPermissionMessage(), nil) + } + imageBillingModel := "" + imageSizeTier := "" + imageInputSize := "" + if imageIntent { + var imageCfgErr error + imageCfg, imageCfgErr := resolveOpenAIResponsesImageBillingConfigDetailedFromBody(normalized, originalModel) + if imageCfgErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, imageCfgErr.Error(), imageCfgErr) + } + imageBillingModel = imageCfg.Model + imageSizeTier = imageCfg.SizeTier + imageInputSize = imageCfg.InputSize + } + + // Apply OpenAI Fast Policy on the response.create frame using the same + // evaluator/normalize/scope rules as the HTTP entrypoints. This is the + // single integration point for all WS ingress turns (first + follow-up + // frames flow through here). + // + // Model fallback: first turn still requires model at the handler layer； + // follow-up response.create frames may omit it and then reuse + // ingressSessionOriginalModel. We always write a concrete upstream model + // before evaluating policy, so whitelist / filter behavior remains stable. + policyApplied, blocked, policyErr := s.applyOpenAIFastPolicyToWSResponseCreate(ctx, account, upstreamModel, normalized) + if policyErr != nil { + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", policyErr) + } + if blocked != nil { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalPolicyDenied) + // Send a Realtime-style error event to the client first, then + // signal the handler to close the connection with PolicyViolation. + // We intentionally do NOT forward this frame upstream. + // + // coder/websocket@v1.8.14 Conn.Write is synchronous and flushes + // the underlying bufio writer before returning (write.go:42 → + // 307-311), and the subsequent close handshake re-acquires the + // same writeFrameMu, so the error event is guaranteed to reach + // the kernel send buffer before any close frame is queued. + eventBytes := buildOpenAIFastPolicyBlockedWSEvent(blocked) + if eventBytes != nil { + writeCtx, cancel := context.WithTimeout(ctx, s.openAIWSWriteTimeout()) + _ = clientConn.Write(writeCtx, coderws.MessageText, eventBytes) + cancel() + } + return openAIWSClientPayload{}, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + blocked.Message, + blocked, + ) + } + normalized = policyApplied + ingressSessionOriginalModel = originalModel + + return openAIWSClientPayload{ + payloadRaw: normalized, + rawForHash: trimmed, + promptCacheKey: promptCacheKey, + previousResponseID: previousResponseID, + originalModel: originalModel, + imageBillingModel: imageBillingModel, + imageSizeTier: imageSizeTier, + imageInputSize: imageInputSize, + payloadBytes: len(normalized), + }, nil + } + + writeClientMessage := func(message []byte) error { + writeCtx, cancel := context.WithTimeout(ctx, s.openAIWSWriteTimeout()) + defer cancel() + return clientConn.Write(writeCtx, coderws.MessageText, message) + } + + readClientMessage := func() ([]byte, error) { + msgType, payload, readErr := clientConn.Read(ctx) + if readErr != nil { + return nil, readErr + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + return nil, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + fmt.Sprintf("unsupported websocket client message type: %s", msgType.String()), + nil, + ) + } + return payload, nil + } + + firstPayload, err := parseClientPayload(firstClientMessage) + if err != nil { + return err + } + + turnState := strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader)) + stateStore := s.getOpenAIWSStateStore() + groupID := getOpenAIGroupIDFromContext(c) + storeDisabledConnMode := s.openAIWSStoreDisabledConnMode() + sessionHash := "" + preferredConnID := "" + storeDisabled := false + refreshIngressRouteState := func(payload openAIWSClientPayload) { + sessionHash = s.GenerateSessionHash(c, payload.rawForHash) + if turnState == "" && stateStore != nil && sessionHash != "" { + if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { + turnState = savedTurnState + } + } + + preferredConnID = "" + if stateStore != nil && payload.previousResponseID != "" { + if connID, ok := stateStore.GetResponseConn(payload.previousResponseID); ok { + preferredConnID = connID + } + } + + storeDisabled = s.isOpenAIWSStoreDisabledInRequestRaw(payload.payloadRaw, account) + if stateStore != nil && storeDisabled && payload.previousResponseID == "" && sessionHash != "" { + if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok { + preferredConnID = connID + } + } + } + refreshIngressRouteState(firstPayload) + + if forceHTTPBridge || s.shouldBridgeOpenAIWSHTTP(account, firstPayload.payloadBytes, firstPayload.previousResponseID) { + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_start account_id=%d account_type=%s payload_bytes=%d threshold_bytes=%d has_session_hash=%v store_disabled=%v", + account.ID, + account.Type, + firstPayload.payloadBytes, + s.openAIWSHTTPBridgeThresholdBytes(), + sessionHash != "", + storeDisabled, + ) + currentBridgePayload := firstPayload + var bridgeReplayInput []json.RawMessage + bridgeReplayInputExists := false + for turn := 1; ; turn++ { + if turn > 1 && hooks != nil && hooks.BeforeRequest != nil { + if err := hooks.BeforeRequest(turn, currentBridgePayload.payloadRaw, currentBridgePayload.originalModel); err != nil { + return err + } + } + if hooks != nil && hooks.BeforeTurn != nil { + if err := hooks.BeforeTurn(turn); err != nil { + return err + } + } + if turnState != "" && c != nil && c.Request != nil { + c.Request.Header.Set(openAIWSTurnStateHeader, turnState) + } + bridgePayloadRaw := currentBridgePayload.payloadRaw + bridgePayloadBytes := currentBridgePayload.payloadBytes + needsBridgeReplay := currentBridgePayload.previousResponseID != "" || openAIWSRawPayloadHasToolCallOutput(currentBridgePayload.payloadRaw) + turnReplayInput, turnReplayInputExists, replayInputErr := buildOpenAIWSReplayInputSequence( + bridgeReplayInput, + bridgeReplayInputExists, + currentBridgePayload.payloadRaw, + needsBridgeReplay, + ) + if replayInputErr != nil { + return fmt.Errorf("build websocket http bridge replay input: %w", replayInputErr) + } + if needsBridgeReplay && turnReplayInputExists { + updatedPayload, setInputErr := setOpenAIWSPayloadInputSequence( + currentBridgePayload.payloadRaw, + turnReplayInput, + true, + ) + if setInputErr != nil { + return fmt.Errorf("set websocket http bridge replay input: %w", setInputErr) + } + bridgePayloadRaw = updatedPayload + bridgePayloadBytes = len(updatedPayload) + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_replay_input account_id=%d turn=%d input_items=%d previous_response_id_present=%v has_tool_output=%v", + account.ID, + turn, + len(turnReplayInput), + currentBridgePayload.previousResponseID != "", + openAIWSRawPayloadHasToolCallOutput(currentBridgePayload.payloadRaw), + ) + } + result, bridgeErr := s.proxyOpenAIWSHTTPBridgeTurn( + ctx, + c, + account, + token, + bridgePayloadRaw, + bridgePayloadBytes, + currentBridgePayload.originalModel, + currentBridgePayload.imageBillingModel, + currentBridgePayload.imageSizeTier, + currentBridgePayload.imageInputSize, + turn, + writeClientMessage, + ) + if hooks != nil && hooks.AfterTurn != nil { + hooks.AfterTurn(turn, result, bridgeErr) + } + if bridgeErr != nil { + return bridgeErr + } + if result == nil { + return errors.New("websocket http bridge turn result is nil") + } + bridgeReplayInput = cloneOpenAIWSRawMessages(turnReplayInput) + bridgeReplayInputExists = turnReplayInputExists + if result.wsReplayInputExists { + bridgeReplayInput = append(bridgeReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...) + bridgeReplayInputExists = true + } + if bridgeTurnState := strings.TrimSpace(result.ResponseHeaders.Get(openAIWSTurnStateHeader)); bridgeTurnState != "" { + turnState = bridgeTurnState + if stateStore != nil && sessionHash != "" { + stateStore.BindSessionTurnState(groupID, sessionHash, bridgeTurnState, s.openAIWSSessionStickyTTL()) + } + } + responseID := strings.TrimSpace(result.RequestID) + if responseID != "" && stateStore != nil { + ttl := s.openAIWSResponseStickyTTL() + logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) + } + nextClientMessage, readErr := readClientMessage() + if readErr != nil { + if isOpenAIWSClientDisconnectError(readErr) { + closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) + logOpenAIWSModeInfo( + "ingress_ws_http_bridge_client_closed account_id=%d close_status=%s close_reason=%s", + account.ID, + closeStatus, + truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), + ) + return nil + } + return fmt.Errorf("read client websocket request: %w", readErr) + } + nextPayload, parseErr := parseClientPayload(nextClientMessage) + if parseErr != nil { + return parseErr + } + currentBridgePayload = nextPayload + } + } + + wsHeaders, _, buildHdrErr := s.buildOpenAIWSHeaders(ctx, c, account, token, wsDecision, isCodexCLI, turnState, strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)), firstPayload.promptCacheKey) + if buildHdrErr != nil { + return fmt.Errorf("build ws headers: %w", buildHdrErr) + } + baseAcquireReq := openAIWSAcquireRequest{ + Account: account, + WSURL: wsURL, + Headers: wsHeaders, + ProxyURL: func() string { + if account.ProxyID != nil && account.Proxy != nil { + return account.Proxy.URL() + } + return "" + }(), + ForceNewConn: false, + } + pool := s.getOpenAIWSConnPool() + if pool == nil { + return errors.New("openai ws conn pool is nil") + } + + logOpenAIWSModeInfo( + "ingress_ws_protocol_confirm account_id=%d account_type=%s transport=%s ws_host=%s ws_path=%s ws_mode=%s store_disabled=%v has_session_hash=%v has_previous_response_id=%v", + account.ID, + account.Type, + normalizeOpenAIWSLogValue(string(wsDecision.Transport)), + wsHost, + wsPath, + normalizeOpenAIWSLogValue(ingressMode), + storeDisabled, + sessionHash != "", + firstPayload.previousResponseID != "", + ) + + if debugEnabled { + logOpenAIWSModeDebug( + "ingress_ws_start account_id=%d account_type=%s transport=%s ws_host=%s preferred_conn_id=%s has_session_hash=%v has_previous_response_id=%v store_disabled=%v", + account.ID, + account.Type, + normalizeOpenAIWSLogValue(string(wsDecision.Transport)), + wsHost, + truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), + sessionHash != "", + firstPayload.previousResponseID != "", + storeDisabled, + ) + } + if firstPayload.previousResponseID != "" { + firstPreviousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(firstPayload.previousResponseID) + logOpenAIWSModeInfo( + "ingress_ws_continuation_probe account_id=%d turn=%d previous_response_id=%s previous_response_id_kind=%s preferred_conn_id=%s session_hash=%s header_session_id=%s header_conversation_id=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v store_disabled=%v", + account.ID, + 1, + truncateOpenAIWSLogValue(firstPayload.previousResponseID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(firstPreviousResponseIDKind), + truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(sessionHash, 12), + openAIWSHeaderValueForLog(baseAcquireReq.Headers, "session_id"), + openAIWSHeaderValueForLog(baseAcquireReq.Headers, "conversation_id"), + turnState != "", + len(turnState), + firstPayload.promptCacheKey != "", + storeDisabled, + ) + } + + acquireTimeout := s.openAIWSAcquireTimeout() + if acquireTimeout <= 0 { + acquireTimeout = 30 * time.Second + } + + acquireTurnLease := func(turn int, preferred string, forcePreferredConn bool) (*openAIWSConnLease, error) { + req := cloneOpenAIWSAcquireRequest(baseAcquireReq) + req.PreferredConnID = strings.TrimSpace(preferred) + req.ForcePreferredConn = forcePreferredConn + // dedicated 模式下每次获取均新建连接，避免跨会话复用残留上下文。 + req.ForceNewConn = dedicatedMode + acquireCtx, acquireCancel := context.WithTimeout(ctx, acquireTimeout) + lease, acquireErr := pool.Acquire(acquireCtx, req) + acquireCancel() + if acquireErr != nil { + dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(acquireErr) + logOpenAIWSModeInfo( + "ingress_ws_upstream_acquire_fail account_id=%d turn=%d reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_preferred_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v", + account.ID, + turn, + normalizeOpenAIWSLogValue(classifyOpenAIWSAcquireError(acquireErr)), + dialStatus, + dialClass, + dialCloseStatus, + truncateOpenAIWSLogValue(dialCloseReason, openAIWSHeaderValueMaxLen), + dialRespServer, + dialRespVia, + dialRespCFRay, + dialRespReqID, + truncateOpenAIWSLogValue(acquireErr.Error(), openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(preferred, openAIWSIDValueMaxLen), + forcePreferredConn, + wsHost, + wsPath, + account.ProxyID != nil && account.Proxy != nil, + ) + var dialErr *openAIWSDialError + if errors.As(acquireErr, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusTooManyRequests { + s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(acquireErr.Error())) + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusTooManyRequests, + ResponseHeaders: cloneHeader(dialErr.ResponseHeaders), + } + } + if errors.Is(acquireErr, errOpenAIWSPreferredConnUnavailable) { + return nil, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + "upstream continuation connection is unavailable; please restart the conversation", + acquireErr, + ) + } + if errors.Is(acquireErr, context.DeadlineExceeded) || errors.Is(acquireErr, errOpenAIWSConnQueueFull) { + return nil, NewOpenAIWSClientCloseError( + coderws.StatusTryAgainLater, + "upstream websocket is busy, please retry later", + acquireErr, + ) + } + return nil, acquireErr + } + connID := strings.TrimSpace(lease.ConnID()) + if handshakeTurnState := strings.TrimSpace(lease.HandshakeHeader(openAIWSTurnStateHeader)); handshakeTurnState != "" { + turnState = handshakeTurnState + if stateStore != nil && sessionHash != "" { + stateStore.BindSessionTurnState(groupID, sessionHash, handshakeTurnState, s.openAIWSSessionStickyTTL()) + } + updatedHeaders := cloneHeader(baseAcquireReq.Headers) + if updatedHeaders == nil { + updatedHeaders = make(http.Header) + } + updatedHeaders.Set(openAIWSTurnStateHeader, handshakeTurnState) + baseAcquireReq.Headers = updatedHeaders + } + logOpenAIWSModeInfo( + "ingress_ws_upstream_connected account_id=%d turn=%d conn_id=%s conn_reused=%v conn_pick_ms=%d queue_wait_ms=%d preferred_conn_id=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + lease.Reused(), + lease.ConnPickDuration().Milliseconds(), + lease.QueueWaitDuration().Milliseconds(), + truncateOpenAIWSLogValue(preferred, openAIWSIDValueMaxLen), + ) + return lease, nil + } + + sendAndRelay := func(turn int, lease *openAIWSConnLease, payload []byte, payloadBytes int, originalModel string, imageBillingModel string, imageSizeTier string, imageInputSize string) (*OpenAIForwardResult, error) { + if lease == nil { + return nil, errors.New("upstream websocket lease is nil") + } + turnStart := time.Now() + wroteDownstream := false + if err := lease.WriteJSONWithContextTimeout(ctx, json.RawMessage(payload), s.openAIWSWriteTimeout()); err != nil { + return nil, wrapOpenAIWSIngressTurnError( + "write_upstream", + fmt.Errorf("write upstream websocket request: %w", err), + false, + ) + } + if debugEnabled { + logOpenAIWSModeDebug( + "ingress_ws_turn_request_sent account_id=%d turn=%d conn_id=%s payload_bytes=%d", + account.ID, + turn, + truncateOpenAIWSLogValue(lease.ConnID(), openAIWSIDValueMaxLen), + payloadBytes, + ) + } + + responseID := "" + usage := OpenAIUsage{} + imageCounter := newOpenAIImageOutputCounter() + var firstTokenMs *int + reqStream := openAIWSPayloadBoolFromRaw(payload, "stream", true) + turnPreviousResponseID := openAIWSPayloadStringFromRaw(payload, "previous_response_id") + turnPreviousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(turnPreviousResponseID) + turnPromptCacheKey := openAIWSPayloadStringFromRaw(payload, "prompt_cache_key") + turnStoreDisabled := s.isOpenAIWSStoreDisabledInRequestRaw(payload, account) + turnHasFunctionCallOutput := openAIWSRawPayloadHasToolCallOutput(payload) + eventCount := 0 + tokenEventCount := 0 + terminalEventCount := 0 + replayCollector := &openAIWSToolCallReplayCollector{} + firstEventType := "" + lastEventType := "" + needModelReplace := false + clientDisconnected := false + mappedModel := "" + var mappedModelBytes []byte + if originalModel != "" { + mappedModel = normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)) + needModelReplace = mappedModel != "" && mappedModel != originalModel + if needModelReplace { + mappedModelBytes = []byte(mappedModel) + } + } + for { + upstreamMessage, readErr := lease.ReadMessageWithContextTimeout(ctx, s.openAIWSReadTimeout()) + if readErr != nil { + lease.MarkBroken() + return nil, wrapOpenAIWSIngressTurnError( + "read_upstream", + fmt.Errorf("read upstream websocket event: %w", readErr), + wroteDownstream, + ) + } + + eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage) + if responseID == "" && eventResponseID != "" { + responseID = eventResponseID + } + if eventType != "" { + eventCount++ + if firstEventType == "" { + firstEventType = eventType + } + lastEventType = eventType + } + if eventType == "error" { + errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) + s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw) + fallbackReason, _ := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) + errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) + recoverablePrevNotFound := fallbackReason == openAIWSIngressStagePreviousResponseNotFound && + turnPreviousResponseID != "" && + !turnHasFunctionCallOutput && + s.openAIWSIngressPreviousResponseRecoveryEnabled() && + !wroteDownstream + if recoverablePrevNotFound { + // 可恢复场景使用非 error 关键字日志，避免被 LegacyPrintf 误判为 ERROR 级别。 + logOpenAIWSModeInfo( + "ingress_ws_prev_response_recoverable account_id=%d turn=%d conn_id=%s idx=%d reason=%s code=%s type=%s message=%s previous_response_id=%s previous_response_id_kind=%s response_id=%s store_disabled=%v has_prompt_cache_key=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(lease.ConnID(), openAIWSIDValueMaxLen), + eventCount, + truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen), + errCode, + errType, + errMessage, + truncateOpenAIWSLogValue(turnPreviousResponseID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(turnPreviousResponseIDKind), + truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), + turnStoreDisabled, + turnPromptCacheKey != "", + ) + } else { + logOpenAIWSModeInfo( + "ingress_ws_error_event account_id=%d turn=%d conn_id=%s idx=%d fallback_reason=%s err_code=%s err_type=%s err_message=%s previous_response_id=%s previous_response_id_kind=%s response_id=%s store_disabled=%v has_prompt_cache_key=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(lease.ConnID(), openAIWSIDValueMaxLen), + eventCount, + truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen), + errCode, + errType, + errMessage, + truncateOpenAIWSLogValue(turnPreviousResponseID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(turnPreviousResponseIDKind), + truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), + turnStoreDisabled, + turnPromptCacheKey != "", + ) + } + // previous_response_not_found 在 ingress 模式支持单次恢复重试： + // 不把该 error 直接下发客户端，而是由上层去掉 previous_response_id 后重放当前 turn。 + if recoverablePrevNotFound { + lease.MarkBroken() + errMsg := strings.TrimSpace(errMsgRaw) + if errMsg == "" { + errMsg = "previous response not found" + } + return nil, wrapOpenAIWSIngressTurnError( + openAIWSIngressStagePreviousResponseNotFound, + errors.New(errMsg), + false, + ) + } + if !wroteDownstream && isOpenAIWSRateLimitError(errCodeRaw, errTypeRaw, errMsgRaw) { + lease.MarkBroken() + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusTooManyRequests, + ResponseBody: append([]byte(nil), upstreamMessage...), + ResponseHeaders: cloneHeader(lease.HandshakeHeaders()), + } + } + } + isTokenEvent := isOpenAIWSTokenEvent(eventType) + if isTokenEvent { + tokenEventCount++ + } + isTerminalEvent := isOpenAIWSTerminalEvent(eventType) + if isTerminalEvent { + terminalEventCount++ + } + if firstTokenMs == nil && isTokenEvent { + ms := int(time.Since(turnStart).Milliseconds()) + firstTokenMs = &ms + } + if openAIWSEventShouldParseUsage(eventType) { + parseOpenAIWSResponseUsageFromCompletedEvent(upstreamMessage, &usage) + } + imageCounter.AddSSEData(upstreamMessage) + + if eventType == "response.failed" { + if hit, code, msg := detectOpenAICyberPolicy(upstreamMessage); hit { + MarkOpsCyberPolicy(c, CyberPolicyMark{ + Code: code, + Message: msg, + Body: truncateString(string(upstreamMessage), 4096), + UpstreamStatus: http.StatusOK, + UpstreamInTok: usage.InputTokens, + UpstreamOutTok: usage.OutputTokens, + }) + } + } + + if !clientDisconnected { + if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && bytes.Contains(upstreamMessage, mappedModelBytes) { + upstreamMessage = replaceOpenAIWSMessageModel(upstreamMessage, mappedModel, originalModel) + } + if openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(upstreamMessage) { + if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(upstreamMessage); changed { + upstreamMessage = corrected + } + } + replayCollector.AddEvent(eventType, upstreamMessage) + if err := writeClientMessage(upstreamMessage); err != nil { + if isOpenAIWSClientDisconnectError(err) { + clientDisconnected = true + closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err) + logOpenAIWSModeInfo( + "ingress_ws_client_disconnected_drain account_id=%d turn=%d conn_id=%s close_status=%s close_reason=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(lease.ConnID(), openAIWSIDValueMaxLen), + closeStatus, + truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), + ) + } else { + return nil, wrapOpenAIWSIngressTurnError( + "write_client", + fmt.Errorf("write client websocket event: %w", err), + wroteDownstream, + ) + } + } else { + wroteDownstream = true + } + } + if isTerminalEvent { + // 客户端已断连时，上游连接的 session 状态不可信，标记 broken 避免回池复用。 + if clientDisconnected { + lease.MarkBroken() + } + firstTokenMsValue := -1 + if firstTokenMs != nil { + firstTokenMsValue = *firstTokenMs + } + if debugEnabled { + logOpenAIWSModeDebug( + "ingress_ws_turn_completed account_id=%d turn=%d conn_id=%s response_id=%s 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(lease.ConnID(), openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), + time.Since(turnStart).Milliseconds(), + eventCount, + tokenEventCount, + terminalEventCount, + truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen), + firstTokenMsValue, + clientDisconnected, + ) + } + imageCount := imageCounter.Count() + result := &OpenAIForwardResult{ + RequestID: responseID, + Usage: usage, + Model: originalModel, + UpstreamModel: mappedModel, + ServiceTier: extractOpenAIServiceTierFromBody(payload), + ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, originalModel), payload, mappedModel), + Stream: reqStream, + OpenAIWSMode: true, + ResponseHeaders: lease.HandshakeHeaders(), + Duration: time.Since(turnStart), + FirstTokenMs: firstTokenMs, + } + if replayInput := replayCollector.Items(); len(replayInput) > 0 { + result.wsReplayInput = replayInput + result.wsReplayInputExists = true + } + if imageCount > 0 { + result.ImageCount = imageCount + result.ImageSize = imageSizeTier + result.ImageInputSize = imageInputSize + result.ImageOutputSizes = imageCounter.Sizes() + result.BillingModel = imageBillingModel + } + return result, nil + } + } + } + + currentPayload := firstPayload.payloadRaw + currentOriginalModel := firstPayload.originalModel + currentImageBillingModel := firstPayload.imageBillingModel + currentImageSizeTier := firstPayload.imageSizeTier + currentImageInputSize := firstPayload.imageInputSize + currentPayloadBytes := firstPayload.payloadBytes + isStrictAffinityTurn := func(payload []byte) bool { + if !storeDisabled { + return false + } + return strings.TrimSpace(openAIWSPayloadStringFromRaw(payload, "previous_response_id")) != "" + } + var sessionLease *openAIWSConnLease + sessionConnID := "" + pinnedSessionConnID := "" + unpinSessionConn := func(connID string) { + connID = strings.TrimSpace(connID) + if connID == "" || pinnedSessionConnID != connID { + return + } + pool.UnpinConn(account.ID, connID) + pinnedSessionConnID = "" + } + pinSessionConn := func(connID string) { + if !storeDisabled { + return + } + connID = strings.TrimSpace(connID) + if connID == "" || pinnedSessionConnID == connID { + return + } + if pinnedSessionConnID != "" { + pool.UnpinConn(account.ID, pinnedSessionConnID) + pinnedSessionConnID = "" + } + if pool.PinConn(account.ID, connID) { + pinnedSessionConnID = connID + } + } + // lastTurnClean 标记最后一轮 sendAndRelay 是否正常完成（收到终端事件且客户端未断连）。 + // 所有异常路径（读写错误、error 事件、客户端断连）已在各自分支或上层（L3403）中 MarkBroken， + // 因此 releaseSessionLease 中只需在非正常结束时 MarkBroken。 + lastTurnClean := false + releaseSessionLease := func() { + if sessionLease == nil { + return + } + if !lastTurnClean { + sessionLease.MarkBroken() + } + unpinSessionConn(sessionConnID) + sessionLease.Release() + if debugEnabled { + logOpenAIWSModeDebug( + "ingress_ws_upstream_released account_id=%d conn_id=%s", + account.ID, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + ) + } + } + defer releaseSessionLease() + + turn := 1 + turnRetry := 0 + turnPrevRecoveryTried := false + lastTurnFinishedAt := time.Time{} + lastTurnResponseID := "" + lastTurnPayload := []byte(nil) + var lastTurnStrictState *openAIWSIngressPreviousTurnStrictState + lastTurnReplayInput := []json.RawMessage(nil) + lastTurnReplayInputExists := false + currentTurnReplayInput := []json.RawMessage(nil) + currentTurnReplayInputExists := false + skipBeforeTurn := false + hasCurrentOrReplayFunctionCallOutput := func(payload []byte) bool { + if openAIWSRawPayloadHasToolCallOutput(payload) { + return true + } + return currentTurnReplayInputExists && openAIWSRawItemsHasFunctionCallOutput(currentTurnReplayInput) + } + resetSessionLease := func(markBroken bool) { + if sessionLease == nil { + return + } + if markBroken { + sessionLease.MarkBroken() + } + releaseSessionLease() + sessionLease = nil + sessionConnID = "" + preferredConnID = "" + } + recoverIngressPrevResponseNotFound := func(relayErr error, turn int, connID string) bool { + if !isOpenAIWSIngressPreviousResponseNotFound(relayErr) { + return false + } + if turnPrevRecoveryTried || !s.openAIWSIngressPreviousResponseRecoveryEnabled() { + return false + } + // 携带 function_call_output 的请求不能丢弃 previous_response_id： + // 上游 API 需要 response chain 来匹配 tool_result 与之前的 tool_use， + // 丢弃后会导致 "No tool call found for function call output" 400 错误。 + if hasCurrentOrReplayFunctionCallOutput(currentPayload) { + return false + } + if isStrictAffinityTurn(currentPayload) { + // Layer 2：严格亲和链路命中 previous_response_not_found 时，降级为“去掉 previous_response_id 后重放一次”。 + // 该错误说明续链锚点已失效，继续 strict fail-close 只会直接中断本轮请求。 + logOpenAIWSModeInfo( + "ingress_ws_prev_response_recovery_layer2 account_id=%d turn=%d conn_id=%s store_disabled_conn_mode=%s action=drop_previous_response_id_retry", + account.ID, + turn, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(storeDisabledConnMode), + ) + } + turnPrevRecoveryTried = true + updatedPayload, removed, dropErr := dropPreviousResponseIDFromRawPayload(currentPayload) + if dropErr != nil || !removed { + reason := "not_removed" + if dropErr != nil { + reason = "drop_error" + } + logOpenAIWSModeInfo( + "ingress_ws_prev_response_recovery_skip account_id=%d turn=%d conn_id=%s reason=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(reason), + ) + return false + } + updatedWithInput, setInputErr := setOpenAIWSPayloadInputSequence( + updatedPayload, + currentTurnReplayInput, + currentTurnReplayInputExists, + ) + if setInputErr != nil { + logOpenAIWSModeInfo( + "ingress_ws_prev_response_recovery_skip account_id=%d turn=%d conn_id=%s reason=set_full_input_error cause=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(setInputErr.Error(), openAIWSLogValueMaxLen), + ) + return false + } + logOpenAIWSModeInfo( + "ingress_ws_prev_response_recovery account_id=%d turn=%d conn_id=%s action=drop_previous_response_id retry=1", + account.ID, + turn, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + ) + currentPayload = updatedWithInput + currentPayloadBytes = len(updatedWithInput) + resetSessionLease(true) + skipBeforeTurn = true + return true + } + retryIngressTurn := func(relayErr error, turn int, connID string) bool { + if !isOpenAIWSIngressTurnRetryable(relayErr) || turnRetry >= 1 { + return false + } + if isStrictAffinityTurn(currentPayload) { + logOpenAIWSModeInfo( + "ingress_ws_turn_retry_skip account_id=%d turn=%d conn_id=%s reason=strict_affinity", + account.ID, + turn, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + ) + return false + } + turnRetry++ + logOpenAIWSModeInfo( + "ingress_ws_turn_retry account_id=%d turn=%d retry=%d reason=%s conn_id=%s", + account.ID, + turn, + turnRetry, + truncateOpenAIWSLogValue(openAIWSIngressTurnRetryReason(relayErr), openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + ) + resetSessionLease(true) + skipBeforeTurn = true + return true + } + for { + if turn > 1 && !skipBeforeTurn && hooks != nil && hooks.BeforeRequest != nil { + if err := hooks.BeforeRequest(turn, currentPayload, currentOriginalModel); err != nil { + return err + } + } + if !skipBeforeTurn && hooks != nil && hooks.BeforeTurn != nil { + if err := hooks.BeforeTurn(turn); err != nil { + return err + } + } + skipBeforeTurn = false + currentPreviousResponseID := openAIWSPayloadStringFromRaw(currentPayload, "previous_response_id") + expectedPrev := strings.TrimSpace(lastTurnResponseID) + toolSignals := ToolContinuationSignals{ + HasFunctionCallOutput: openAIWSRawPayloadHasToolCallOutput(currentPayload), + } + if toolSignals.HasFunctionCallOutput { + var currentReqBody map[string]any + if err := json.Unmarshal(currentPayload, ¤tReqBody); err == nil { + toolSignals = AnalyzeToolContinuationSignals(currentReqBody) + } + } + hasFunctionCallOutput := toolSignals.HasFunctionCallOutput + // store=false + function_call_output 场景必须有续链锚点。 + // 若客户端未传 previous_response_id，优先回填上一轮响应 ID，避免上游报 call_id 无法关联。 + if shouldInferIngressFunctionCallOutputPreviousResponseID( + storeDisabled, + turn, + toolSignals, + currentPreviousResponseID, + expectedPrev, + ) { + updatedPayload, setPrevErr := setPreviousResponseIDToRawPayload(currentPayload, expectedPrev) + if setPrevErr != nil { + logOpenAIWSModeInfo( + "ingress_ws_function_call_output_prev_infer_skip account_id=%d turn=%d conn_id=%s reason=set_previous_response_id_error cause=%s expected_previous_response_id=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(setPrevErr.Error(), openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), + ) + } else { + currentPayload = updatedPayload + currentPayloadBytes = len(updatedPayload) + currentPreviousResponseID = expectedPrev + logOpenAIWSModeInfo( + "ingress_ws_function_call_output_prev_infer account_id=%d turn=%d conn_id=%s action=set_previous_response_id previous_response_id=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), + ) + } + } + nextReplayInput, nextReplayInputExists, replayInputErr := buildOpenAIWSReplayInputSequence( + lastTurnReplayInput, + lastTurnReplayInputExists, + currentPayload, + currentPreviousResponseID != "", + ) + if replayInputErr != nil { + logOpenAIWSModeInfo( + "ingress_ws_replay_input_skip account_id=%d turn=%d conn_id=%s reason=build_error cause=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(replayInputErr.Error(), openAIWSLogValueMaxLen), + ) + currentTurnReplayInput = nil + currentTurnReplayInputExists = false + } else { + currentTurnReplayInput = nextReplayInput + currentTurnReplayInputExists = nextReplayInputExists + } + replayHasFunctionCallOutput := currentTurnReplayInputExists && + openAIWSRawItemsHasFunctionCallOutput(currentTurnReplayInput) + hasFunctionCallOutput = hasFunctionCallOutput || replayHasFunctionCallOutput + if storeDisabled && turn > 1 && currentPreviousResponseID != "" { + shouldKeepPreviousResponseID := false + strictReason := "" + var strictErr error + if lastTurnStrictState != nil { + shouldKeepPreviousResponseID, strictReason, strictErr = shouldKeepIngressPreviousResponseIDWithStrictState( + lastTurnStrictState, + currentPayload, + lastTurnResponseID, + hasFunctionCallOutput, + ) + } else { + shouldKeepPreviousResponseID, strictReason, strictErr = shouldKeepIngressPreviousResponseID( + lastTurnPayload, + currentPayload, + lastTurnResponseID, + hasFunctionCallOutput, + ) + } + if strictErr != nil { + logOpenAIWSModeInfo( + "ingress_ws_prev_response_strict_eval account_id=%d turn=%d conn_id=%s action=keep_previous_response_id reason=%s cause=%s previous_response_id=%s expected_previous_response_id=%s has_function_call_output=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(strictReason), + truncateOpenAIWSLogValue(strictErr.Error(), openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), + hasFunctionCallOutput, + ) + } else if !shouldKeepPreviousResponseID { + updatedPayload, removed, dropErr := dropPreviousResponseIDFromRawPayload(currentPayload) + if dropErr != nil || !removed { + dropReason := "not_removed" + if dropErr != nil { + dropReason = "drop_error" + } + logOpenAIWSModeInfo( + "ingress_ws_prev_response_strict_eval account_id=%d turn=%d conn_id=%s action=keep_previous_response_id reason=%s drop_reason=%s previous_response_id=%s expected_previous_response_id=%s has_function_call_output=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(strictReason), + normalizeOpenAIWSLogValue(dropReason), + truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), + hasFunctionCallOutput, + ) + } else { + updatedWithInput, setInputErr := setOpenAIWSPayloadInputSequence( + updatedPayload, + currentTurnReplayInput, + currentTurnReplayInputExists, + ) + if setInputErr != nil { + logOpenAIWSModeInfo( + "ingress_ws_prev_response_strict_eval account_id=%d turn=%d conn_id=%s action=keep_previous_response_id reason=%s drop_reason=set_full_input_error previous_response_id=%s expected_previous_response_id=%s cause=%s has_function_call_output=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(strictReason), + truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(setInputErr.Error(), openAIWSLogValueMaxLen), + hasFunctionCallOutput, + ) + } else { + currentPayload = updatedWithInput + currentPayloadBytes = len(updatedWithInput) + logOpenAIWSModeInfo( + "ingress_ws_prev_response_strict_eval account_id=%d turn=%d conn_id=%s action=drop_previous_response_id_full_create reason=%s previous_response_id=%s expected_previous_response_id=%s has_function_call_output=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(strictReason), + truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), + hasFunctionCallOutput, + ) + currentPreviousResponseID = "" + } + } + } + } + forcePreferredConn := isStrictAffinityTurn(currentPayload) + if sessionLease == nil { + acquiredLease, acquireErr := acquireTurnLease(turn, preferredConnID, forcePreferredConn) + if acquireErr != nil { + return fmt.Errorf("acquire upstream websocket: %w", acquireErr) + } + sessionLease = acquiredLease + sessionConnID = strings.TrimSpace(sessionLease.ConnID()) + if storeDisabled { + pinSessionConn(sessionConnID) + } else { + unpinSessionConn(sessionConnID) + } + } + shouldPreflightPing := turn > 1 && sessionLease != nil && turnRetry == 0 + if shouldPreflightPing && openAIWSIngressPreflightPingIdle > 0 && !lastTurnFinishedAt.IsZero() { + if time.Since(lastTurnFinishedAt) < openAIWSIngressPreflightPingIdle { + shouldPreflightPing = false + } + } + if shouldPreflightPing { + if pingErr := sessionLease.PingWithTimeout(openAIWSConnHealthCheckTO); pingErr != nil { + logOpenAIWSModeInfo( + "ingress_ws_upstream_preflight_ping_fail account_id=%d turn=%d conn_id=%s cause=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(pingErr.Error(), openAIWSLogValueMaxLen), + ) + if forcePreferredConn { + // 携带 function_call_output 的请求不能丢弃 previous_response_id： + // 上游 API 需要 response chain 来匹配 tool_result 与之前的 tool_use， + // 除非 replay input 已经包含与每个 tool_result 匹配的 tool_use 上下文。 + hasFCOutput := hasFunctionCallOutput + hasReplayToolContext := hasFCOutput && + currentTurnReplayInputExists && + openAIWSRawItemsHaveToolCallContextForOutputs(currentTurnReplayInput) + if !turnPrevRecoveryTried && currentPreviousResponseID != "" && (!hasFCOutput || hasReplayToolContext) { + updatedPayload, removed, dropErr := dropPreviousResponseIDFromRawPayload(currentPayload) + if dropErr != nil || !removed { + reason := "not_removed" + if dropErr != nil { + reason = "drop_error" + } + logOpenAIWSModeInfo( + "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=%s previous_response_id=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(reason), + truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + ) + } else { + updatedWithInput, setInputErr := setOpenAIWSPayloadInputSequence( + updatedPayload, + currentTurnReplayInput, + currentTurnReplayInputExists, + ) + if setInputErr != nil { + logOpenAIWSModeInfo( + "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=set_full_input_error previous_response_id=%s cause=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(setInputErr.Error(), openAIWSLogValueMaxLen), + ) + } else { + logOpenAIWSModeInfo( + "ingress_ws_preflight_ping_recovery account_id=%d turn=%d conn_id=%s action=drop_previous_response_id_retry previous_response_id=%s has_function_call_output=%v has_replay_tool_context=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + hasFCOutput, + hasReplayToolContext, + ) + turnPrevRecoveryTried = true + currentPayload = updatedWithInput + currentPayloadBytes = len(updatedWithInput) + resetSessionLease(true) + skipBeforeTurn = true + continue + } + } + } + if hasFCOutput && currentPreviousResponseID != "" { + reason := "function_call_output_missing_replay_context" + if hasReplayToolContext { + reason = "function_call_output_replay_not_applied" + } + logOpenAIWSModeInfo( + "ingress_ws_preflight_ping_recovery_skip account_id=%d turn=%d conn_id=%s reason=%s action=fail_close previous_response_id=%s has_replay_tool_context=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + reason, + truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + hasReplayToolContext, + ) + } + resetSessionLease(true) + return NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + "upstream continuation connection is unavailable; please restart the conversation", + pingErr, + ) + } + resetSessionLease(true) + + acquiredLease, acquireErr := acquireTurnLease(turn, preferredConnID, forcePreferredConn) + if acquireErr != nil { + return fmt.Errorf("acquire upstream websocket after preflight ping fail: %w", acquireErr) + } + sessionLease = acquiredLease + sessionConnID = strings.TrimSpace(sessionLease.ConnID()) + if storeDisabled { + pinSessionConn(sessionConnID) + } + } + } + connID := sessionConnID + if currentPreviousResponseID != "" { + chainedFromLast := expectedPrev != "" && currentPreviousResponseID == expectedPrev + currentPreviousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(currentPreviousResponseID) + logOpenAIWSModeInfo( + "ingress_ws_turn_chain account_id=%d turn=%d conn_id=%s previous_response_id=%s previous_response_id_kind=%s last_turn_response_id=%s chained_from_last=%v preferred_conn_id=%s header_session_id=%s header_conversation_id=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v store_disabled=%v", + account.ID, + turn, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(currentPreviousResponseID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(currentPreviousResponseIDKind), + truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), + chainedFromLast, + truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), + openAIWSHeaderValueForLog(baseAcquireReq.Headers, "session_id"), + openAIWSHeaderValueForLog(baseAcquireReq.Headers, "conversation_id"), + turnState != "", + len(turnState), + openAIWSPayloadStringFromRaw(currentPayload, "prompt_cache_key") != "", + storeDisabled, + ) + } + + result, relayErr := sendAndRelay(turn, sessionLease, currentPayload, currentPayloadBytes, currentOriginalModel, currentImageBillingModel, currentImageSizeTier, currentImageInputSize) + if relayErr != nil { + lastTurnClean = false + if recoverIngressPrevResponseNotFound(relayErr, turn, connID) { + continue + } + if retryIngressTurn(relayErr, turn, connID) { + continue + } + finalErr := relayErr + if unwrapped := errors.Unwrap(relayErr); unwrapped != nil { + finalErr = unwrapped + } + if hooks != nil && hooks.AfterTurn != nil { + hooks.AfterTurn(turn, nil, finalErr) + } + sessionLease.MarkBroken() + return finalErr + } + turnRetry = 0 + turnPrevRecoveryTried = false + lastTurnFinishedAt = time.Now() + lastTurnClean = true + if hooks != nil && hooks.AfterTurn != nil { + hooks.AfterTurn(turn, result, nil) + } + if result == nil { + return errors.New("websocket turn result is nil") + } + responseID := strings.TrimSpace(result.RequestID) + lastTurnResponseID = responseID + lastTurnPayload = cloneOpenAIWSPayloadBytes(currentPayload) + lastTurnReplayInput = cloneOpenAIWSRawMessages(currentTurnReplayInput) + lastTurnReplayInputExists = currentTurnReplayInputExists + if result.wsReplayInputExists { + lastTurnReplayInput = append(lastTurnReplayInput, cloneOpenAIWSRawMessages(result.wsReplayInput)...) + lastTurnReplayInputExists = true + } + nextStrictState, strictStateErr := buildOpenAIWSIngressPreviousTurnStrictState(currentPayload) + if strictStateErr != nil { + lastTurnStrictState = nil + logOpenAIWSModeInfo( + "ingress_ws_prev_response_strict_state_skip account_id=%d turn=%d conn_id=%s reason=build_error cause=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(strictStateErr.Error(), openAIWSLogValueMaxLen), + ) + } else { + lastTurnStrictState = nextStrictState + } + + if responseID != "" && stateStore != nil { + ttl := s.openAIWSResponseStickyTTL() + logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) + stateStore.BindResponseConn(responseID, connID, ttl) + } + if stateStore != nil && storeDisabled && sessionHash != "" { + stateStore.BindSessionConn(groupID, sessionHash, connID, s.openAIWSSessionStickyTTL()) + } + if connID != "" { + preferredConnID = connID + } + + nextClientMessage, readErr := readClientMessage() + if readErr != nil { + if isOpenAIWSClientDisconnectError(readErr) { + closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) + logOpenAIWSModeInfo( + "ingress_ws_client_closed account_id=%d conn_id=%s close_status=%s close_reason=%s", + account.ID, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + closeStatus, + truncateOpenAIWSLogValue(closeReason, openAIWSHeaderValueMaxLen), + ) + return nil + } + return fmt.Errorf("read client websocket request: %w", readErr) + } + + nextPayload, parseErr := parseClientPayload(nextClientMessage) + if parseErr != nil { + return parseErr + } + if nextPayload.promptCacheKey != "" { + // ingress 会话在整个客户端 WS 生命周期内复用同一上游连接； + // prompt_cache_key 对握手头的更新仅在未来需要重新建连时生效。 + updatedHeaders, _, updHdrErr := s.buildOpenAIWSHeaders(ctx, c, account, token, wsDecision, isCodexCLI, turnState, strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)), nextPayload.promptCacheKey) + if updHdrErr != nil { + logOpenAIWSModeInfo("ingress_ws_update_headers_failed account_id=%d err=%v", account.ID, updHdrErr) + } else { + baseAcquireReq.Headers = updatedHeaders + } + } + if nextPayload.previousResponseID != "" { + expectedPrev := strings.TrimSpace(lastTurnResponseID) + chainedFromLast := expectedPrev != "" && nextPayload.previousResponseID == expectedPrev + nextPreviousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(nextPayload.previousResponseID) + logOpenAIWSModeInfo( + "ingress_ws_next_turn_chain account_id=%d turn=%d next_turn=%d conn_id=%s previous_response_id=%s previous_response_id_kind=%s last_turn_response_id=%s chained_from_last=%v has_prompt_cache_key=%v store_disabled=%v", + account.ID, + turn, + turn+1, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(nextPayload.previousResponseID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(nextPreviousResponseIDKind), + truncateOpenAIWSLogValue(expectedPrev, openAIWSIDValueMaxLen), + chainedFromLast, + nextPayload.promptCacheKey != "", + storeDisabled, + ) + } + if stateStore != nil && nextPayload.previousResponseID != "" { + if stickyConnID, ok := stateStore.GetResponseConn(nextPayload.previousResponseID); ok { + if sessionConnID != "" && stickyConnID != "" && stickyConnID != sessionConnID { + logOpenAIWSModeInfo( + "ingress_ws_keep_session_conn account_id=%d turn=%d conn_id=%s sticky_conn_id=%s previous_response_id=%s", + account.ID, + turn, + truncateOpenAIWSLogValue(sessionConnID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(stickyConnID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(nextPayload.previousResponseID, openAIWSIDValueMaxLen), + ) + } else { + preferredConnID = stickyConnID + } + } + } + currentPayload = nextPayload.payloadRaw + currentOriginalModel = nextPayload.originalModel + currentImageBillingModel = nextPayload.imageBillingModel + currentImageSizeTier = nextPayload.imageSizeTier + currentImageInputSize = nextPayload.imageInputSize + currentPayloadBytes = nextPayload.payloadBytes + storeDisabled = s.isOpenAIWSStoreDisabledInRequestRaw(currentPayload, account) + if !storeDisabled { + unpinSessionConn(sessionConnID) + } + turn++ + } +} diff --git a/backend/internal/service/openai_ws_forwarder_logutil.go b/backend/internal/service/openai_ws_forwarder_logutil.go new file mode 100644 index 0000000000..611938091d --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_logutil.go @@ -0,0 +1,690 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "sort" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "go.uber.org/zap" +) + +func normalizeOpenAIWSLogValue(value string) string { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + return "-" + } + return openAIWSLogValueReplacer.Replace(trimmed) +} + +func truncateOpenAIWSLogValue(value string, maxLen int) string { + normalized := normalizeOpenAIWSLogValue(value) + if normalized == "-" || maxLen <= 0 { + return normalized + } + if len(normalized) <= maxLen { + return normalized + } + return normalized[:maxLen] + "..." +} + +func openAIWSHeaderValueForLog(headers http.Header, key string) string { + if headers == nil { + return "-" + } + return truncateOpenAIWSLogValue(headers.Get(key), openAIWSHeaderValueMaxLen) +} + +func hasOpenAIWSHeader(headers http.Header, key string) bool { + if headers == nil { + return false + } + return strings.TrimSpace(headers.Get(key)) != "" +} + +type openAIWSSessionHeaderResolution struct { + SessionID string + ConversationID string + SessionSource string + ConversationSource string +} + +func resolveOpenAIWSSessionHeaders(c *gin.Context, promptCacheKey string) openAIWSSessionHeaderResolution { + resolution := openAIWSSessionHeaderResolution{ + SessionSource: "none", + ConversationSource: "none", + } + if c != nil && c.Request != nil { + if sessionID := strings.TrimSpace(c.Request.Header.Get("session_id")); sessionID != "" { + resolution.SessionID = sessionID + resolution.SessionSource = "header_session_id" + } + if conversationID := strings.TrimSpace(c.Request.Header.Get("conversation_id")); conversationID != "" { + resolution.ConversationID = conversationID + resolution.ConversationSource = "header_conversation_id" + if resolution.SessionID == "" { + resolution.SessionID = conversationID + resolution.SessionSource = "header_conversation_id" + } + } + } + + cacheKey := strings.TrimSpace(promptCacheKey) + if cacheKey != "" { + if resolution.SessionID == "" { + resolution.SessionID = cacheKey + resolution.SessionSource = "prompt_cache_key" + } + } + return resolution +} + +func shouldLogOpenAIWSEvent(idx int, eventType string) bool { + if idx <= openAIWSEventLogHeadLimit { + return true + } + if openAIWSEventLogEveryN > 0 && idx%openAIWSEventLogEveryN == 0 { + return true + } + if eventType == "error" || isOpenAIWSTerminalEvent(eventType) { + return true + } + return false +} + +func shouldLogOpenAIWSBufferedEvent(idx int) bool { + if idx <= openAIWSBufferLogHeadLimit { + return true + } + if openAIWSBufferLogEveryN > 0 && idx%openAIWSBufferLogEveryN == 0 { + return true + } + return false +} + +func openAIWSEventMayContainModel(eventType string) bool { + switch eventType { + case "response.created", + "response.in_progress", + "response.completed", + "response.done", + "response.failed", + "response.incomplete", + "response.cancelled", + "response.canceled": + return true + default: + trimmed := strings.TrimSpace(eventType) + if trimmed == eventType { + return false + } + switch trimmed { + case "response.created", + "response.in_progress", + "response.completed", + "response.done", + "response.failed", + "response.incomplete", + "response.cancelled", + "response.canceled": + return true + default: + return false + } + } +} + +func openAIWSEventMayContainToolCalls(eventType string) bool { + eventType = strings.TrimSpace(eventType) + if eventType == "" { + return false + } + if strings.Contains(eventType, "function_call") || strings.Contains(eventType, "tool_call") { + return true + } + switch eventType { + case "response.output_item.added", "response.output_item.done", "response.completed", "response.done": + return true + default: + return false + } +} + +func openAIWSEventShouldParseUsage(eventType string) bool { + switch strings.TrimSpace(eventType) { + case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled": + return true + default: + return false + } +} + +func parseOpenAIWSEventEnvelope(message []byte) (eventType string, responseID string, response gjson.Result) { + if len(message) == 0 { + return "", "", gjson.Result{} + } + values := gjson.GetManyBytes(message, "type", "response.id", "id", "response") + eventType = strings.TrimSpace(values[0].String()) + if id := strings.TrimSpace(values[1].String()); id != "" { + responseID = id + } else { + responseID = strings.TrimSpace(values[2].String()) + } + return eventType, responseID, values[3] +} + +func openAIWSMessageLikelyContainsToolCalls(message []byte) bool { + if len(message) == 0 { + return false + } + return bytes.Contains(message, []byte(`"tool_calls"`)) || + bytes.Contains(message, []byte(`"tool_call"`)) || + bytes.Contains(message, []byte(`"function_call"`)) +} + +func parseOpenAIWSResponseUsageFromCompletedEvent(message []byte, usage *OpenAIUsage) { + if usage == nil || len(message) == 0 { + return + } + if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(message); ok { + *usage = parsedUsage + } +} + +func parseOpenAIWSErrorEventFields(message []byte) (code string, errType string, errMessage string) { + if len(message) == 0 { + return "", "", "" + } + values := gjson.GetManyBytes(message, "error.code", "error.type", "error.message") + return strings.TrimSpace(values[0].String()), strings.TrimSpace(values[1].String()), strings.TrimSpace(values[2].String()) +} + +func summarizeOpenAIWSErrorEventFieldsFromRaw(codeRaw, errTypeRaw, errMessageRaw string) (code string, errType string, errMessage string) { + code = truncateOpenAIWSLogValue(codeRaw, openAIWSLogValueMaxLen) + errType = truncateOpenAIWSLogValue(errTypeRaw, openAIWSLogValueMaxLen) + errMessage = truncateOpenAIWSLogValue(errMessageRaw, openAIWSLogValueMaxLen) + return code, errType, errMessage +} + +func summarizeOpenAIWSErrorEventFields(message []byte) (code string, errType string, errMessage string) { + if len(message) == 0 { + return "-", "-", "-" + } + return summarizeOpenAIWSErrorEventFieldsFromRaw(parseOpenAIWSErrorEventFields(message)) +} + +func summarizeOpenAIWSPayloadKeySizes(payload map[string]any, topN int) string { + if len(payload) == 0 { + return "-" + } + type keySize struct { + Key string + Size int + } + sizes := make([]keySize, 0, len(payload)) + for key, value := range payload { + size := estimateOpenAIWSPayloadValueSize(value, openAIWSPayloadSizeEstimateDepth) + sizes = append(sizes, keySize{Key: key, Size: size}) + } + sort.Slice(sizes, func(i, j int) bool { + if sizes[i].Size == sizes[j].Size { + return sizes[i].Key < sizes[j].Key + } + return sizes[i].Size > sizes[j].Size + }) + + if topN <= 0 || topN > len(sizes) { + topN = len(sizes) + } + parts := make([]string, 0, topN) + for idx := 0; idx < topN; idx++ { + item := sizes[idx] + parts = append(parts, fmt.Sprintf("%s:%d", item.Key, item.Size)) + } + return strings.Join(parts, ",") +} + +func estimateOpenAIWSPayloadValueSize(value any, depth int) int { + if depth <= 0 { + return -1 + } + switch v := value.(type) { + case nil: + return 0 + case string: + return len(v) + case []byte: + return len(v) + case bool: + return 1 + case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64: + return 8 + case float32, float64: + return 8 + case map[string]any: + if len(v) == 0 { + return 2 + } + total := 2 + count := 0 + for key, item := range v { + count++ + if count > openAIWSPayloadSizeEstimateMaxItems { + return -1 + } + itemSize := estimateOpenAIWSPayloadValueSize(item, depth-1) + if itemSize < 0 { + return -1 + } + total += len(key) + itemSize + 3 + if total > openAIWSPayloadSizeEstimateMaxBytes { + return -1 + } + } + return total + case []any: + if len(v) == 0 { + return 2 + } + total := 2 + limit := len(v) + if limit > openAIWSPayloadSizeEstimateMaxItems { + return -1 + } + for i := 0; i < limit; i++ { + itemSize := estimateOpenAIWSPayloadValueSize(v[i], depth-1) + if itemSize < 0 { + return -1 + } + total += itemSize + 1 + if total > openAIWSPayloadSizeEstimateMaxBytes { + return -1 + } + } + return total + default: + raw, err := json.Marshal(v) + if err != nil { + return -1 + } + if len(raw) > openAIWSPayloadSizeEstimateMaxBytes { + return -1 + } + return len(raw) + } +} + +func openAIWSPayloadString(payload map[string]any, key string) string { + if len(payload) == 0 { + return "" + } + raw, ok := payload[key] + if !ok { + return "" + } + switch v := raw.(type) { + case nil: + return "" + case string: + return strings.TrimSpace(v) + case []byte: + return strings.TrimSpace(string(v)) + default: + return "" + } +} + +func openAIWSPayloadStringFromRaw(payload []byte, key string) string { + if len(payload) == 0 || strings.TrimSpace(key) == "" { + return "" + } + return strings.TrimSpace(gjson.GetBytes(payload, key).String()) +} + +func openAIWSPayloadBoolFromRaw(payload []byte, key string, defaultValue bool) bool { + if len(payload) == 0 || strings.TrimSpace(key) == "" { + return defaultValue + } + value := gjson.GetBytes(payload, key) + if !value.Exists() { + return defaultValue + } + if value.Type != gjson.True && value.Type != gjson.False { + return defaultValue + } + return value.Bool() +} + +func openAIWSSessionHashesFromID(sessionID string) (string, string) { + return deriveOpenAISessionHashes(sessionID) +} + +func extractOpenAIWSImageURL(value any) string { + switch v := value.(type) { + case string: + return strings.TrimSpace(v) + case map[string]any: + if raw, ok := v["url"].(string); ok { + return strings.TrimSpace(raw) + } + } + return "" +} + +func summarizeOpenAIWSInput(input any) string { + items, ok := input.([]any) + if !ok || len(items) == 0 { + return "-" + } + + itemCount := len(items) + textChars := 0 + imageDataURLs := 0 + imageDataURLChars := 0 + imageRemoteURLs := 0 + + handleContentItem := func(contentItem map[string]any) { + contentType, _ := contentItem["type"].(string) + switch strings.TrimSpace(contentType) { + case "input_text", "output_text", "text": + if text, ok := contentItem["text"].(string); ok { + textChars += len(text) + } + case "input_image": + imageURL := extractOpenAIWSImageURL(contentItem["image_url"]) + if imageURL == "" { + return + } + if strings.HasPrefix(strings.ToLower(imageURL), "data:image/") { + imageDataURLs++ + imageDataURLChars += len(imageURL) + return + } + imageRemoteURLs++ + } + } + + handleInputItem := func(inputItem map[string]any) { + if content, ok := inputItem["content"].([]any); ok { + for _, rawContent := range content { + contentItem, ok := rawContent.(map[string]any) + if !ok { + continue + } + handleContentItem(contentItem) + } + return + } + + itemType, _ := inputItem["type"].(string) + switch strings.TrimSpace(itemType) { + case "input_text", "output_text", "text": + if text, ok := inputItem["text"].(string); ok { + textChars += len(text) + } + case "input_image": + imageURL := extractOpenAIWSImageURL(inputItem["image_url"]) + if imageURL == "" { + return + } + if strings.HasPrefix(strings.ToLower(imageURL), "data:image/") { + imageDataURLs++ + imageDataURLChars += len(imageURL) + return + } + imageRemoteURLs++ + } + } + + for _, rawItem := range items { + inputItem, ok := rawItem.(map[string]any) + if !ok { + continue + } + handleInputItem(inputItem) + } + + return fmt.Sprintf( + "items=%d,text_chars=%d,image_data_urls=%d,image_data_url_chars=%d,image_remote_urls=%d", + itemCount, + textChars, + imageDataURLs, + imageDataURLChars, + imageRemoteURLs, + ) +} + +func dropOpenAIWSPayloadKey(payload map[string]any, key string, removed *[]string) { + if len(payload) == 0 || strings.TrimSpace(key) == "" { + return + } + if _, exists := payload[key]; !exists { + return + } + delete(payload, key) + *removed = append(*removed, key) +} + +// applyOpenAIWSRetryPayloadStrategy 在 WS 连续失败时仅移除无语义字段， +// 避免重试成功却改变原始请求语义。 +// 注意：prompt_cache_key 不应在重试中移除；它常用于会话稳定标识（session_id 兜底）。 +func applyOpenAIWSRetryPayloadStrategy(payload map[string]any, attempt int) (strategy string, removedKeys []string) { + if len(payload) == 0 { + return "empty", nil + } + if attempt <= 1 { + return "full", nil + } + + removed := make([]string, 0, 2) + if attempt >= 2 { + dropOpenAIWSPayloadKey(payload, "include", &removed) + } + + if len(removed) == 0 { + return "full", nil + } + sort.Strings(removed) + return "trim_optional_fields", removed +} + +func logOpenAIWSModeInfo(format string, args ...any) { + logger.LegacyPrintf("service.openai_gateway", "[OpenAI WS Mode][openai_ws_mode=true] "+format, args...) +} + +func isOpenAIWSModeDebugEnabled() bool { + return logger.L().Core().Enabled(zap.DebugLevel) +} + +func logOpenAIWSModeDebug(format string, args ...any) { + if !isOpenAIWSModeDebugEnabled() { + return + } + logger.LegacyPrintf("service.openai_gateway", "[debug] [OpenAI WS Mode][openai_ws_mode=true] "+format, args...) +} + +func logOpenAIWSBindResponseAccountWarn(groupID, accountID int64, responseID string, err error) { + if err == nil { + return + } + logger.L().Warn( + "openai.ws_bind_response_account_failed", + zap.Int64("group_id", groupID), + zap.Int64("account_id", accountID), + zap.String("response_id", truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen)), + zap.Error(err), + ) +} + +func summarizeOpenAIWSReadCloseError(err error) (status string, reason string) { + if err == nil { + return "-", "-" + } + statusCode := coderws.CloseStatus(err) + if statusCode == -1 { + return "-", "-" + } + closeStatus := fmt.Sprintf("%d(%s)", int(statusCode), statusCode.String()) + closeReason := "-" + var closeErr coderws.CloseError + if errors.As(err, &closeErr) { + reasonText := strings.TrimSpace(closeErr.Reason) + if reasonText != "" { + closeReason = normalizeOpenAIWSLogValue(reasonText) + } + } + return normalizeOpenAIWSLogValue(closeStatus), closeReason +} + +func unwrapOpenAIWSDialBaseError(err error) error { + if err == nil { + return nil + } + var dialErr *openAIWSDialError + if errors.As(err, &dialErr) && dialErr != nil && dialErr.Err != nil { + return dialErr.Err + } + return err +} + +func openAIWSDialRespHeaderForLog(err error, key string) string { + var dialErr *openAIWSDialError + if !errors.As(err, &dialErr) || dialErr == nil || dialErr.ResponseHeaders == nil { + return "-" + } + return truncateOpenAIWSLogValue(dialErr.ResponseHeaders.Get(key), openAIWSHeaderValueMaxLen) +} + +func classifyOpenAIWSDialError(err error) string { + if err == nil { + return "-" + } + baseErr := unwrapOpenAIWSDialBaseError(err) + if baseErr == nil { + return "-" + } + if errors.Is(baseErr, context.DeadlineExceeded) { + return "ctx_deadline_exceeded" + } + if errors.Is(baseErr, context.Canceled) { + return "ctx_canceled" + } + var netErr net.Error + if errors.As(baseErr, &netErr) && netErr.Timeout() { + return "net_timeout" + } + if status := coderws.CloseStatus(baseErr); status != -1 { + return normalizeOpenAIWSLogValue(fmt.Sprintf("ws_close_%d", int(status))) + } + message := strings.ToLower(strings.TrimSpace(baseErr.Error())) + switch { + case strings.Contains(message, "handshake not finished"): + return "handshake_not_finished" + case strings.Contains(message, "bad handshake"): + return "bad_handshake" + case strings.Contains(message, "connection refused"): + return "connection_refused" + case strings.Contains(message, "no such host"): + return "dns_not_found" + case strings.Contains(message, "tls"): + return "tls_error" + case strings.Contains(message, "i/o timeout"): + return "io_timeout" + case strings.Contains(message, "context deadline exceeded"): + return "ctx_deadline_exceeded" + default: + return "dial_error" + } +} + +func summarizeOpenAIWSDialError(err error) ( + statusCode int, + dialClass string, + closeStatus string, + closeReason string, + respServer string, + respVia string, + respCFRay string, + respRequestID string, +) { + dialClass = "-" + closeStatus = "-" + closeReason = "-" + respServer = "-" + respVia = "-" + respCFRay = "-" + respRequestID = "-" + if err == nil { + return + } + var dialErr *openAIWSDialError + if errors.As(err, &dialErr) && dialErr != nil { + statusCode = dialErr.StatusCode + respServer = openAIWSDialRespHeaderForLog(err, "server") + respVia = openAIWSDialRespHeaderForLog(err, "via") + respCFRay = openAIWSDialRespHeaderForLog(err, "cf-ray") + respRequestID = openAIWSDialRespHeaderForLog(err, "x-request-id") + } + dialClass = normalizeOpenAIWSLogValue(classifyOpenAIWSDialError(err)) + closeStatus, closeReason = summarizeOpenAIWSReadCloseError(unwrapOpenAIWSDialBaseError(err)) + return +} + +func isOpenAIWSClientDisconnectError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || errors.Is(err, context.Canceled) { + return true + } + switch coderws.CloseStatus(err) { + case coderws.StatusNormalClosure, coderws.StatusGoingAway, coderws.StatusNoStatusRcvd, coderws.StatusAbnormalClosure: + return true + } + message := strings.ToLower(strings.TrimSpace(err.Error())) + if message == "" { + return false + } + return strings.Contains(message, "failed to read frame header: eof") || + strings.Contains(message, "unexpected eof") || + strings.Contains(message, "use of closed network connection") || + strings.Contains(message, "connection reset by peer") || + strings.Contains(message, "broken pipe") || + strings.Contains(message, "an established connection was aborted") +} + +func classifyOpenAIWSReadFallbackReason(err error) string { + if err == nil { + return "read_event" + } + switch coderws.CloseStatus(err) { + case coderws.StatusPolicyViolation: + return "policy_violation" + case coderws.StatusMessageTooBig: + return "message_too_big" + default: + return "read_event" + } +} + +func sortedKeys(m map[string]any) []string { + if len(m) == 0 { + return nil + } + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} diff --git a/backend/internal/service/openai_ws_forwarder_payload.go b/backend/internal/service/openai_ws_forwarder_payload.go new file mode 100644 index 0000000000..cd84e6f25b --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_payload.go @@ -0,0 +1,705 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/url" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/openai" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func (s *OpenAIGatewayService) buildOpenAIResponsesWSURL(account *Account) (string, error) { + if account == nil { + return "", errors.New("account is nil") + } + var targetURL string + switch account.Type { + case AccountTypeOAuth: + targetURL = chatgptCodexURL + case AccountTypeAPIKey: + baseURL := account.GetOpenAIBaseURL() + if baseURL == "" { + targetURL = openaiPlatformAPIURL + } else { + validatedURL, err := s.validateUpstreamBaseURL(baseURL) + if err != nil { + return "", err + } + targetURL = buildOpenAIResponsesURL(validatedURL) + } + default: + targetURL = openaiPlatformAPIURL + } + + parsed, err := url.Parse(strings.TrimSpace(targetURL)) + if err != nil { + return "", fmt.Errorf("invalid target url: %w", err) + } + switch strings.ToLower(parsed.Scheme) { + case "https": + parsed.Scheme = "wss" + case "http": + parsed.Scheme = "ws" + case "wss", "ws": + // 保持不变 + default: + return "", fmt.Errorf("unsupported scheme for ws: %s", parsed.Scheme) + } + return parsed.String(), nil +} + +func (s *OpenAIGatewayService) buildOpenAIWSHeaders( + ctx context.Context, + c *gin.Context, + account *Account, + token string, + decision OpenAIWSProtocolDecision, + isCodexCLI bool, + turnState string, + turnMetadata string, + promptCacheKey string, +) (http.Header, openAIWSSessionHeaderResolution, error) { + headers := make(http.Header) + headers.Set("authorization", "Bearer "+token) + + sessionResolution := resolveOpenAIWSSessionHeaders(c, promptCacheKey) + if c != nil && c.Request != nil { + if v := strings.TrimSpace(c.Request.Header.Get("accept-language")); v != "" { + headers.Set("accept-language", v) + } + } + // OAuth 账号：将 apiKeyID 混入 session 标识符，防止跨用户会话碰撞。 + if account != nil && account.Type == AccountTypeOAuth { + apiKeyID := getAPIKeyIDFromContext(c) + if sessionResolution.SessionID != "" { + headers.Set("session_id", isolateOpenAISessionID(apiKeyID, sessionResolution.SessionID)) + } + if sessionResolution.ConversationID != "" { + headers.Set("conversation_id", isolateOpenAISessionID(apiKeyID, sessionResolution.ConversationID)) + } + } else { + if sessionResolution.SessionID != "" { + headers.Set("session_id", sessionResolution.SessionID) + } + if sessionResolution.ConversationID != "" { + headers.Set("conversation_id", sessionResolution.ConversationID) + } + } + if state := strings.TrimSpace(turnState); state != "" { + headers.Set(openAIWSTurnStateHeader, state) + } + if metadata := strings.TrimSpace(turnMetadata); metadata != "" { + headers.Set(openAIWSTurnMetadataHeader, metadata) + } + + if account != nil && account.Type == AccountTypeOAuth { + if err := resolveAndSetOpenAIChatGPTAccountHeaders(ctx, s.accountRepo, headers, account); err != nil { + return nil, sessionResolution, fmt.Errorf("resolve chatgpt account headers: %w", err) + } + headers.Set("originator", resolveOpenAIUpstreamOriginator(c, isCodexCLI)) + } + + betaValue := openAIWSBetaV2Value + if decision.Transport == OpenAIUpstreamTransportResponsesWebsocket { + betaValue = openAIWSBetaV1Value + } + headers.Set("OpenAI-Beta", betaValue) + + customUA := "" + if account != nil { + customUA = account.GetOpenAIUserAgent() + } + if strings.TrimSpace(customUA) != "" { + headers.Set("user-agent", customUA) + } else if c != nil { + if ua := strings.TrimSpace(c.GetHeader("User-Agent")); ua != "" { + headers.Set("user-agent", ua) + } + } + if s != nil && s.cfg != nil && s.cfg.Gateway.ForceCodexCLI { + headers.Set("user-agent", codexCLIUserAgent) + } + if account != nil && account.Type == AccountTypeOAuth && !openai.IsCodexCLIRequest(headers.Get("user-agent")) { + headers.Set("user-agent", codexCLIUserAgent) + } + + // 账号级请求头覆写（仅 openai api_key 账号启用时生效；OAuth 路径 no-op）。 + // 覆盖所有 WS 模式（ctx_pool/dedicated/passthrough）的握手头。 + account.ApplyHeaderOverrides(headers) + + return headers, sessionResolution, nil +} + +func (s *OpenAIGatewayService) buildOpenAIWSCreatePayload(reqBody map[string]any, account *Account) map[string]any { + // OpenAI WS Mode 协议：response.create 字段与 HTTP /responses 基本一致。 + // 保留 stream 字段（与 Codex CLI 一致），仅移除 background。 + payload := make(map[string]any, len(reqBody)+1) + for k, v := range reqBody { + payload[k] = v + } + + delete(payload, "background") + if _, exists := payload["stream"]; !exists { + payload["stream"] = true + } + payload["type"] = "response.create" + + // OAuth 默认保持 store=false，避免误依赖服务端历史。 + if account != nil && account.Type == AccountTypeOAuth && !s.isOpenAIWSStoreRecoveryAllowed(account) { + payload["store"] = false + } + return payload +} + +func setOpenAIWSTurnMetadata(payload map[string]any, turnMetadata string) { + if len(payload) == 0 { + return + } + metadata := strings.TrimSpace(turnMetadata) + if metadata == "" { + return + } + + switch existing := payload["client_metadata"].(type) { + case map[string]any: + existing[openAIWSTurnMetadataHeader] = metadata + payload["client_metadata"] = existing + case map[string]string: + next := make(map[string]any, len(existing)+1) + for k, v := range existing { + next[k] = v + } + next[openAIWSTurnMetadataHeader] = metadata + payload["client_metadata"] = next + default: + payload["client_metadata"] = map[string]any{ + openAIWSTurnMetadataHeader: metadata, + } + } +} + +func (s *OpenAIGatewayService) isOpenAIWSStoreRecoveryAllowed(account *Account) bool { + if account != nil && account.IsOpenAIWSAllowStoreRecoveryEnabled() { + return true + } + if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.AllowStoreRecovery { + return true + } + return false +} + +func (s *OpenAIGatewayService) isOpenAIWSStoreDisabledInRequest(reqBody map[string]any, account *Account) bool { + if account != nil && account.Type == AccountTypeOAuth && !s.isOpenAIWSStoreRecoveryAllowed(account) { + return true + } + if len(reqBody) == 0 { + return false + } + rawStore, ok := reqBody["store"] + if !ok { + return false + } + storeEnabled, ok := rawStore.(bool) + if !ok { + return false + } + return !storeEnabled +} + +func (s *OpenAIGatewayService) isOpenAIWSStoreDisabledInRequestRaw(reqBody []byte, account *Account) bool { + if account != nil && account.Type == AccountTypeOAuth && !s.isOpenAIWSStoreRecoveryAllowed(account) { + return true + } + if len(reqBody) == 0 { + return false + } + storeValue := gjson.GetBytes(reqBody, "store") + if !storeValue.Exists() { + return false + } + if storeValue.Type != gjson.True && storeValue.Type != gjson.False { + return false + } + return !storeValue.Bool() +} + +func (s *OpenAIGatewayService) openAIWSStoreDisabledConnMode() string { + if s == nil || s.cfg == nil { + return openAIWSStoreDisabledConnModeStrict + } + mode := strings.ToLower(strings.TrimSpace(s.cfg.Gateway.OpenAIWS.StoreDisabledConnMode)) + switch mode { + case openAIWSStoreDisabledConnModeStrict, openAIWSStoreDisabledConnModeAdaptive, openAIWSStoreDisabledConnModeOff: + return mode + case "": + // 兼容旧配置：仅配置了布尔开关时按旧语义推导。 + if s.cfg.Gateway.OpenAIWS.StoreDisabledForceNewConn { + return openAIWSStoreDisabledConnModeStrict + } + return openAIWSStoreDisabledConnModeOff + default: + return openAIWSStoreDisabledConnModeStrict + } +} + +func shouldForceNewConnOnStoreDisabled(mode, lastFailureReason string) bool { + switch mode { + case openAIWSStoreDisabledConnModeOff: + return false + case openAIWSStoreDisabledConnModeAdaptive: + reason := strings.TrimPrefix(strings.TrimSpace(lastFailureReason), "prewarm_") + switch reason { + case "policy_violation", "message_too_big", "auth_failed", "write_request", "write": + return true + default: + return false + } + default: + return true + } +} + +func dropPreviousResponseIDFromRawPayload(payload []byte) ([]byte, bool, error) { + return dropPreviousResponseIDFromRawPayloadWithDeleteFn(payload, sjson.DeleteBytes) +} + +func dropPreviousResponseIDFromRawPayloadWithDeleteFn( + payload []byte, + deleteFn func([]byte, string) ([]byte, error), +) ([]byte, bool, error) { + if len(payload) == 0 { + return payload, false, nil + } + if !gjson.GetBytes(payload, "previous_response_id").Exists() { + return payload, false, nil + } + if deleteFn == nil { + deleteFn = sjson.DeleteBytes + } + + updated := payload + for i := 0; i < openAIWSMaxPrevResponseIDDeletePasses && + gjson.GetBytes(updated, "previous_response_id").Exists(); i++ { + next, err := deleteFn(updated, "previous_response_id") + if err != nil { + return payload, false, err + } + updated = next + } + return updated, !gjson.GetBytes(updated, "previous_response_id").Exists(), nil +} + +func setPreviousResponseIDToRawPayload(payload []byte, previousResponseID string) ([]byte, error) { + normalizedPrevID := strings.TrimSpace(previousResponseID) + if len(payload) == 0 || normalizedPrevID == "" { + return payload, nil + } + updated, err := sjson.SetBytes(payload, "previous_response_id", normalizedPrevID) + if err == nil { + return updated, nil + } + + var reqBody map[string]any + if unmarshalErr := json.Unmarshal(payload, &reqBody); unmarshalErr != nil { + return nil, err + } + reqBody["previous_response_id"] = normalizedPrevID + rebuilt, marshalErr := json.Marshal(reqBody) + if marshalErr != nil { + return nil, marshalErr + } + return rebuilt, nil +} + +func shouldInferIngressFunctionCallOutputPreviousResponseID( + storeDisabled bool, + turn int, + signals ToolContinuationSignals, + currentPreviousResponseID string, + expectedPreviousResponseID string, +) bool { + if !storeDisabled || turn <= 1 || !signals.HasFunctionCallOutput { + return false + } + if strings.TrimSpace(currentPreviousResponseID) != "" { + return false + } + if signals.HasFunctionCallOutputMissingCallID { + return false + } + // If the client already sent the actual tool-call context, treat this as + // a full replay / self-contained continuation payload rather than + // downgrading it into an inferred delta continuation. item_reference alone + // is not enough on the store=false WS path: it still needs a valid prior + // response anchor so upstream can resolve the referenced function_call. + if signals.HasToolCallContext { + return false + } + return strings.TrimSpace(expectedPreviousResponseID) != "" +} + +func alignStoreDisabledPreviousResponseID( + payload []byte, + expectedPreviousResponseID string, +) ([]byte, bool, error) { + if len(payload) == 0 { + return payload, false, nil + } + expected := strings.TrimSpace(expectedPreviousResponseID) + if expected == "" { + return payload, false, nil + } + current := openAIWSPayloadStringFromRaw(payload, "previous_response_id") + if current == "" || current == expected { + return payload, false, nil + } + + withoutPrev, removed, dropErr := dropPreviousResponseIDFromRawPayload(payload) + if dropErr != nil { + return payload, false, dropErr + } + if !removed { + return payload, false, nil + } + updated, setErr := setPreviousResponseIDToRawPayload(withoutPrev, expected) + if setErr != nil { + return payload, false, setErr + } + return updated, true, nil +} + +func cloneOpenAIWSPayloadBytes(payload []byte) []byte { + if len(payload) == 0 { + return nil + } + cloned := make([]byte, len(payload)) + copy(cloned, payload) + return cloned +} + +func cloneOpenAIWSRawMessages(items []json.RawMessage) []json.RawMessage { + if items == nil { + return nil + } + cloned := make([]json.RawMessage, 0, len(items)) + for idx := range items { + cloned = append(cloned, json.RawMessage(cloneOpenAIWSPayloadBytes(items[idx]))) + } + return cloned +} + +func normalizeOpenAIWSJSONForCompare(raw []byte) ([]byte, error) { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 { + return nil, errors.New("json is empty") + } + var decoded any + if err := json.Unmarshal(trimmed, &decoded); err != nil { + return nil, err + } + return json.Marshal(decoded) +} + +func normalizeOpenAIWSJSONForCompareOrRaw(raw []byte) []byte { + normalized, err := normalizeOpenAIWSJSONForCompare(raw) + if err != nil { + return bytes.TrimSpace(raw) + } + return normalized +} + +func normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(payload []byte) ([]byte, error) { + if len(payload) == 0 { + return nil, errors.New("payload is empty") + } + var decoded map[string]any + if err := json.Unmarshal(payload, &decoded); err != nil { + return nil, err + } + delete(decoded, "input") + delete(decoded, "previous_response_id") + return json.Marshal(decoded) +} + +func openAIWSExtractNormalizedInputSequence(payload []byte) ([]json.RawMessage, bool, error) { + if len(payload) == 0 { + return nil, false, nil + } + inputValue := gjson.GetBytes(payload, "input") + if !inputValue.Exists() { + return nil, false, nil + } + if inputValue.Type == gjson.JSON { + raw := strings.TrimSpace(inputValue.Raw) + if strings.HasPrefix(raw, "[") { + var items []json.RawMessage + if err := json.Unmarshal([]byte(raw), &items); err != nil { + return nil, true, err + } + return items, true, nil + } + return []json.RawMessage{json.RawMessage(raw)}, true, nil + } + if inputValue.Type == gjson.String { + encoded, _ := json.Marshal(inputValue.String()) + return []json.RawMessage{encoded}, true, nil + } + return []json.RawMessage{json.RawMessage(inputValue.Raw)}, true, nil +} + +func openAIWSInputIsPrefixExtended(previousPayload, currentPayload []byte) (bool, error) { + previousItems, previousExists, prevErr := openAIWSExtractNormalizedInputSequence(previousPayload) + if prevErr != nil { + return false, prevErr + } + currentItems, currentExists, currentErr := openAIWSExtractNormalizedInputSequence(currentPayload) + if currentErr != nil { + return false, currentErr + } + if !previousExists && !currentExists { + return true, nil + } + if !previousExists { + return len(currentItems) == 0, nil + } + if !currentExists { + return len(previousItems) == 0, nil + } + if len(currentItems) < len(previousItems) { + return false, nil + } + + for idx := range previousItems { + previousNormalized := normalizeOpenAIWSJSONForCompareOrRaw(previousItems[idx]) + currentNormalized := normalizeOpenAIWSJSONForCompareOrRaw(currentItems[idx]) + if !bytes.Equal(previousNormalized, currentNormalized) { + return false, nil + } + } + return true, nil +} + +func openAIWSRawItemsHasPrefix(items []json.RawMessage, prefix []json.RawMessage) bool { + if len(prefix) == 0 { + return true + } + if len(items) < len(prefix) { + return false + } + for idx := range prefix { + previousNormalized := normalizeOpenAIWSJSONForCompareOrRaw(prefix[idx]) + currentNormalized := normalizeOpenAIWSJSONForCompareOrRaw(items[idx]) + if !bytes.Equal(previousNormalized, currentNormalized) { + return false + } + } + return true +} + +func openAIWSRawItemsHasFunctionCallOutput(items []json.RawMessage) bool { + for _, item := range items { + if isCodexToolCallOutputItemType(gjson.GetBytes(item, "type").String()) { + return true + } + } + return false +} + +func openAIWSRawItemsHaveToolCallContextForOutputs(items []json.RawMessage) bool { + if len(items) == 0 { + return false + } + contextCallIDs := make(map[string]struct{}) + outputCallIDs := make(map[string]struct{}) + for _, item := range items { + itemType := gjson.GetBytes(item, "type").String() + callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String()) + switch { + case isCodexToolCallContextItemType(itemType): + if callID != "" { + contextCallIDs[callID] = struct{}{} + } + case isCodexToolCallOutputItemType(itemType): + if callID == "" { + return false + } + outputCallIDs[callID] = struct{}{} + } + } + if len(outputCallIDs) == 0 || len(contextCallIDs) == 0 { + return false + } + for callID := range outputCallIDs { + if _, ok := contextCallIDs[callID]; !ok { + return false + } + } + return true +} + +func openAIWSRawPayloadHasToolCallOutput(payload []byte) bool { + if len(payload) == 0 { + return false + } + input := gjson.GetBytes(payload, "input") + if !input.Exists() { + return false + } + if input.IsArray() { + for _, item := range input.Array() { + if isCodexToolCallOutputItemType(item.Get("type").String()) { + return true + } + } + return false + } + if input.Type == gjson.JSON { + return isCodexToolCallOutputItemType(input.Get("type").String()) + } + return false +} + +func buildOpenAIWSReplayInputSequence( + previousFullInput []json.RawMessage, + previousFullInputExists bool, + currentPayload []byte, + hasPreviousResponseID bool, +) ([]json.RawMessage, bool, error) { + currentItems, currentExists, currentErr := openAIWSExtractNormalizedInputSequence(currentPayload) + if currentErr != nil { + return nil, false, currentErr + } + if !hasPreviousResponseID { + return cloneOpenAIWSRawMessages(currentItems), currentExists, nil + } + if !previousFullInputExists { + return cloneOpenAIWSRawMessages(currentItems), currentExists, nil + } + if !currentExists || len(currentItems) == 0 { + return cloneOpenAIWSRawMessages(previousFullInput), true, nil + } + if openAIWSRawItemsHasPrefix(currentItems, previousFullInput) { + return cloneOpenAIWSRawMessages(currentItems), true, nil + } + merged := make([]json.RawMessage, 0, len(previousFullInput)+len(currentItems)) + merged = append(merged, cloneOpenAIWSRawMessages(previousFullInput)...) + merged = append(merged, cloneOpenAIWSRawMessages(currentItems)...) + return merged, true, nil +} + +func setOpenAIWSPayloadInputSequence( + payload []byte, + fullInput []json.RawMessage, + fullInputExists bool, +) ([]byte, error) { + if !fullInputExists { + return payload, nil + } + // Preserve [] vs null semantics when input exists but is empty. + inputForMarshal := fullInput + if inputForMarshal == nil { + inputForMarshal = []json.RawMessage{} + } + inputRaw, marshalErr := json.Marshal(inputForMarshal) + if marshalErr != nil { + return nil, marshalErr + } + return sjson.SetRawBytes(payload, "input", inputRaw) +} + +func shouldKeepIngressPreviousResponseID( + previousPayload []byte, + currentPayload []byte, + lastTurnResponseID string, + hasFunctionCallOutput bool, +) (bool, string, error) { + if hasFunctionCallOutput { + return true, "has_function_call_output", nil + } + currentPreviousResponseID := strings.TrimSpace(openAIWSPayloadStringFromRaw(currentPayload, "previous_response_id")) + if currentPreviousResponseID == "" { + return false, "missing_previous_response_id", nil + } + expectedPreviousResponseID := strings.TrimSpace(lastTurnResponseID) + if expectedPreviousResponseID == "" { + return false, "missing_last_turn_response_id", nil + } + if currentPreviousResponseID != expectedPreviousResponseID { + return false, "previous_response_id_mismatch", nil + } + if len(previousPayload) == 0 { + return false, "missing_previous_turn_payload", nil + } + + previousComparable, previousComparableErr := normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(previousPayload) + if previousComparableErr != nil { + return false, "non_input_compare_error", previousComparableErr + } + currentComparable, currentComparableErr := normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(currentPayload) + if currentComparableErr != nil { + return false, "non_input_compare_error", currentComparableErr + } + if !bytes.Equal(previousComparable, currentComparable) { + return false, "non_input_changed", nil + } + return true, "strict_incremental_ok", nil +} + +type openAIWSIngressPreviousTurnStrictState struct { + nonInputComparable []byte +} + +func buildOpenAIWSIngressPreviousTurnStrictState(payload []byte) (*openAIWSIngressPreviousTurnStrictState, error) { + if len(payload) == 0 { + return nil, nil + } + nonInputComparable, nonInputErr := normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(payload) + if nonInputErr != nil { + return nil, nonInputErr + } + return &openAIWSIngressPreviousTurnStrictState{ + nonInputComparable: nonInputComparable, + }, nil +} + +func shouldKeepIngressPreviousResponseIDWithStrictState( + previousState *openAIWSIngressPreviousTurnStrictState, + currentPayload []byte, + lastTurnResponseID string, + hasFunctionCallOutput bool, +) (bool, string, error) { + if hasFunctionCallOutput { + return true, "has_function_call_output", nil + } + currentPreviousResponseID := strings.TrimSpace(openAIWSPayloadStringFromRaw(currentPayload, "previous_response_id")) + if currentPreviousResponseID == "" { + return false, "missing_previous_response_id", nil + } + expectedPreviousResponseID := strings.TrimSpace(lastTurnResponseID) + if expectedPreviousResponseID == "" { + return false, "missing_last_turn_response_id", nil + } + if currentPreviousResponseID != expectedPreviousResponseID { + return false, "previous_response_id_mismatch", nil + } + if previousState == nil { + return false, "missing_previous_turn_payload", nil + } + + currentComparable, currentComparableErr := normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(currentPayload) + if currentComparableErr != nil { + return false, "non_input_compare_error", currentComparableErr + } + if !bytes.Equal(previousState.nonInputComparable, currentComparable) { + return false, "non_input_changed", nil + } + return true, "strict_incremental_ok", nil +} diff --git a/backend/internal/service/openai_ws_forwarder_support.go b/backend/internal/service/openai_ws_forwarder_support.go new file mode 100644 index 0000000000..a2fc28d9e7 --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_support.go @@ -0,0 +1,659 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func (s *OpenAIGatewayService) isOpenAIWSGeneratePrewarmEnabled() bool { + return s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.PrewarmGenerateEnabled +} + +// performOpenAIWSGeneratePrewarm 在 WSv2 下执行可选的 generate=false 预热。 +// 预热默认关闭，仅在配置开启后生效；失败时按可恢复错误回退到 HTTP。 +func (s *OpenAIGatewayService) performOpenAIWSGeneratePrewarm( + ctx context.Context, + lease *openAIWSConnLease, + decision OpenAIWSProtocolDecision, + payload map[string]any, + previousResponseID string, + reqBody map[string]any, + account *Account, + stateStore OpenAIWSStateStore, + groupID int64, +) error { + if s == nil { + return nil + } + if lease == nil || account == nil { + logOpenAIWSModeInfo("prewarm_skip reason=invalid_state has_lease=%v has_account=%v", lease != nil, account != nil) + return nil + } + connID := strings.TrimSpace(lease.ConnID()) + if !s.isOpenAIWSGeneratePrewarmEnabled() { + return nil + } + if decision.Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { + logOpenAIWSModeInfo( + "prewarm_skip account_id=%d conn_id=%s reason=transport_not_v2 transport=%s", + account.ID, + connID, + normalizeOpenAIWSLogValue(string(decision.Transport)), + ) + return nil + } + if strings.TrimSpace(previousResponseID) != "" { + logOpenAIWSModeInfo( + "prewarm_skip account_id=%d conn_id=%s reason=has_previous_response_id previous_response_id=%s", + account.ID, + connID, + truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen), + ) + return nil + } + if lease.IsPrewarmed() { + logOpenAIWSModeInfo("prewarm_skip account_id=%d conn_id=%s reason=already_prewarmed", account.ID, connID) + return nil + } + if NeedsToolContinuation(reqBody) { + logOpenAIWSModeInfo("prewarm_skip account_id=%d conn_id=%s reason=tool_continuation", account.ID, connID) + return nil + } + prewarmStart := time.Now() + logOpenAIWSModeInfo("prewarm_start account_id=%d conn_id=%s", account.ID, connID) + + prewarmPayload := make(map[string]any, len(payload)+1) + for k, v := range payload { + prewarmPayload[k] = v + } + prewarmPayload["generate"] = false + prewarmPayloadJSON := payloadAsJSONBytes(prewarmPayload) + + if err := lease.WriteJSONWithContextTimeout(ctx, prewarmPayload, s.openAIWSWriteTimeout()); err != nil { + lease.MarkBroken() + logOpenAIWSModeInfo( + "prewarm_write_fail account_id=%d conn_id=%s cause=%s", + account.ID, + connID, + truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), + ) + return wrapOpenAIWSFallback("prewarm_write", err) + } + logOpenAIWSModeInfo("prewarm_write_sent account_id=%d conn_id=%s payload_bytes=%d", account.ID, connID, len(prewarmPayloadJSON)) + + prewarmResponseID := "" + prewarmEventCount := 0 + prewarmTerminalCount := 0 + for { + message, readErr := lease.ReadMessageWithContextTimeout(ctx, s.openAIWSReadTimeout()) + if readErr != nil { + lease.MarkBroken() + closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) + logOpenAIWSModeInfo( + "prewarm_read_fail account_id=%d conn_id=%s close_status=%s close_reason=%s cause=%s events=%d", + account.ID, + connID, + closeStatus, + closeReason, + truncateOpenAIWSLogValue(readErr.Error(), openAIWSLogValueMaxLen), + prewarmEventCount, + ) + return wrapOpenAIWSFallback("prewarm_"+classifyOpenAIWSReadFallbackReason(readErr), readErr) + } + + eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(message) + if eventType == "" { + continue + } + prewarmEventCount++ + if prewarmResponseID == "" && eventResponseID != "" { + prewarmResponseID = eventResponseID + } + if prewarmEventCount <= openAIWSPrewarmEventLogHead || eventType == "error" || isOpenAIWSTerminalEvent(eventType) { + logOpenAIWSModeInfo( + "prewarm_event account_id=%d conn_id=%s idx=%d type=%s bytes=%d", + account.ID, + connID, + prewarmEventCount, + truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen), + len(message), + ) + } + + if eventType == "error" { + errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message) + s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw) + errMsg := strings.TrimSpace(errMsgRaw) + if errMsg == "" { + errMsg = "OpenAI websocket prewarm error" + } + fallbackReason, canFallback := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) + errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) + logOpenAIWSModeInfo( + "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", + account.ID, + connID, + prewarmEventCount, + truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen), + canFallback, + errCode, + errType, + errMessage, + ) + lease.MarkBroken() + if canFallback { + return wrapOpenAIWSFallback("prewarm_"+fallbackReason, errors.New(errMsg)) + } + return wrapOpenAIWSFallback("prewarm_error_event", errors.New(errMsg)) + } + + if isOpenAIWSTerminalEvent(eventType) { + prewarmTerminalCount++ + break + } + } + + lease.MarkPrewarmed() + if prewarmResponseID != "" && stateStore != nil { + ttl := s.openAIWSResponseStickyTTL() + logOpenAIWSBindResponseAccountWarn(groupID, account.ID, prewarmResponseID, stateStore.BindResponseAccount(ctx, groupID, prewarmResponseID, account.ID, ttl)) + stateStore.BindResponseConn(prewarmResponseID, lease.ConnID(), ttl) + } + logOpenAIWSModeInfo( + "prewarm_done account_id=%d conn_id=%s response_id=%s events=%d terminal_events=%d duration_ms=%d", + account.ID, + connID, + truncateOpenAIWSLogValue(prewarmResponseID, openAIWSIDValueMaxLen), + prewarmEventCount, + prewarmTerminalCount, + time.Since(prewarmStart).Milliseconds(), + ) + return nil +} + +func payloadAsJSON(payload map[string]any) string { + return string(payloadAsJSONBytes(payload)) +} + +func payloadAsJSONBytes(payload map[string]any) []byte { + if len(payload) == 0 { + return []byte("{}") + } + body, err := json.Marshal(payload) + if err != nil { + return []byte("{}") + } + return body +} + +func isOpenAIWSTerminalEvent(eventType string) bool { + switch strings.TrimSpace(eventType) { + case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled": + return true + default: + return false + } +} + +func isOpenAIWSTokenEvent(eventType string) bool { + eventType = strings.TrimSpace(eventType) + if eventType == "" { + return false + } + switch eventType { + case "response.created", "response.in_progress", "response.output_item.added", "response.output_item.done": + return false + } + if strings.Contains(eventType, ".delta") { + return true + } + if strings.HasPrefix(eventType, "response.output_text") { + return true + } + if strings.HasPrefix(eventType, "response.output") { + return true + } + // 终止事件（response.completed/done/failed/...）由 isOpenAIWSTerminalEvent 单独处理。 + // 不能把它们当作 token event，否则当上游没有可识别的 delta 时， + // firstTokenMs 会被填到终止时刻，等于把"总耗时"误报为"首 token 延迟"。 + return false +} + +func replaceOpenAIWSMessageModel(message []byte, fromModel, toModel string) []byte { + if len(message) == 0 { + return message + } + if strings.TrimSpace(fromModel) == "" || strings.TrimSpace(toModel) == "" || fromModel == toModel { + return message + } + if !bytes.Contains(message, []byte(`"model"`)) || !bytes.Contains(message, []byte(fromModel)) { + return message + } + modelValues := gjson.GetManyBytes(message, "model", "response.model") + replaceModel := modelValues[0].Exists() && modelValues[0].Str == fromModel + replaceResponseModel := modelValues[1].Exists() && modelValues[1].Str == fromModel + if !replaceModel && !replaceResponseModel { + return message + } + updated := message + if replaceModel { + if next, err := sjson.SetBytes(updated, "model", toModel); err == nil { + updated = next + } + } + if replaceResponseModel { + if next, err := sjson.SetBytes(updated, "response.model", toModel); err == nil { + updated = next + } + } + return updated +} + +func populateOpenAIUsageFromResponseJSON(body []byte, usage *OpenAIUsage) { + if usage == nil || len(body) == 0 { + return + } + values := gjson.GetManyBytes( + body, + "usage.input_tokens", + "usage.output_tokens", + "usage.input_tokens_details.cached_tokens", + ) + usage.InputTokens = int(values[0].Int()) + usage.OutputTokens = int(values[1].Int()) + usage.CacheReadInputTokens = int(values[2].Int()) +} + +func getOpenAIGroupIDFromContext(c *gin.Context) int64 { + if c == nil { + return 0 + } + value, exists := c.Get("api_key") + if !exists { + return 0 + } + apiKey, ok := value.(*APIKey) + if !ok || apiKey == nil || apiKey.GroupID == nil { + return 0 + } + return *apiKey.GroupID +} + +// SelectAccountByPreviousResponseID 按 previous_response_id 命中账号粘连。 +// 未命中或账号不可用时返回 (nil, nil)，由调用方继续走常规调度。 +func (s *OpenAIGatewayService) SelectAccountByPreviousResponseID( + ctx context.Context, + groupID *int64, + previousResponseID string, + requestedModel string, + excludedIDs map[int64]struct{}, + requireCompact bool, +) (*AccountSelectionResult, error) { + return s.selectAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, "", requireCompact) +} + +func (s *OpenAIGatewayService) selectAccountByPreviousResponseIDForCapability( + ctx context.Context, + groupID *int64, + previousResponseID string, + requestedModel string, + excludedIDs map[int64]struct{}, + requiredCapability OpenAIEndpointCapability, + requireCompact bool, +) (*AccountSelectionResult, error) { + if s == nil { + return nil, nil + } + accountID, account, responseID, store := s.resolveAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact) + if accountID <= 0 || account == nil || store == nil { + return nil, nil + } + + result, acquireErr := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency) + if acquireErr == nil && result.Acquired { + logOpenAIWSBindResponseAccountWarn( + derefGroupID(groupID), + accountID, + responseID, + store.BindResponseAccount(ctx, derefGroupID(groupID), responseID, accountID, s.openAIWSResponseStickyTTL()), + ) + return &AccountSelectionResult{ + Account: account, + Acquired: true, + ReleaseFunc: result.ReleaseFunc, + }, nil + } + + cfg := s.schedulingConfig() + if s.concurrencyService != nil { + return &AccountSelectionResult{ + Account: account, + WaitPlan: &AccountWaitPlan{ + AccountID: accountID, + MaxConcurrency: account.Concurrency, + Timeout: cfg.StickySessionWaitTimeout, + MaxWaiting: cfg.StickySessionMaxWaiting, + }, + }, nil + } + return nil, nil +} + +func (s *OpenAIGatewayService) ResolveAccountIDByPreviousResponseIDForScheduler( + ctx context.Context, + groupID *int64, + previousResponseID string, + requestedModel string, + excludedIDs map[int64]struct{}, + requiredCapability OpenAIEndpointCapability, + requireCompact bool, +) int64 { + accountID, _, _, _ := s.resolveAccountByPreviousResponseIDForCapability(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact) + return accountID +} + +func (s *OpenAIGatewayService) resolveAccountByPreviousResponseIDForCapability( + ctx context.Context, + groupID *int64, + previousResponseID string, + requestedModel string, + excludedIDs map[int64]struct{}, + requiredCapability OpenAIEndpointCapability, + requireCompact bool, +) (int64, *Account, string, OpenAIWSStateStore) { + if s == nil { + return 0, nil, "", nil + } + responseID := strings.TrimSpace(previousResponseID) + if responseID == "" { + return 0, nil, "", nil + } + store := s.getOpenAIWSStateStore() + if store == nil { + return 0, nil, "", nil + } + + accountID, err := store.GetResponseAccount(ctx, derefGroupID(groupID), responseID) + if err != nil || accountID <= 0 { + return 0, nil, "", nil + } + if excludedIDs != nil { + if _, excluded := excludedIDs[accountID]; excluded { + return 0, nil, "", nil + } + } + + account, err := s.getSchedulableAccount(ctx, accountID) + if err != nil || account == nil { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + // 非 WSv2 场景（如 force_http/全局关闭）不应使用 previous_response_id 粘连， + // 以保持“回滚到 HTTP”后的历史行为一致性。 + if s.getOpenAIWSProtocolResolver().Resolve(account).Transport != OpenAIUpstreamTransportResponsesWebsocketV2 { + return 0, nil, "", nil + } + if shouldClearStickySession(account, requestedModel) || !account.IsOpenAI() || !account.IsSchedulable() { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if requestedModel != "" && !account.IsModelSupported(requestedModel) { + return 0, nil, "", nil + } + if !account.SupportsOpenAIEndpointCapability(requiredCapability) { + return 0, nil, "", nil + } + // Quota auto-pause must also gate the previous_response_id sticky path; otherwise an + // account over its 5h/7d threshold keeps serving the same response chain even though + // normal scheduling skips it. Pause is transient, so fall through to normal scheduling + // without deleting the binding (the window may reset before the next turn). + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused { + return 0, nil, "", nil + } + if s.schedulerSnapshot != nil && s.accountRepo != nil { + latest, latestErr := s.accountRepo.GetByID(ctx, account.ID) + if latestErr != nil || latest == nil { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if shouldClearStickySession(latest, requestedModel) || !latest.IsOpenAI() || !latest.IsSchedulable() { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + if requestedModel != "" && !latest.IsModelSupported(requestedModel) { + return 0, nil, "", nil + } + if !latest.SupportsOpenAIEndpointCapability(requiredCapability) { + return 0, nil, "", nil + } + if paused, _ := shouldAutoPauseOpenAIAccountByQuota(ctx, latest); paused { + return 0, nil, "", nil + } + if s.isOpenAIAccountRuntimeBlocked(latest) { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + account = latest + } + if requireCompact && openAICompactSupportTier(account) == 0 { + _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) + return 0, nil, "", nil + } + return accountID, account, responseID, store +} + +func classifyOpenAIWSAcquireError(err error) string { + if err == nil { + return "acquire_conn" + } + var dialErr *openAIWSDialError + if errors.As(err, &dialErr) { + switch dialErr.StatusCode { + case 426: + return "upgrade_required" + case 401, 403: + return "auth_failed" + case 429: + return "upstream_rate_limited" + } + if dialErr.StatusCode >= 500 { + return "upstream_5xx" + } + return "dial_failed" + } + if errors.Is(err, errOpenAIWSConnQueueFull) { + return "conn_queue_full" + } + if errors.Is(err, errOpenAIWSPreferredConnUnavailable) { + return "preferred_conn_unavailable" + } + if errors.Is(err, context.DeadlineExceeded) { + return "acquire_timeout" + } + return "acquire_conn" +} + +func isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw string) bool { + code := strings.ToLower(strings.TrimSpace(codeRaw)) + errType := strings.ToLower(strings.TrimSpace(errTypeRaw)) + msg := strings.ToLower(strings.TrimSpace(msgRaw)) + + if strings.Contains(errType, "rate_limit") || strings.Contains(errType, "usage_limit") { + return true + } + if strings.Contains(code, "rate_limit") || strings.Contains(code, "usage_limit") || strings.Contains(code, "insufficient_quota") { + return true + } + if strings.Contains(msg, "usage limit") && strings.Contains(msg, "reached") { + return true + } + if strings.Contains(msg, "rate limit") && (strings.Contains(msg, "reached") || strings.Contains(msg, "exceeded")) { + return true + } + return false +} + +func (s *OpenAIGatewayService) persistOpenAIWSRateLimitSignal(ctx context.Context, account *Account, headers http.Header, responseBody []byte, codeRaw, errTypeRaw, msgRaw string) { + if s == nil || s.rateLimitService == nil || account == nil || account.Platform != PlatformOpenAI { + return + } + if !isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw) { + return + } + s.handleOpenAIAccountUpstreamError(ctx, account, http.StatusTooManyRequests, headers, responseBody) +} + +func classifyOpenAIWSErrorEventFromRaw(codeRaw, errTypeRaw, msgRaw string) (string, bool) { + code := strings.ToLower(strings.TrimSpace(codeRaw)) + errType := strings.ToLower(strings.TrimSpace(errTypeRaw)) + msg := strings.ToLower(strings.TrimSpace(msgRaw)) + + switch code { + case "upgrade_required": + return "upgrade_required", true + case "websocket_not_supported", "websocket_unsupported": + return "ws_unsupported", true + case "websocket_connection_limit_reached": + return "ws_connection_limit_reached", true + case "invalid_encrypted_content": + return "invalid_encrypted_content", true + case "previous_response_not_found": + return "previous_response_not_found", true + } + if isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw) { + return "upstream_rate_limited", false + } + if strings.Contains(msg, "upgrade required") || strings.Contains(msg, "status 426") { + return "upgrade_required", true + } + if strings.Contains(errType, "upgrade") { + return "upgrade_required", true + } + if strings.Contains(msg, "websocket") && strings.Contains(msg, "unsupported") { + return "ws_unsupported", true + } + if strings.Contains(msg, "connection limit") && strings.Contains(msg, "websocket") { + return "ws_connection_limit_reached", true + } + if strings.Contains(msg, "invalid_encrypted_content") || + (strings.Contains(msg, "encrypted content") && strings.Contains(msg, "could not be verified")) { + return "invalid_encrypted_content", true + } + if strings.Contains(msg, "previous_response_not_found") || + (strings.Contains(msg, "previous response") && strings.Contains(msg, "not found")) { + return "previous_response_not_found", true + } + if strings.Contains(errType, "server_error") || strings.Contains(code, "server_error") { + return "upstream_error_event", true + } + return "event_error", false +} + +func classifyOpenAIWSErrorEvent(message []byte) (string, bool) { + if len(message) == 0 { + return "event_error", false + } + return classifyOpenAIWSErrorEventFromRaw(parseOpenAIWSErrorEventFields(message)) +} + +func openAIWSErrorHTTPStatusFromRaw(codeRaw, errTypeRaw string) int { + code := strings.ToLower(strings.TrimSpace(codeRaw)) + errType := strings.ToLower(strings.TrimSpace(errTypeRaw)) + switch { + case strings.Contains(errType, "invalid_request"), + strings.Contains(code, "invalid_request"), + strings.Contains(code, "bad_request"), + code == "invalid_encrypted_content", + code == "previous_response_not_found": + return http.StatusBadRequest + case strings.Contains(errType, "authentication"), + strings.Contains(code, "invalid_api_key"), + strings.Contains(code, "unauthorized"): + return http.StatusUnauthorized + case strings.Contains(errType, "permission"), + strings.Contains(code, "forbidden"): + return http.StatusForbidden + case isOpenAIWSRateLimitError(codeRaw, errTypeRaw, ""): + return http.StatusTooManyRequests + default: + return http.StatusBadGateway + } +} + +func openAIWSErrorHTTPStatus(message []byte) int { + if len(message) == 0 { + return http.StatusBadGateway + } + codeRaw, errTypeRaw, _ := parseOpenAIWSErrorEventFields(message) + return openAIWSErrorHTTPStatusFromRaw(codeRaw, errTypeRaw) +} + +func (s *OpenAIGatewayService) openAIWSFallbackCooldown() time.Duration { + if s == nil || s.cfg == nil { + return 30 * time.Second + } + seconds := s.cfg.Gateway.OpenAIWS.FallbackCooldownSeconds + if seconds <= 0 { + return 0 + } + return time.Duration(seconds) * time.Second +} + +func (s *OpenAIGatewayService) isOpenAIWSFallbackCooling(accountID int64) bool { + if s == nil || accountID <= 0 { + return false + } + cooldown := s.openAIWSFallbackCooldown() + if cooldown <= 0 { + return false + } + rawUntil, ok := s.openaiWSFallbackUntil.Load(accountID) + if !ok || rawUntil == nil { + return false + } + until, ok := rawUntil.(time.Time) + if !ok || until.IsZero() { + s.openaiWSFallbackUntil.Delete(accountID) + return false + } + if time.Now().Before(until) { + return true + } + s.openaiWSFallbackUntil.Delete(accountID) + return false +} + +func (s *OpenAIGatewayService) markOpenAIWSFallbackCooling(accountID int64, _ string) { + if s == nil || accountID <= 0 { + return + } + cooldown := s.openAIWSFallbackCooldown() + if cooldown <= 0 { + return + } + s.openaiWSFallbackUntil.Store(accountID, time.Now().Add(cooldown)) +} + +func (s *OpenAIGatewayService) clearOpenAIWSFallbackCooling(accountID int64) { + if s == nil || accountID <= 0 { + return + } + s.openaiWSFallbackUntil.Delete(accountID) +} diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go new file mode 100644 index 0000000000..ea8c0a0722 --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -0,0 +1,732 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/url" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +func (s *OpenAIGatewayService) forwardOpenAIWSV2( + ctx context.Context, + c *gin.Context, + account *Account, + reqBody map[string]any, + token string, + decision OpenAIWSProtocolDecision, + isCodexCLI bool, + reqStream bool, + originalModel string, + mappedModel string, + startTime time.Time, + attempt int, + lastFailureReason string, +) (*OpenAIForwardResult, error) { + if s == nil || account == nil { + return nil, wrapOpenAIWSFallback("invalid_state", errors.New("service or account is nil")) + } + + wsURL, err := s.buildOpenAIResponsesWSURL(account) + if err != nil { + return nil, wrapOpenAIWSFallback("build_ws_url", err) + } + wsHost := "-" + wsPath := "-" + if parsed, parseErr := url.Parse(wsURL); parseErr == nil && parsed != nil { + if h := strings.TrimSpace(parsed.Host); h != "" { + wsHost = normalizeOpenAIWSLogValue(h) + } + if p := strings.TrimSpace(parsed.Path); p != "" { + wsPath = normalizeOpenAIWSLogValue(p) + } + } + logOpenAIWSModeDebug( + "dial_target account_id=%d account_type=%s ws_host=%s ws_path=%s", + account.ID, + account.Type, + wsHost, + wsPath, + ) + + payload := s.buildOpenAIWSCreatePayload(reqBody, account) + payloadStrategy, removedKeys := applyOpenAIWSRetryPayloadStrategy(payload, attempt) + previousResponseID := openAIWSPayloadString(payload, "previous_response_id") + previousResponseIDKind := ClassifyOpenAIPreviousResponseIDKind(previousResponseID) + promptCacheKey := openAIWSPayloadString(payload, "prompt_cache_key") + _, hasTools := payload["tools"] + debugEnabled := isOpenAIWSModeDebugEnabled() + payloadBytes := -1 + resolvePayloadBytes := func() int { + if payloadBytes >= 0 { + return payloadBytes + } + payloadBytes = len(payloadAsJSONBytes(payload)) + return payloadBytes + } + streamValue := "-" + if raw, ok := payload["stream"]; ok { + streamValue = normalizeOpenAIWSLogValue(strings.TrimSpace(fmt.Sprintf("%v", raw))) + } + turnState := "" + turnMetadata := "" + if c != nil && c.Request != nil { + turnState = strings.TrimSpace(c.GetHeader(openAIWSTurnStateHeader)) + turnMetadata = strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)) + } + setOpenAIWSTurnMetadata(payload, turnMetadata) + payloadEventType := openAIWSPayloadString(payload, "type") + if payloadEventType == "" { + payloadEventType = "response.create" + } + if s.shouldEmitOpenAIWSPayloadSchema(attempt) { + logOpenAIWSModeInfo( + "[debug] payload_schema account_id=%d attempt=%d event=%s payload_keys=%s payload_bytes=%d payload_key_sizes=%s input_summary=%s stream=%s payload_strategy=%s removed_keys=%s has_previous_response_id=%v has_prompt_cache_key=%v has_tools=%v", + account.ID, + attempt, + payloadEventType, + normalizeOpenAIWSLogValue(strings.Join(sortedKeys(payload), ",")), + resolvePayloadBytes(), + normalizeOpenAIWSLogValue(summarizeOpenAIWSPayloadKeySizes(payload, openAIWSPayloadKeySizeTopN)), + normalizeOpenAIWSLogValue(summarizeOpenAIWSInput(payload["input"])), + streamValue, + normalizeOpenAIWSLogValue(payloadStrategy), + normalizeOpenAIWSLogValue(strings.Join(removedKeys, ",")), + previousResponseID != "", + promptCacheKey != "", + hasTools, + ) + } + + stateStore := s.getOpenAIWSStateStore() + groupID := getOpenAIGroupIDFromContext(c) + sessionHash := s.GenerateSessionHash(c, nil) + if sessionHash == "" { + var legacySessionHash string + sessionHash, legacySessionHash = openAIWSSessionHashesFromID(promptCacheKey) + attachOpenAILegacySessionHashToGin(c, legacySessionHash) + } + if turnState == "" && stateStore != nil && sessionHash != "" { + if savedTurnState, ok := stateStore.GetSessionTurnState(groupID, sessionHash); ok { + turnState = savedTurnState + } + } + preferredConnID := "" + if stateStore != nil && previousResponseID != "" { + if connID, ok := stateStore.GetResponseConn(previousResponseID); ok { + preferredConnID = connID + } + } + storeDisabled := s.isOpenAIWSStoreDisabledInRequest(reqBody, account) + if stateStore != nil && storeDisabled && previousResponseID == "" && sessionHash != "" { + if connID, ok := stateStore.GetSessionConn(groupID, sessionHash); ok { + preferredConnID = connID + } + } + storeDisabledConnMode := s.openAIWSStoreDisabledConnMode() + forceNewConnByPolicy := shouldForceNewConnOnStoreDisabled(storeDisabledConnMode, lastFailureReason) + forceNewConn := forceNewConnByPolicy && storeDisabled && previousResponseID == "" && sessionHash != "" && preferredConnID == "" + wsHeaders, sessionResolution, buildHdrErr := s.buildOpenAIWSHeaders(ctx, c, account, token, decision, isCodexCLI, turnState, turnMetadata, promptCacheKey) + if buildHdrErr != nil { + return nil, fmt.Errorf("build ws headers: %w", buildHdrErr) + } + logOpenAIWSModeDebug( + "acquire_start account_id=%d account_type=%s transport=%s preferred_conn_id=%s has_previous_response_id=%v session_hash=%s has_turn_state=%v turn_state_len=%d has_turn_metadata=%v turn_metadata_len=%d store_disabled=%v store_disabled_conn_mode=%s retry_last_reason=%s force_new_conn=%v header_user_agent=%s header_openai_beta=%s header_originator=%s header_accept_language=%s header_session_id=%s header_conversation_id=%s session_id_source=%s conversation_id_source=%s has_prompt_cache_key=%v has_chatgpt_account_id=%v has_authorization=%v has_session_id=%v has_conversation_id=%v proxy_enabled=%v", + account.ID, + account.Type, + normalizeOpenAIWSLogValue(string(decision.Transport)), + truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), + previousResponseID != "", + truncateOpenAIWSLogValue(sessionHash, 12), + turnState != "", + len(turnState), + turnMetadata != "", + len(turnMetadata), + storeDisabled, + normalizeOpenAIWSLogValue(storeDisabledConnMode), + truncateOpenAIWSLogValue(lastFailureReason, openAIWSLogValueMaxLen), + forceNewConn, + openAIWSHeaderValueForLog(wsHeaders, "user-agent"), + openAIWSHeaderValueForLog(wsHeaders, "openai-beta"), + openAIWSHeaderValueForLog(wsHeaders, "originator"), + openAIWSHeaderValueForLog(wsHeaders, "accept-language"), + openAIWSHeaderValueForLog(wsHeaders, "session_id"), + openAIWSHeaderValueForLog(wsHeaders, "conversation_id"), + normalizeOpenAIWSLogValue(sessionResolution.SessionSource), + normalizeOpenAIWSLogValue(sessionResolution.ConversationSource), + promptCacheKey != "", + hasOpenAIWSHeader(wsHeaders, "chatgpt-account-id"), + hasOpenAIWSHeader(wsHeaders, "authorization"), + hasOpenAIWSHeader(wsHeaders, "session_id"), + hasOpenAIWSHeader(wsHeaders, "conversation_id"), + account.ProxyID != nil && account.Proxy != nil, + ) + + acquireCtx, acquireCancel := context.WithTimeout(ctx, s.openAIWSAcquireTimeout()) + defer acquireCancel() + + lease, err := s.getOpenAIWSConnPool().Acquire(acquireCtx, openAIWSAcquireRequest{ + Account: account, + WSURL: wsURL, + Headers: wsHeaders, + PreferredConnID: preferredConnID, + ForceNewConn: forceNewConn, + ProxyURL: func() string { + if account.ProxyID != nil && account.Proxy != nil { + return account.Proxy.URL() + } + return "" + }(), + }) + if err != nil { + dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(err) + logOpenAIWSModeInfo( + "acquire_fail account_id=%d account_type=%s transport=%s reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_new_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v", + account.ID, + account.Type, + normalizeOpenAIWSLogValue(string(decision.Transport)), + normalizeOpenAIWSLogValue(classifyOpenAIWSAcquireError(err)), + dialStatus, + dialClass, + dialCloseStatus, + truncateOpenAIWSLogValue(dialCloseReason, openAIWSHeaderValueMaxLen), + dialRespServer, + dialRespVia, + dialRespCFRay, + dialRespReqID, + truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), + forceNewConn, + wsHost, + wsPath, + account.ProxyID != nil && account.Proxy != nil, + ) + var dialErr *openAIWSDialError + if errors.As(err, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusTooManyRequests { + s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error())) + } + return nil, wrapOpenAIWSFallback(classifyOpenAIWSAcquireError(err), err) + } + // cleanExit 标记正常终端事件退出，此时上游不会再发送帧，连接可安全归还复用。 + // 所有异常路径（读写错误、error 事件等）已在各自分支中提前调用 MarkBroken， + // 因此 defer 中只需处理正常退出时不 MarkBroken 即可。 + cleanExit := false + defer func() { + if !cleanExit { + lease.MarkBroken() + } + lease.Release() + }() + connID := strings.TrimSpace(lease.ConnID()) + logOpenAIWSModeDebug( + "connected account_id=%d account_type=%s transport=%s conn_id=%s conn_reused=%v conn_pick_ms=%d queue_wait_ms=%d has_previous_response_id=%v", + account.ID, + account.Type, + normalizeOpenAIWSLogValue(string(decision.Transport)), + connID, + lease.Reused(), + lease.ConnPickDuration().Milliseconds(), + lease.QueueWaitDuration().Milliseconds(), + previousResponseID != "", + ) + if previousResponseID != "" { + logOpenAIWSModeInfo( + "continuation_probe account_id=%d account_type=%s conn_id=%s previous_response_id=%s previous_response_id_kind=%s preferred_conn_id=%s conn_reused=%v store_disabled=%v session_hash=%s header_session_id=%s header_conversation_id=%s session_id_source=%s conversation_id_source=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v", + account.ID, + account.Type, + truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(previousResponseIDKind), + truncateOpenAIWSLogValue(preferredConnID, openAIWSIDValueMaxLen), + lease.Reused(), + storeDisabled, + truncateOpenAIWSLogValue(sessionHash, 12), + openAIWSHeaderValueForLog(wsHeaders, "session_id"), + openAIWSHeaderValueForLog(wsHeaders, "conversation_id"), + normalizeOpenAIWSLogValue(sessionResolution.SessionSource), + normalizeOpenAIWSLogValue(sessionResolution.ConversationSource), + turnState != "", + len(turnState), + promptCacheKey != "", + ) + } + if c != nil { + SetOpsLatencyMs(c, OpsOpenAIWSConnPickMsKey, lease.ConnPickDuration().Milliseconds()) + SetOpsLatencyMs(c, OpsOpenAIWSQueueWaitMsKey, lease.QueueWaitDuration().Milliseconds()) + c.Set(OpsOpenAIWSConnReusedKey, lease.Reused()) + if connID != "" { + c.Set(OpsOpenAIWSConnIDKey, connID) + } + } + + handshakeTurnState := strings.TrimSpace(lease.HandshakeHeader(openAIWSTurnStateHeader)) + logOpenAIWSModeDebug( + "handshake account_id=%d conn_id=%s has_turn_state=%v turn_state_len=%d", + account.ID, + connID, + handshakeTurnState != "", + len(handshakeTurnState), + ) + if handshakeTurnState != "" { + if stateStore != nil && sessionHash != "" { + stateStore.BindSessionTurnState(groupID, sessionHash, handshakeTurnState, s.openAIWSSessionStickyTTL()) + } + if c != nil { + c.Header(http.CanonicalHeaderKey(openAIWSTurnStateHeader), handshakeTurnState) + } + } + + if err := s.performOpenAIWSGeneratePrewarm( + ctx, + lease, + decision, + payload, + previousResponseID, + reqBody, + account, + stateStore, + groupID, + ); err != nil { + return nil, err + } + + if err := lease.WriteJSONWithContextTimeout(ctx, payload, s.openAIWSWriteTimeout()); err != nil { + lease.MarkBroken() + logOpenAIWSModeInfo( + "write_request_fail account_id=%d conn_id=%s cause=%s payload_bytes=%d", + account.ID, + connID, + truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), + resolvePayloadBytes(), + ) + return nil, wrapOpenAIWSFallback("write_request", err) + } + if debugEnabled { + logOpenAIWSModeDebug( + "write_request_sent account_id=%d conn_id=%s stream=%v payload_bytes=%d previous_response_id=%s", + account.ID, + connID, + reqStream, + resolvePayloadBytes(), + truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen), + ) + } + + usage := &OpenAIUsage{} + imageCounter := newOpenAIImageOutputCounter() + var firstTokenMs *int + responseID := "" + var finalResponse []byte + wroteDownstream := false + needModelReplace := originalModel != mappedModel + var mappedModelBytes []byte + if needModelReplace && mappedModel != "" { + mappedModelBytes = []byte(mappedModel) + } + bufferedStreamEvents := make([][]byte, 0, 4) + eventCount := 0 + tokenEventCount := 0 + terminalEventCount := 0 + bufferedEventCount := 0 + flushedBufferedEventCount := 0 + firstEventType := "" + lastEventType := "" + + var flusher http.Flusher + if reqStream { + if s.responseHeaderFilter != nil { + responseheaders.WriteFilteredHeaders(c.Writer.Header(), http.Header{}, s.responseHeaderFilter) + } + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + f, ok := c.Writer.(http.Flusher) + if !ok { + lease.MarkBroken() + return nil, wrapOpenAIWSFallback("streaming_not_supported", errors.New("streaming not supported")) + } + flusher = f + } + + clientDisconnected := false + flushBatchSize := s.openAIWSEventFlushBatchSize() + flushInterval := s.openAIWSEventFlushInterval() + pendingFlushEvents := 0 + lastFlushAt := time.Now() + flushStreamWriter := func(force bool) { + if clientDisconnected || flusher == nil || pendingFlushEvents <= 0 { + return + } + if !force && flushBatchSize > 1 && pendingFlushEvents < flushBatchSize { + if flushInterval <= 0 || time.Since(lastFlushAt) < flushInterval { + return + } + } + flusher.Flush() + pendingFlushEvents = 0 + lastFlushAt = time.Now() + } + emitStreamMessage := func(message []byte, forceFlush bool) { + if clientDisconnected { + return + } + frame := make([]byte, 0, len(message)+8) + frame = append(frame, "data: "...) + frame = append(frame, message...) + frame = append(frame, '\n', '\n') + _, wErr := c.Writer.Write(frame) + if wErr == nil { + wroteDownstream = true + pendingFlushEvents++ + flushStreamWriter(forceFlush) + return + } + clientDisconnected = true + logger.LegacyPrintf("service.openai_gateway", "[OpenAI WS Mode] client disconnected, continue draining upstream: account=%d", account.ID) + } + flushBufferedStreamEvents := func(reason string) { + if len(bufferedStreamEvents) == 0 { + return + } + flushed := len(bufferedStreamEvents) + for _, buffered := range bufferedStreamEvents { + emitStreamMessage(buffered, false) + } + bufferedStreamEvents = bufferedStreamEvents[:0] + flushStreamWriter(true) + flushedBufferedEventCount += flushed + if debugEnabled { + logOpenAIWSModeDebug( + "buffer_flush account_id=%d conn_id=%s reason=%s flushed=%d total_flushed=%d client_disconnected=%v", + account.ID, + connID, + truncateOpenAIWSLogValue(reason, openAIWSLogValueMaxLen), + flushed, + flushedBufferedEventCount, + clientDisconnected, + ) + } + } + + readTimeout := s.openAIWSReadTimeout() + + for { + message, readErr := lease.ReadMessageWithContextTimeout(ctx, readTimeout) + if readErr != nil { + lease.MarkBroken() + closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) + logOpenAIWSModeInfo( + "read_fail account_id=%d conn_id=%s wrote_downstream=%v close_status=%s close_reason=%s cause=%s events=%d token_events=%d terminal_events=%d buffered_pending=%d buffered_flushed=%d first_event=%s last_event=%s", + account.ID, + connID, + wroteDownstream, + closeStatus, + closeReason, + truncateOpenAIWSLogValue(readErr.Error(), openAIWSLogValueMaxLen), + eventCount, + tokenEventCount, + terminalEventCount, + len(bufferedStreamEvents), + flushedBufferedEventCount, + truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen), + ) + if !wroteDownstream { + return nil, wrapOpenAIWSFallback(classifyOpenAIWSReadFallbackReason(readErr), readErr) + } + if clientDisconnected { + break + } + setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(readErr.Error()), "") + return nil, fmt.Errorf("openai ws read event: %w", readErr) + } + + eventType, eventResponseID, responseField := parseOpenAIWSEventEnvelope(message) + if eventType == "" { + continue + } + eventCount++ + if firstEventType == "" { + firstEventType = eventType + } + lastEventType = eventType + + if responseID == "" && eventResponseID != "" { + responseID = eventResponseID + } + + isTokenEvent := isOpenAIWSTokenEvent(eventType) + if isTokenEvent { + tokenEventCount++ + } + isTerminalEvent := isOpenAIWSTerminalEvent(eventType) + if isTerminalEvent { + terminalEventCount++ + } + if firstTokenMs == nil && isTokenEvent { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + if debugEnabled && shouldLogOpenAIWSEvent(eventCount, eventType) { + logOpenAIWSModeDebug( + "event_received account_id=%d conn_id=%s idx=%d type=%s bytes=%d token=%v terminal=%v buffered_pending=%d", + account.ID, + connID, + eventCount, + truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen), + len(message), + isTokenEvent, + isTerminalEvent, + len(bufferedStreamEvents), + ) + } + + if !clientDisconnected { + if needModelReplace && len(mappedModelBytes) > 0 && openAIWSEventMayContainModel(eventType) && bytes.Contains(message, mappedModelBytes) { + message = replaceOpenAIWSMessageModel(message, mappedModel, originalModel) + } + if openAIWSEventMayContainToolCalls(eventType) && openAIWSMessageLikelyContainsToolCalls(message) { + if corrected, changed := s.toolCorrector.CorrectToolCallsInSSEBytes(message); changed { + message = corrected + } + } + } + if openAIWSEventShouldParseUsage(eventType) { + parseOpenAIWSResponseUsageFromCompletedEvent(message, usage) + } + imageCounter.AddSSEData(message) + + if eventType == "response.failed" { + if hit, code, msg := detectOpenAICyberPolicy(message); hit { + MarkOpsCyberPolicy(c, CyberPolicyMark{ + Code: code, + Message: msg, + Body: truncateString(string(message), 4096), + UpstreamStatus: http.StatusOK, + UpstreamInTok: usage.InputTokens, + UpstreamOutTok: usage.OutputTokens, + }) + } + } + + if eventType == "error" { + errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message) + s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw) + errMsg := strings.TrimSpace(errMsgRaw) + if errMsg == "" { + errMsg = "Upstream websocket error" + } + fallbackReason, canFallback := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) + errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) + logOpenAIWSModeInfo( + "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", + account.ID, + connID, + eventCount, + truncateOpenAIWSLogValue(fallbackReason, openAIWSLogValueMaxLen), + canFallback, + errCode, + errType, + errMessage, + ) + if fallbackReason == "previous_response_not_found" { + logOpenAIWSModeInfo( + "previous_response_not_found_diag account_id=%d account_type=%s conn_id=%s previous_response_id=%s previous_response_id_kind=%s response_id=%s event_idx=%d req_stream=%v store_disabled=%v conn_reused=%v session_hash=%s header_session_id=%s header_conversation_id=%s session_id_source=%s conversation_id_source=%s has_turn_state=%v turn_state_len=%d has_prompt_cache_key=%v err_code=%s err_type=%s err_message=%s", + account.ID, + account.Type, + connID, + truncateOpenAIWSLogValue(previousResponseID, openAIWSIDValueMaxLen), + normalizeOpenAIWSLogValue(previousResponseIDKind), + truncateOpenAIWSLogValue(responseID, openAIWSIDValueMaxLen), + eventCount, + reqStream, + storeDisabled, + lease.Reused(), + truncateOpenAIWSLogValue(sessionHash, 12), + openAIWSHeaderValueForLog(wsHeaders, "session_id"), + openAIWSHeaderValueForLog(wsHeaders, "conversation_id"), + normalizeOpenAIWSLogValue(sessionResolution.SessionSource), + normalizeOpenAIWSLogValue(sessionResolution.ConversationSource), + turnState != "", + len(turnState), + promptCacheKey != "", + errCode, + errType, + errMessage, + ) + } + // error 事件后连接不再可复用，避免回池后污染下一请求。 + lease.MarkBroken() + if !wroteDownstream && canFallback { + return nil, wrapOpenAIWSFallback(fallbackReason, errors.New(errMsg)) + } + statusCode := openAIWSErrorHTTPStatusFromRaw(errCodeRaw, errTypeRaw) + setOpsUpstreamError(c, statusCode, errMsg, "") + if reqStream && !clientDisconnected { + flushBufferedStreamEvents("error_event") + emitStreamMessage(message, true) + } + if !reqStream { + c.JSON(statusCode, gin.H{ + "error": gin.H{ + "type": "upstream_error", + "message": errMsg, + }, + }) + } + return nil, fmt.Errorf("openai ws error event: %s", errMsg) + } + + if reqStream { + // 在首个 token 前先缓冲事件（如 response.created）， + // 以便上游早期断连时仍可安全回退到 HTTP，不给下游发送半截流。 + shouldBuffer := firstTokenMs == nil && !isTokenEvent && !isTerminalEvent + if shouldBuffer { + buffered := make([]byte, len(message)) + copy(buffered, message) + bufferedStreamEvents = append(bufferedStreamEvents, buffered) + bufferedEventCount++ + if debugEnabled && shouldLogOpenAIWSBufferedEvent(bufferedEventCount) { + logOpenAIWSModeDebug( + "buffer_enqueue account_id=%d conn_id=%s idx=%d event_idx=%d event_type=%s buffer_size=%d", + account.ID, + connID, + bufferedEventCount, + eventCount, + truncateOpenAIWSLogValue(eventType, openAIWSLogValueMaxLen), + len(bufferedStreamEvents), + ) + } + } else { + flushBufferedStreamEvents(eventType) + emitStreamMessage(message, isTerminalEvent) + } + } else { + if responseField.Exists() && responseField.Type == gjson.JSON { + finalResponse = []byte(responseField.Raw) + } + } + + if isTerminalEvent { + cleanExit = true + break + } + } + + if !reqStream { + if len(finalResponse) == 0 { + logOpenAIWSModeInfo( + "missing_final_response account_id=%d conn_id=%s events=%d token_events=%d terminal_events=%d wrote_downstream=%v", + account.ID, + connID, + eventCount, + tokenEventCount, + terminalEventCount, + wroteDownstream, + ) + if !wroteDownstream { + return nil, wrapOpenAIWSFallback("missing_final_response", errors.New("no terminal response payload")) + } + return nil, errors.New("ws finished without final response") + } + + if needModelReplace { + finalResponse = s.replaceModelInResponseBody(finalResponse, mappedModel, originalModel) + } + finalResponse = s.correctToolCallsInResponseBody(finalResponse) + populateOpenAIUsageFromResponseJSON(finalResponse, usage) + if responseID == "" { + responseID = strings.TrimSpace(gjson.GetBytes(finalResponse, "id").String()) + } + + c.Data(http.StatusOK, "application/json", finalResponse) + } else { + flushStreamWriter(true) + } + + if responseID != "" && stateStore != nil { + ttl := s.openAIWSResponseStickyTTL() + logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) + stateStore.BindResponseConn(responseID, lease.ConnID(), ttl) + } + if stateStore != nil && storeDisabled && sessionHash != "" { + stateStore.BindSessionConn(groupID, sessionHash, lease.ConnID(), s.openAIWSSessionStickyTTL()) + } + firstTokenMsValue := -1 + if firstTokenMs != nil { + firstTokenMsValue = *firstTokenMs + } + logOpenAIWSModeDebug( + "completed account_id=%d conn_id=%s response_id=%s stream=%v duration_ms=%d events=%d token_events=%d terminal_events=%d buffered_events=%d buffered_flushed=%d first_event=%s last_event=%s first_token_ms=%d wrote_downstream=%v client_disconnected=%v", + account.ID, + connID, + truncateOpenAIWSLogValue(strings.TrimSpace(responseID), openAIWSIDValueMaxLen), + reqStream, + time.Since(startTime).Milliseconds(), + eventCount, + tokenEventCount, + terminalEventCount, + bufferedEventCount, + flushedBufferedEventCount, + truncateOpenAIWSLogValue(firstEventType, openAIWSLogValueMaxLen), + truncateOpenAIWSLogValue(lastEventType, openAIWSLogValueMaxLen), + firstTokenMsValue, + wroteDownstream, + clientDisconnected, + ) + + return &OpenAIForwardResult{ + RequestID: responseID, + Usage: *usage, + Model: originalModel, + UpstreamModel: mappedModel, + ImageCount: imageCounter.Count(), + ImageOutputSizes: imageCounter.Sizes(), + ServiceTier: extractOpenAIServiceTier(reqBody), + ReasoningEffort: extractOpenAIReasoningEffort(reqBody, originalModel), + Stream: reqStream, + OpenAIWSMode: true, + ResponseHeaders: lease.HandshakeHeaders(), + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + }, nil +} + +// ProxyResponsesWebSocketFromClient 处理客户端入站 WebSocket（OpenAI Responses WS Mode）并转发到上游。 +// 当前实现按“单请求 -> 终止事件 -> 下一请求”的顺序代理，适配 Codex CLI 的 turn 模式。 +// stripCodexSparkImageGenerationToolFromRawPayload removes the image_generation +// tool from a raw /responses payload when the upstream model is gpt-5.3-codex-spark. +// Spark rejects that tool upstream with HTTP 400 (invalid_request_error, param=tools); +// Codex clients advertise it by default. Returns the (possibly unchanged) payload, +// whether it changed, and any JSON decode error. +func stripCodexSparkImageGenerationToolFromRawPayload(payload []byte, model string) ([]byte, bool, error) { + if !isCodexSparkModel(model) || !openAIRequestBodyHasImageGenerationTool(payload) { + return payload, false, nil + } + return stripOpenAIImageGenerationToolFromRawPayload(payload) +} + +func stripOpenAIImageGenerationToolFromRawPayload(payload []byte) ([]byte, bool, error) { + payloadMap := make(map[string]any) + if err := json.Unmarshal(payload, &payloadMap); err != nil { + return payload, false, err + } + if !stripOpenAIImageGenerationTools(payloadMap) { + return payload, false, nil + } + rebuilt, err := json.Marshal(payloadMap) + if err != nil { + return payload, false, err + } + return rebuilt, true, nil +}