diff --git a/common/json.go b/common/json.go index 1625be6d51f7..68bbe7ea7c2c 100644 --- a/common/json.go +++ b/common/json.go @@ -6,6 +6,8 @@ import ( "io" ) +type RawMessage = json.RawMessage + func Unmarshal(data []byte, v any) error { return json.Unmarshal(data, v) } @@ -22,7 +24,7 @@ func Marshal(v any) ([]byte, error) { return json.Marshal(v) } -func GetJsonType(data json.RawMessage) string { +func GetJsonType(data RawMessage) string { trimmed := bytes.TrimSpace(data) if len(trimmed) == 0 { return "unknown" @@ -45,7 +47,7 @@ func GetJsonType(data json.RawMessage) string { } // JsonRawMessageToString returns JSON strings as their decoded value and other JSON values as raw text. -func JsonRawMessageToString(data json.RawMessage) string { +func JsonRawMessageToString(data RawMessage) string { trimmed := bytes.TrimSpace(data) if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { return "" diff --git a/common/rate-limit.go b/common/rate-limit.go index 301c101c9748..73fea0835a37 100644 --- a/common/rate-limit.go +++ b/common/rate-limit.go @@ -45,6 +45,9 @@ func (l *InMemoryRateLimiter) clearExpiredItems() { func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration int64) bool { l.mutex.Lock() defer l.mutex.Unlock() + if maxRequestNum == 0 { + return true + } // [old <-- new] queue, ok := l.store[key] now := time.Now().Unix() @@ -68,3 +71,19 @@ func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration in } return true } + +// Check reports whether a request would be allowed without recording it. +// The duration parameter's unit is seconds. +func (l *InMemoryRateLimiter) Check(key string, maxRequestNum int, duration int64) bool { + l.mutex.Lock() + defer l.mutex.Unlock() + if maxRequestNum == 0 { + return true + } + queue, ok := l.store[key] + if !ok || len(*queue) < maxRequestNum { + return true + } + now := time.Now().Unix() + return now-(*queue)[0] >= duration +} diff --git a/controller/relay.go b/controller/relay.go index 1d14dcc6f880..cda442b93415 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -7,7 +7,6 @@ import ( "log" "net/http" "strings" - "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -65,6 +64,22 @@ func geminiRelayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewA return err } +func ResponsesWebSocket(c *gin.Context) { + requestId := c.GetString(common.RequestIdKey) + ws, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + return + } + defer ws.Close() + + if newAPIError := relay.ResponsesWebSocketHelper(c, ws); newAPIError != nil { + errorPreview := common.LocalLogPreview(newAPIError.Error()) + logger.LogError(c, fmt.Sprintf("responses websocket relay error: %s", errorPreview)) + newAPIError.SetMessage(common.MessageWithRequestId(newAPIError.Error(), requestId)) + helper.WssError(c, ws, newAPIError.ToOpenAIError()) + } +} + func Relay(c *gin.Context, relayFormat types.RelayFormat) { requestId := c.GetString(common.RequestIdKey) @@ -248,7 +263,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } var upgrader = websocket.Upgrader{ - Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol + Subprotocols: []string{"realtime", "responses"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol CheckOrigin: func(r *http.Request) bool { return true // 允许跨域 }, @@ -322,82 +337,11 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service } func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { - if openaiErr == nil { - return false - } - if service.ShouldSkipRetryAfterChannelAffinityFailure(c) { - return false - } - if types.IsChannelError(openaiErr) { - return true - } - if types.IsSkipRetryError(openaiErr) { - return false - } - if retryTimes <= 0 { - return false - } - if _, ok := c.Get("specific_channel_id"); ok { - return false - } - code := openaiErr.StatusCode - if code >= 200 && code < 300 { - return false - } - if code < 100 || code > 599 { - return true - } - if operation_setting.IsAlwaysSkipRetryCode(openaiErr.GetErrorCode()) { - return false - } - return operation_setting.ShouldRetryByStatusCode(code) + return service.ShouldRetryRelayError(c, openaiErr, retryTimes) } func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) { - logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.Error()))) - // 不要使用context获取渠道信息，异步处理时可能会出现渠道信息不一致的情况 - // do not use context to get channel info, there may be inconsistent channel info when processing asynchronously - if service.ShouldDisableChannel(err) && channelError.AutoBan { - gopool.Go(func() { - service.DisableChannel(channelError, err.ErrorWithStatusCode()) - }) - } - - if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) { - // 保存错误日志到mysql中 - userId := c.GetInt("id") - tokenName := c.GetString("token_name") - modelName := c.GetString("original_model") - tokenId := c.GetInt("token_id") - userGroup := c.GetString("group") - channelId := c.GetInt("channel_id") - other := make(map[string]interface{}) - if c.Request != nil && c.Request.URL != nil { - other["request_path"] = c.Request.URL.Path - } - other["error_type"] = err.GetErrorType() - other["error_code"] = err.GetErrorCode() - other["status_code"] = err.StatusCode - other["channel_id"] = channelId - other["channel_name"] = c.GetString("channel_name") - other["channel_type"] = c.GetInt("channel_type") - adminInfo := make(map[string]interface{}) - adminInfo["use_channel"] = c.GetStringSlice("use_channel") - isMultiKey := common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey) - if isMultiKey { - adminInfo["is_multi_key"] = true - adminInfo["multi_key_index"] = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex) - } - service.AppendChannelAffinityAdminInfo(c, adminInfo) - other["admin_info"] = adminInfo - startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime) - if startTime.IsZero() { - startTime = time.Now() - } - useTimeSeconds := int(time.Since(startTime).Seconds()) - model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other) - } - + service.ProcessChannelError(c, channelError, err) } func RelayMidjourney(c *gin.Context) { diff --git a/dto/openai_response.go b/dto/openai_response.go index 0e6b818dbd8b..868f6f9f96f2 100644 --- a/dto/openai_response.go +++ b/dto/openai_response.go @@ -228,11 +228,12 @@ type Usage struct { UsageSemantic string `json:"usage_semantic,omitempty"` UsageSource string `json:"usage_source,omitempty"` - PromptTokensDetails InputTokenDetails `json:"prompt_tokens_details"` - CompletionTokenDetails OutputTokenDetails `json:"completion_tokens_details"` - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` - InputTokensDetails *InputTokenDetails `json:"input_tokens_details"` + PromptTokensDetails InputTokenDetails `json:"prompt_tokens_details"` + CompletionTokenDetails OutputTokenDetails `json:"completion_tokens_details"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + InputTokensDetails *InputTokenDetails `json:"input_tokens_details"` + OutputTokensDetails *OutputTokenDetails `json:"output_tokens_details,omitempty"` // claude cache 1h ClaudeCacheCreation5mTokens int `json:"claude_cache_creation_5_m_tokens"` diff --git a/middleware/auth.go b/middleware/auth.go index 23d933fbe0c1..ef394788221b 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -276,20 +276,7 @@ func TokenAuthReadOnly() func(c *gin.Context) { func TokenAuth() func(c *gin.Context) { return func(c *gin.Context) { // 先检测是否为ws - if c.Request.Header.Get("Sec-WebSocket-Protocol") != "" { - // Sec-WebSocket-Protocol: realtime, openai-insecure-api-key.sk-xxx, openai-beta.realtime-v1 - // read sk from Sec-WebSocket-Protocol - key := c.Request.Header.Get("Sec-WebSocket-Protocol") - parts := strings.Split(key, ",") - for _, part := range parts { - part = strings.TrimSpace(part) - if strings.HasPrefix(part, "openai-insecure-api-key") { - key = strings.TrimPrefix(part, "openai-insecure-api-key.") - break - } - } - c.Request.Header.Set("Authorization", "Bearer "+key) - } + applyWebSocketSubprotocolAuthorization(c.Request.Header) // 检查path包含/v1/messages 或 /v1/models if strings.Contains(c.Request.URL.Path, "/v1/messages") || strings.Contains(c.Request.URL.Path, "/v1/models") { anthropicKey := c.Request.Header.Get("x-api-key") @@ -406,6 +393,31 @@ func TokenAuth() func(c *gin.Context) { } } +func applyWebSocketSubprotocolAuthorization(header http.Header) bool { + key, ok := apiKeyFromWebSocketSubprotocol(header.Get("Sec-WebSocket-Protocol")) + if !ok { + return false + } + header.Set("Authorization", "Bearer "+key) + return true +} + +func apiKeyFromWebSocketSubprotocol(protocols string) (string, bool) { + if protocols == "" { + return "", false + } + const insecureAPIKeyPrefix = "openai-insecure-api-key." + parts := strings.Split(protocols, ",") + for _, part := range parts { + part = strings.TrimSpace(part) + if strings.HasPrefix(part, insecureAPIKeyPrefix) { + key := strings.TrimPrefix(part, insecureAPIKeyPrefix) + return key, key != "" + } + } + return "", false +} + func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) error { if token == nil { return fmt.Errorf("token is nil") diff --git a/middleware/auth_test.go b/middleware/auth_test.go new file mode 100644 index 000000000000..edef280b5996 --- /dev/null +++ b/middleware/auth_test.go @@ -0,0 +1,86 @@ +package middleware + +import ( + "net/http" + "testing" +) + +func TestAPIKeyFromWebSocketSubprotocol(t *testing.T) { + tests := []struct { + name string + protocols string + wantKey string + wantOK bool + }{ + { + name: "responses protocol only", + protocols: "responses", + wantOK: false, + }, + { + name: "realtime protocol only", + protocols: "realtime", + wantOK: false, + }, + { + name: "responses with insecure key", + protocols: "responses, openai-insecure-api-key.sk-test", + wantKey: "sk-test", + wantOK: true, + }, + { + name: "realtime with beta and insecure key", + protocols: "realtime, openai-insecure-api-key.sk-realtime, openai-beta.realtime-v1", + wantKey: "sk-realtime", + wantOK: true, + }, + { + name: "empty insecure key", + protocols: "responses, openai-insecure-api-key.", + wantOK: false, + }, + { + name: "bare insecure marker is not a key", + protocols: "openai-insecure-api-key", + wantOK: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotKey, gotOK := apiKeyFromWebSocketSubprotocol(tt.protocols) + if gotOK != tt.wantOK { + t.Fatalf("ok = %v, want %v", gotOK, tt.wantOK) + } + if gotKey != tt.wantKey { + t.Fatalf("key = %q, want %q", gotKey, tt.wantKey) + } + }) + } +} + +func TestApplyWebSocketSubprotocolAuthorizationDoesNotOverrideProtocolOnly(t *testing.T) { + header := http.Header{} + header.Set("Authorization", "Bearer sk-original") + header.Set("Sec-WebSocket-Protocol", "responses") + + if applyWebSocketSubprotocolAuthorization(header) { + t.Fatal("authorization was unexpectedly applied") + } + if got := header.Get("Authorization"); got != "Bearer sk-original" { + t.Fatalf("Authorization = %q, want original bearer", got) + } +} + +func TestApplyWebSocketSubprotocolAuthorizationOverridesWithInsecureKey(t *testing.T) { + header := http.Header{} + header.Set("Authorization", "Bearer sk-original") + header.Set("Sec-WebSocket-Protocol", "responses, openai-insecure-api-key.sk-from-protocol") + + if !applyWebSocketSubprotocolAuthorization(header) { + t.Fatal("authorization was not applied") + } + if got := header.Get("Authorization"); got != "Bearer sk-from-protocol" { + t.Fatalf("Authorization = %q, want protocol bearer", got) + } +} diff --git a/middleware/model-rate-limit.go b/middleware/model-rate-limit.go index 80a3995df097..62b37e4d32d0 100644 --- a/middleware/model-rate-limit.go +++ b/middleware/model-rate-limit.go @@ -5,12 +5,14 @@ import ( "fmt" "net/http" "strconv" + "strings" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common/limiter" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/go-redis/redis/v8" @@ -21,6 +23,8 @@ const ( ModelRequestRateLimitSuccessCountMark = "MRRLS" ) +type ModelRequestRateLimitCommit func(success bool) + // 检查Redis中的请求限制 func checkRedisRateLimit(ctx context.Context, rdb *redis.Client, key string, maxCount int, duration int64) (bool, error) { // 如果maxCount为0，表示不限制 @@ -74,6 +78,107 @@ func recordRedisRequest(ctx context.Context, rdb *redis.Client, key string, maxC rdb.Expire(ctx, key, time.Duration(setting.ModelRequestRateLimitDurationMinutes)*time.Minute) } +func modelRequestRateLimitConfig(c *gin.Context) (duration int64, totalMaxCount int, successMaxCount int) { + duration = int64(setting.ModelRequestRateLimitDurationMinutes * 60) + totalMaxCount = setting.ModelRequestRateLimitCount + successMaxCount = setting.ModelRequestRateLimitSuccessCount + + group := common.GetContextKeyString(c, constant.ContextKeyTokenGroup) + if group == "" { + group = common.GetContextKeyString(c, constant.ContextKeyUserGroup) + } + groupTotalCount, groupSuccessCount, found := setting.GetGroupRateLimit(group) + if found { + totalMaxCount = groupTotalCount + successMaxCount = groupSuccessCount + } + return duration, totalMaxCount, successMaxCount +} + +func newModelRateLimitError(message string, statusCode int) *types.NewAPIError { + return types.NewErrorWithStatusCode( + fmt.Errorf("%s", message), + types.ErrorCodeInvalidRequest, + statusCode, + types.ErrOptionWithSkipRetry(), + types.ErrOptionWithNoRecordErrorLog(), + ) +} + +func CheckModelRequestRateLimit(c *gin.Context) (ModelRequestRateLimitCommit, *types.NewAPIError) { + if !setting.ModelRequestRateLimitEnabled { + return func(bool) {}, nil + } + + duration, totalMaxCount, successMaxCount := modelRequestRateLimitConfig(c) + userId := strconv.Itoa(c.GetInt("id")) + + if common.RedisEnabled { + ctx := context.Background() + rdb := common.RDB + successKey := fmt.Sprintf("rateLimit:%s:%s", ModelRequestRateLimitSuccessCountMark, userId) + allowed, err := checkRedisRateLimit(ctx, rdb, successKey, successMaxCount, duration) + if err != nil { + fmt.Println("检查成功请求数限制失败:", err.Error()) + return nil, newModelRateLimitError("rate_limit_check_failed", http.StatusInternalServerError) + } + if !allowed { + return nil, newModelRateLimitError(fmt.Sprintf("您已达到请求数限制：%d分钟内最多请求%d次", setting.ModelRequestRateLimitDurationMinutes, successMaxCount), http.StatusTooManyRequests) + } + + if totalMaxCount > 0 { + totalKey := fmt.Sprintf("rateLimit:%s", userId) + tb := limiter.New(ctx, rdb) + allowed, err = tb.Allow( + ctx, + totalKey, + limiter.WithCapacity(int64(totalMaxCount)*duration), + limiter.WithRate(int64(totalMaxCount)), + limiter.WithRequested(duration), + ) + if err != nil { + fmt.Println("检查总请求数限制失败:", err.Error()) + return nil, newModelRateLimitError("rate_limit_check_failed", http.StatusInternalServerError) + } + if !allowed { + return nil, newModelRateLimitError(fmt.Sprintf("您已达到总请求数限制：%d分钟内最多请求%d次，包括失败次数，请检查您的请求是否正确", setting.ModelRequestRateLimitDurationMinutes, totalMaxCount), http.StatusTooManyRequests) + } + } + + return func(success bool) { + if success { + recordRedisRequest(ctx, rdb, successKey, successMaxCount) + } + }, nil + } + + inMemoryRateLimiter.Init(time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute) + totalKey := ModelRequestRateLimitCountMark + userId + successKey := ModelRequestRateLimitSuccessCountMark + userId + + if totalMaxCount > 0 && !inMemoryRateLimiter.Request(totalKey, totalMaxCount, duration) { + return nil, newModelRateLimitError(fmt.Sprintf("您已达到总请求数限制：%d分钟内最多请求%d次，包括失败次数，请检查您的请求是否正确", setting.ModelRequestRateLimitDurationMinutes, totalMaxCount), http.StatusTooManyRequests) + } + if successMaxCount > 0 && !inMemoryRateLimiter.Check(successKey, successMaxCount, duration) { + return nil, newModelRateLimitError(fmt.Sprintf("您已达到请求数限制：%d分钟内最多请求%d次", setting.ModelRequestRateLimitDurationMinutes, successMaxCount), http.StatusTooManyRequests) + } + + return func(success bool) { + if success && successMaxCount > 0 { + inMemoryRateLimiter.Request(successKey, successMaxCount, duration) + } + }, nil +} + +func isResponsesWebSocketHandshake(c *gin.Context) bool { + return c != nil && + c.Request != nil && + c.Request.Method == http.MethodGet && + c.Request.URL != nil && + c.Request.URL.Path == "/v1/responses" && + strings.EqualFold(c.Request.Header.Get("Upgrade"), "websocket") +} + // Redis限流处理器 func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { return func(c *gin.Context) { @@ -166,35 +271,16 @@ func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) // ModelRequestRateLimit 模型请求限流中间件 func ModelRequestRateLimit() func(c *gin.Context) { return func(c *gin.Context) { - // 在每个请求时检查是否启用限流 - if !setting.ModelRequestRateLimitEnabled { + if isResponsesWebSocketHandshake(c) { c.Next() return } - - // 计算限流参数 - duration := int64(setting.ModelRequestRateLimitDurationMinutes * 60) - totalMaxCount := setting.ModelRequestRateLimitCount - successMaxCount := setting.ModelRequestRateLimitSuccessCount - - // 获取分组 - group := common.GetContextKeyString(c, constant.ContextKeyTokenGroup) - if group == "" { - group = common.GetContextKeyString(c, constant.ContextKeyUserGroup) - } - - //获取分组的限流配置 - groupTotalCount, groupSuccessCount, found := setting.GetGroupRateLimit(group) - if found { - totalMaxCount = groupTotalCount - successMaxCount = groupSuccessCount - } - - // 根据存储类型选择并执行限流处理器 - if common.RedisEnabled { - redisRateLimitHandler(duration, totalMaxCount, successMaxCount)(c) - } else { - memoryRateLimitHandler(duration, totalMaxCount, successMaxCount)(c) + commit, apiErr := CheckModelRequestRateLimit(c) + if apiErr != nil { + abortWithOpenAiMessage(c, apiErr.StatusCode, apiErr.Error(), apiErr.GetErrorCode()) + return } + c.Next() + commit(c.Writer.Status() < 400) } } diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 2665b8d027e9..59a51e1bad13 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -46,12 +46,7 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http // compute usage usage := dto.Usage{} if responsesResponse.Usage != nil { - usage.PromptTokens = responsesResponse.Usage.InputTokens - usage.CompletionTokens = responsesResponse.Usage.OutputTokens - usage.TotalTokens = responsesResponse.Usage.TotalTokens - if responsesResponse.Usage.InputTokensDetails != nil { - usage.PromptTokensDetails.CachedTokens = responsesResponse.Usage.InputTokensDetails.CachedTokens - } + service.ApplyResponsesUsage(&usage, responsesResponse.Usage) } if info == nil || info.ResponsesUsageInfo == nil || info.ResponsesUsageInfo.BuiltInTools == nil { return &usage, nil @@ -93,18 +88,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp case "response.completed": if streamResponse.Response != nil { if streamResponse.Response.Usage != nil { - if streamResponse.Response.Usage.InputTokens != 0 { - usage.PromptTokens = streamResponse.Response.Usage.InputTokens - } - if streamResponse.Response.Usage.OutputTokens != 0 { - usage.CompletionTokens = streamResponse.Response.Usage.OutputTokens - } - if streamResponse.Response.Usage.TotalTokens != 0 { - usage.TotalTokens = streamResponse.Response.Usage.TotalTokens - } - if streamResponse.Response.Usage.InputTokensDetails != nil { - usage.PromptTokensDetails.CachedTokens = streamResponse.Response.Usage.InputTokensDetails.CachedTokens - } + service.ApplyResponsesUsage(usage, streamResponse.Response.Usage) } if streamResponse.Response.HasImageGenerationCall() { c.Set("image_generation_call", true) diff --git a/relay/responses_websocket.go b/relay/responses_websocket.go new file mode 100644 index 000000000000..c516b7859f40 --- /dev/null +++ b/relay/responses_websocket.go @@ -0,0 +1,874 @@ +package relay + +import ( + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + appconstant "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/middleware" + appmodel "github.com/QuantumNous/new-api/model" + relaychannel "github.com/QuantumNous/new-api/relay/channel" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relay/helper" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" +) + +const responsesWSEventTypeResponseCreate = "response.create" + +type responsesWSCreateEvent struct { + Type string `json:"type"` + EventID string `json:"event_id,omitempty"` + Request common.RawMessage `json:"response,omitempty"` +} + +type responsesWSCreateRequest struct { + Request dto.OpenAIResponsesRequest + Generate common.RawMessage +} + +type responsesWSErrorEvent struct { + Type string `json:"type"` + Status int `json:"status"` + EventID string `json:"event_id,omitempty"` + Error *types.OpenAIError `json:"error"` +} + +type responsesWSCallState struct { + info *relaycommon.RelayInfo + usage *dto.Usage + outputText strings.Builder + commitRate middleware.ModelRequestRateLimitCommit +} + +type responsesWSSession struct { + c *gin.Context + client *websocket.Conn + target *websocket.Conn + lockedModel string + lockedChannel *appmodel.Channel + nextEventIndex int + + clientWriteMu sync.Mutex + targetWriteMu sync.Mutex + stateMu sync.Mutex + current *responsesWSCallState +} + +func ResponsesWebSocketHelper(c *gin.Context, client *websocket.Conn) *types.NewAPIError { + session := &responsesWSSession{ + c: c, + client: client, + } + defer session.closeTarget() + defer session.failCurrent() + + for { + messageType, message, err := client.ReadMessage() + if err != nil { + if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + return nil + } + return types.NewError(err, types.ErrorCodeBadRequestBody, types.ErrOptionWithSkipRetry()) + } + + eventType, eventErr := responsesWSEventType(message) + if eventErr != nil { + session.sendError("", newResponsesWSInvalidRequestError(eventErr)) + continue + } + + if eventType != responsesWSEventTypeResponseCreate { + if !session.hasTarget() { + session.sendError("", newResponsesWSInvalidRequestError(errors.New("first responses websocket event must be response.create"))) + continue + } + if err := session.writeTarget(messageType, message); err != nil { + return session.handleControlEventWriteFailure(err) + } + continue + } + + create, eventID, err := normalizeResponsesWSCreateEvent(message) + if err != nil { + session.sendError("", newResponsesWSInvalidRequestError(err)) + continue + } + if create.Request.Model == "" { + session.sendError(eventID, newResponsesWSInvalidRequestError(errors.New("model is required"))) + continue + } + if err := session.handleResponseCreate(create, eventID); err != nil { + session.sendError(eventID, err) + } + } +} + +func responsesWSEventType(message []byte) (string, error) { + var event struct { + Type string `json:"type"` + } + if err := common.Unmarshal(message, &event); err != nil { + return "", fmt.Errorf("invalid websocket event json: %w", err) + } + if strings.TrimSpace(event.Type) == "" { + return "", errors.New("websocket event type is required") + } + return event.Type, nil +} + +func newResponsesWSInvalidRequestError(err error) *types.NewAPIError { + return types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) +} + +func normalizeResponsesWSCreateEvent(message []byte) (responsesWSCreateRequest, string, error) { + var event responsesWSCreateEvent + if err := common.Unmarshal(message, &event); err != nil { + return responsesWSCreateRequest{}, "", err + } + if event.Type != responsesWSEventTypeResponseCreate { + return responsesWSCreateRequest{}, event.EventID, fmt.Errorf("unsupported event type %q", event.Type) + } + + var generate common.RawMessage + var raw map[string]common.RawMessage + if err := common.Unmarshal(message, &raw); err == nil { + if generateRaw, ok := raw["generate"]; ok { + generate = generateRaw + } + } + + payload := event.Request + if len(payload) == 0 { + if err := common.Unmarshal(message, &raw); err != nil { + return responsesWSCreateRequest{}, event.EventID, err + } + delete(raw, "type") + delete(raw, "event_id") + delete(raw, "background") + delete(raw, "generate") + delete(raw, "stream") + delete(raw, "stream_options") + var err error + payload, err = common.Marshal(raw) + if err != nil { + return responsesWSCreateRequest{}, event.EventID, err + } + } else { + var responseMap map[string]common.RawMessage + if err := common.Unmarshal(payload, &responseMap); err == nil { + if len(generate) == 0 { + if generateRaw, ok := responseMap["generate"]; ok { + generate = generateRaw + } + } + if _, exists := responseMap["generate"]; exists { + delete(responseMap, "generate") + if merged, err := common.Marshal(responseMap); err == nil { + payload = merged + } + } + } + } + + var req dto.OpenAIResponsesRequest + if err := common.Unmarshal(payload, &req); err != nil { + return responsesWSCreateRequest{}, event.EventID, err + } + req.Stream = nil + req.StreamOptions = nil + return responsesWSCreateRequest{ + Request: req, + Generate: generate, + }, event.EventID, nil +} + +func (s *responsesWSSession) handleResponseCreate(create responsesWSCreateRequest, eventID string) *types.NewAPIError { + req := create.Request + if s.lockedModel != "" && req.Model != s.lockedModel { + return types.NewErrorWithStatusCode( + fmt.Errorf("responses websocket connection is locked to model %q; got %q", s.lockedModel, req.Model), + types.ErrorCodeInvalidRequest, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + + if s.hasCurrent() { + return types.NewErrorWithStatusCode( + errors.New("another response.create is already in progress on this websocket connection"), + types.ErrorCodeInvalidRequest, + http.StatusConflict, + types.ErrOptionWithSkipRetry(), + ) + } + + commitRate, apiErr := middleware.CheckModelRequestRateLimit(s.c) + if apiErr != nil { + return apiErr + } + + if !s.hasTarget() { + return s.connectAndSendFirst(create, commitRate) + } + + state, payload, apiErr := s.prepareCall(create, commitRate) + if apiErr != nil { + commitRate(false) + return apiErr + } + if !s.tryReserveCurrent(state) { + state.refund(s.c) + commitRate(false) + return types.NewErrorWithStatusCode( + errors.New("another response.create is already in progress on this websocket connection"), + types.ErrorCodeInvalidRequest, + http.StatusConflict, + types.ErrOptionWithSkipRetry(), + ) + } + if err := s.writeTarget(websocket.TextMessage, payload); err != nil { + return s.handleTargetWriteFailureWithState(state, err) + } + return nil +} + +func (s *responsesWSSession) handleControlEventWriteFailure(err error) *types.NewAPIError { + apiErr := s.handleTargetWriteFailure(err) + s.sendError("", apiErr) + return nil +} + +func (s *responsesWSSession) handleTargetWriteFailure(err error) *types.NewAPIError { + s.closeTarget() + apiErr := types.NewError(err, types.ErrorCodeBadResponse) + apiErr, _ = s.processChannelError(s.lockedChannel, apiErr, nil) + return apiErr +} + +func (s *responsesWSSession) handleTargetWriteFailureWithState(state *responsesWSCallState, err error) *types.NewAPIError { + s.finishCall(state, false) + return s.handleTargetWriteFailure(err) +} + +func (s *responsesWSSession) connectAndSendFirst(create responsesWSCreateRequest, commitRate middleware.ModelRequestRateLimitCommit) *types.NewAPIError { + req := create.Request + if err := checkResponsesWSModelAccess(s.c, req.Model); err != nil { + commitRate(false) + return err + } + + retryParam := &service.RetryParam{ + Ctx: s.c, + TokenGroup: common.GetContextKeyString(s.c, appconstant.ContextKeyUsingGroup), + ModelName: req.Model, + Retry: common.GetPointer(0), + } + if retryParam.TokenGroup == "" { + retryParam.TokenGroup = common.GetContextKeyString(s.c, appconstant.ContextKeyTokenGroup) + } + + var lastErr *types.NewAPIError + for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { + channel, apiErr := selectResponsesWSChannel(s.c, req.Model, retryParam) + if apiErr != nil { + lastErr = apiErr + break + } + addResponsesWSUsedChannel(s.c, channel.Id) + + if channel.Type != appconstant.ChannelTypeOpenAI && channel.Type != appconstant.ChannelTypeCodex { + lastErr = types.NewErrorWithStatusCode( + fmt.Errorf("responses websocket only supports OpenAI and Codex channels, got channel type %d", channel.Type), + types.ErrorCodeInvalidRequest, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + continue + } + + state, payload, apiErr := s.prepareCall(create, commitRate) + if apiErr != nil { + commitRate(false) + return apiErr + } + + adaptor := GetAdaptor(state.info.ApiType) + if adaptor == nil { + state.refund(s.c) + apiErr = types.NewError(fmt.Errorf("invalid api type: %d", state.info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry()) + var shouldRetry bool + lastErr, shouldRetry = s.processChannelError(channel, apiErr, retryParam) + if !shouldRetry { + break + } + continue + } + adaptor.Init(state.info) + target, apiErr := dialResponsesWebSocketUpstream(s.c, adaptor, state.info) + if apiErr != nil { + state.refund(s.c) + var shouldRetry bool + lastErr, shouldRetry = s.processChannelError(channel, apiErr, retryParam) + if !shouldRetry { + break + } + continue + } + + s.setTarget(target) + if !s.tryReserveCurrent(state) { + s.closeTarget() + state.refund(s.c) + commitRate(false) + return types.NewErrorWithStatusCode(errors.New("another response.create is already in progress on this websocket connection"), types.ErrorCodeInvalidRequest, http.StatusConflict, types.ErrOptionWithSkipRetry()) + } + if err := s.writeTarget(websocket.TextMessage, payload); err != nil { + s.finishCall(state, false) + s.closeTarget() + apiErr = types.NewError(err, types.ErrorCodeBadResponse) + var shouldRetry bool + lastErr, shouldRetry = s.processChannelError(channel, apiErr, retryParam) + if !shouldRetry { + break + } + continue + } + + s.lockedModel = req.Model + s.lockedChannel = channel + service.RecordChannelAffinity(s.c, channel.Id) + s.startTargetReader() + return nil + } + + if lastErr == nil { + lastErr = types.NewError(errors.New("failed to connect responses websocket upstream"), types.ErrorCodeDoRequestFailed, types.ErrOptionWithSkipRetry()) + } + commitRate(false) + return lastErr +} + +func (s *responsesWSSession) processChannelError(channel *appmodel.Channel, apiErr *types.NewAPIError, retryParam *service.RetryParam) (*types.NewAPIError, bool) { + if apiErr == nil { + return nil, false + } + apiErr = service.NormalizeViolationFeeError(apiErr) + statusCodeMapping := "" + if s.c != nil { + statusCodeMapping = s.c.GetString("status_code_mapping") + } + service.ResetStatusCode(apiErr, statusCodeMapping) + if channel != nil && s.c != nil { + service.ProcessChannelError(s.c, *types.NewChannelError( + channel.Id, + channel.Type, + channel.Name, + channel.ChannelInfo.IsMultiKey, + common.GetContextKeyString(s.c, appconstant.ContextKeyChannelKey), + channel.GetAutoBan(), + ), apiErr) + } + if retryParam == nil { + return apiErr, false + } + return apiErr, service.ShouldRetryRelayError(s.c, apiErr, common.RetryTimes-retryParam.GetRetry()) +} + +func (s *responsesWSSession) prepareCall(create responsesWSCreateRequest, commitRate middleware.ModelRequestRateLimitCommit) (*responsesWSCallState, []byte, *types.NewAPIError) { + req := create.Request + common.SetContextKey(s.c, appconstant.ContextKeyRequestStartTime, time.Now()) + relayInfo := relaycommon.GenRelayInfoResponses(s.c, &req) + relayInfo.RequestId = fmt.Sprintf("%s-ws-%d", relayInfo.RequestId, s.nextEventIndex) + s.nextEventIndex++ + + meta := req.GetTokenCountMeta() + if setting.ShouldCheckPromptSensitive() && meta != nil { + contains, words := service.CheckSensitiveText(meta.CombineText) + if contains { + return nil, nil, types.NewError(fmt.Errorf("user sensitive words detected: %s", strings.Join(words, ", ")), types.ErrorCodeSensitiveWordsDetected, types.ErrOptionWithSkipRetry()) + } + } + + tokens, err := service.EstimateRequestToken(s.c, meta, relayInfo) + if err != nil { + return nil, nil, types.NewError(err, types.ErrorCodeCountTokenFailed) + } + relayInfo.SetEstimatePromptTokens(tokens) + + priceData, err := helper.ModelPriceHelper(s.c, relayInfo, tokens, meta) + if err != nil { + return nil, nil, types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest)) + } + if !priceData.FreeModel { + if apiErr := service.PreConsumeBilling(s.c, priceData.QuotaToPreConsume, relayInfo); apiErr != nil { + return nil, nil, apiErr + } + } + + payload, apiErr := buildResponsesWSCreatePayload(s.c, relayInfo, req, create.Generate) + if apiErr != nil { + if relayInfo.Billing != nil { + relayInfo.Billing.Refund(s.c) + } + return nil, nil, apiErr + } + + return &responsesWSCallState{ + info: relayInfo, + usage: &dto.Usage{}, + commitRate: commitRate, + }, payload, nil +} + +func buildResponsesWSCreatePayload(c *gin.Context, relayInfo *relaycommon.RelayInfo, req dto.OpenAIResponsesRequest, generate common.RawMessage) ([]byte, *types.NewAPIError) { + relayInfo.InitChannelMeta(c) + request, err := common.DeepCopy(&req) + if err != nil { + return nil, types.NewError(fmt.Errorf("failed to copy responses request: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) + } + if err := helper.ModelMappedHelper(c, relayInfo, request); err != nil { + return nil, types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry()) + } + + adaptor := GetAdaptor(relayInfo.ApiType) + if adaptor == nil { + return nil, types.NewError(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry()) + } + adaptor.Init(relayInfo) + convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, relayInfo, *request) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + relaycommon.AppendRequestConversionFromRequest(relayInfo, convertedRequest) + jsonData, err := common.Marshal(convertedRequest) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + jsonData, err = relaycommon.RemoveDisabledFields(jsonData, relayInfo.ChannelOtherSettings, relayInfo.ChannelSetting.PassThroughBodyEnabled) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + jsonData, err = removeResponsesWSTransportFields(jsonData) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + if len(relayInfo.ParamOverride) > 0 { + jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, relayInfo) + if err != nil { + return nil, newAPIErrorFromParamOverride(err) + } + } + + event, err := buildResponsesWSCreateEvent(jsonData, generate) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + return event, nil +} + +func buildResponsesWSCreateEvent(jsonData []byte, generate common.RawMessage) ([]byte, error) { + var event map[string]common.RawMessage + if err := common.Unmarshal(jsonData, &event); err != nil { + return nil, err + } + typeData, err := common.Marshal(responsesWSEventTypeResponseCreate) + if err != nil { + return nil, err + } + event["type"] = typeData + delete(event, "event_id") + delete(event, "background") + delete(event, "stream") + delete(event, "stream_options") + if len(generate) > 0 { + event["generate"] = generate + } + return common.Marshal(event) +} + +func removeResponsesWSTransportFields(jsonData []byte) ([]byte, error) { + var data map[string]any + if err := common.Unmarshal(jsonData, &data); err != nil { + return jsonData, err + } + delete(data, "stream") + delete(data, "stream_options") + delete(data, "background") + return common.Marshal(data) +} + +func dialResponsesWebSocketUpstream(c *gin.Context, adaptor relaychannel.Adaptor, info *relaycommon.RelayInfo) (*websocket.Conn, *types.NewAPIError) { + fullRequestURL, err := adaptor.GetRequestURL(info) + if err != nil { + return nil, types.NewError(fmt.Errorf("get request url failed: %w", err), types.ErrorCodeDoRequestFailed) + } + fullRequestURL = toWebSocketURL(fullRequestURL) + + targetHeader := http.Header{} + if err := adaptor.SetupRequestHeader(c, &targetHeader, info); err != nil { + return nil, types.NewError(fmt.Errorf("setup request header failed: %w", err), types.ErrorCodeDoRequestFailed) + } + headerOverride, err := relaychannel.ResolveHeaderOverride(info, c) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeChannelHeaderOverrideInvalid) + } + for key, value := range headerOverride { + targetHeader.Set(key, value) + } + + targetConn, resp, err := websocket.DefaultDialer.Dial(fullRequestURL, targetHeader) + if err != nil { + statusCode := http.StatusInternalServerError + if resp != nil { + statusCode = resp.StatusCode + } + return nil, types.NewErrorWithStatusCode(fmt.Errorf("dial failed to %s: %w", fullRequestURL, err), types.ErrorCodeDoRequestFailed, statusCode) + } + return targetConn, nil +} + +func toWebSocketURL(raw string) string { + switch { + case strings.HasPrefix(raw, "https://"): + return "wss://" + strings.TrimPrefix(raw, "https://") + case strings.HasPrefix(raw, "http://"): + return "ws://" + strings.TrimPrefix(raw, "http://") + default: + return raw + } +} + +func (s *responsesWSSession) startTargetReader() { + target := s.getTarget() + if target == nil { + return + } + go func() { + for { + messageType, message, err := target.ReadMessage() + if err != nil { + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + logger.LogError(s.c, "responses websocket upstream read failed: "+err.Error()) + } + s.failCurrent() + _ = s.client.Close() + return + } + s.observeUpstreamMessage(message) + if err := s.writeClient(messageType, message); err != nil { + logger.LogError(s.c, "responses websocket client write failed: "+err.Error()) + s.failCurrent() + s.closeTarget() + return + } + } + }() +} + +func (s *responsesWSSession) observeUpstreamMessage(message []byte) { + state := s.getCurrent() + if state == nil { + return + } + state.info.SetFirstResponseTime() + + var streamResponse dto.ResponsesStreamResponse + if err := common.Unmarshal(message, &streamResponse); err != nil { + return + } + + switch streamResponse.Type { + case "response.completed", "response.done", "response.incomplete": + s.applyTerminalResponseUsage(state, streamResponse.Response) + s.finishCall(state, true) + case "response.failed", "response.cancelled", "response.canceled": + s.finishCall(state, false) + case "response.output_text.delta": + state.outputText.WriteString(streamResponse.Delta) + case dto.ResponsesOutputTypeItemDone: + if streamResponse.Item != nil && streamResponse.Item.Type == dto.BuildInCallWebSearchCall { + if state.info != nil && state.info.ResponsesUsageInfo != nil && state.info.ResponsesUsageInfo.BuiltInTools != nil { + if webSearchTool, exists := state.info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview]; exists && webSearchTool != nil { + webSearchTool.CallCount++ + } + } + } + case "error": + s.finishCall(state, false) + } +} + +func (s *responsesWSSession) applyTerminalResponseUsage(state *responsesWSCallState, response *dto.OpenAIResponsesResponse) { + if state == nil || response == nil { + return + } + if response.Usage != nil { + service.ApplyResponsesUsage(state.usage, response.Usage) + } + if response.HasImageGenerationCall() { + s.c.Set("image_generation_call", true) + s.c.Set("image_generation_call_quality", response.GetQuality()) + s.c.Set("image_generation_call_size", response.GetSize()) + } +} + +func (s *responsesWSSession) finishCall(state *responsesWSCallState, success bool) { + if state == nil { + return + } + if !s.clearCurrent(state) { + return + } + if !success { + state.refund(s.c) + if state.commitRate != nil { + state.commitRate(false) + } + return + } + + finalizeResponsesWSUsage(state) + service.PostTextConsumeQuota(s.c, state.info, state.usage, nil) + if state.commitRate != nil { + state.commitRate(true) + } +} + +func finalizeResponsesWSUsage(state *responsesWSCallState) { + if state == nil || state.usage == nil || state.info == nil { + return + } + if state.usage.CompletionTokens == 0 { + if output := state.outputText.String(); output != "" { + state.usage.CompletionTokens = service.CountTextToken(output, state.info.UpstreamModelName) + } + } + if state.usage.PromptTokens == 0 && state.usage.CompletionTokens != 0 { + state.usage.PromptTokens = state.info.GetEstimatePromptTokens() + } + if state.usage.TotalTokens == 0 { + state.usage.TotalTokens = state.usage.PromptTokens + state.usage.CompletionTokens + } +} + +func (state *responsesWSCallState) refund(c *gin.Context) { + if state != nil && state.info != nil && state.info.Billing != nil { + state.info.Billing.Refund(c) + } +} + +func (s *responsesWSSession) tryReserveCurrent(state *responsesWSCallState) bool { + s.stateMu.Lock() + defer s.stateMu.Unlock() + if s.current != nil { + return false + } + s.current = state + return true +} + +func (s *responsesWSSession) hasCurrent() bool { + s.stateMu.Lock() + defer s.stateMu.Unlock() + return s.current != nil +} + +func (s *responsesWSSession) clearCurrent(state *responsesWSCallState) bool { + s.stateMu.Lock() + defer s.stateMu.Unlock() + if state != nil && s.current != state { + return false + } + s.current = nil + return true +} + +func (s *responsesWSSession) getCurrent() *responsesWSCallState { + s.stateMu.Lock() + defer s.stateMu.Unlock() + return s.current +} + +func (s *responsesWSSession) failCurrent() { + state := s.getCurrent() + if state != nil { + s.finishCall(state, false) + } +} + +func (s *responsesWSSession) writeClient(messageType int, message []byte) error { + s.clientWriteMu.Lock() + defer s.clientWriteMu.Unlock() + return s.client.WriteMessage(messageType, message) +} + +func (s *responsesWSSession) hasTarget() bool { + s.targetWriteMu.Lock() + defer s.targetWriteMu.Unlock() + return s.target != nil +} + +func (s *responsesWSSession) getTarget() *websocket.Conn { + s.targetWriteMu.Lock() + defer s.targetWriteMu.Unlock() + return s.target +} + +func (s *responsesWSSession) setTarget(target *websocket.Conn) { + s.targetWriteMu.Lock() + defer s.targetWriteMu.Unlock() + s.target = target +} + +func (s *responsesWSSession) writeTarget(messageType int, message []byte) error { + s.targetWriteMu.Lock() + defer s.targetWriteMu.Unlock() + if s.target == nil { + return errors.New("responses websocket upstream is not connected") + } + return s.target.WriteMessage(messageType, message) +} + +func (s *responsesWSSession) sendError(eventID string, apiErr *types.NewAPIError) { + if apiErr == nil { + return + } + payload, err := buildResponsesWSErrorPayload(eventID, apiErr) + if err != nil { + return + } + _ = s.writeClient(websocket.TextMessage, payload) +} + +func buildResponsesWSErrorPayload(eventID string, apiErr *types.NewAPIError) ([]byte, error) { + if apiErr == nil { + return nil, errors.New("api error is nil") + } + status := apiErr.StatusCode + if status == 0 { + status = http.StatusInternalServerError + } + openaiErr := apiErr.ToOpenAIError() + return common.Marshal(&responsesWSErrorEvent{ + Type: "error", + Status: status, + EventID: eventID, + Error: &openaiErr, + }) +} + +func (s *responsesWSSession) closeTarget() { + s.targetWriteMu.Lock() + defer s.targetWriteMu.Unlock() + if s.target != nil { + _ = s.target.Close() + s.target = nil + } +} + +func checkResponsesWSModelAccess(c *gin.Context, modelName string) *types.NewAPIError { + if !common.GetContextKeyBool(c, appconstant.ContextKeyTokenModelLimitEnabled) { + return nil + } + raw, ok := common.GetContextKey(c, appconstant.ContextKeyTokenModelLimit) + if !ok { + return types.NewErrorWithStatusCode(errors.New("token has no model access"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + tokenModelLimit, ok := raw.(map[string]bool) + if !ok { + tokenModelLimit = map[string]bool{} + } + matchName := ratio_setting.FormatMatchingModelName(modelName) + if _, ok := tokenModelLimit[matchName]; !ok { + return types.NewErrorWithStatusCode(fmt.Errorf("token is not allowed to use model %s", modelName), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + return nil +} + +func selectResponsesWSChannel(c *gin.Context, modelName string, retryParam *service.RetryParam) (*appmodel.Channel, *types.NewAPIError) { + if channelIdRaw, ok := common.GetContextKey(c, appconstant.ContextKeyTokenSpecificChannelId); ok { + channelID, ok := channelIdRaw.(string) + if !ok { + return nil, types.NewErrorWithStatusCode(errors.New("invalid specified channel id"), types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + id, err := strconv.Atoi(channelID) + if err != nil { + return nil, types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + channel, err := appmodel.GetChannelById(id, true) + if err != nil { + return nil, types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + if channel.Status != common.ChannelStatusEnabled { + return nil, types.NewErrorWithStatusCode(errors.New("specified channel is disabled"), types.ErrorCodeGetChannelFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + if err := middleware.SetupContextForSelectedChannel(c, channel, modelName); err != nil { + return nil, err + } + return channel, nil + } + + usingGroup := common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup) + if usingGroup == "" { + usingGroup = retryParam.TokenGroup + } + + if retryParam.GetRetry() == 0 { + if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelName, usingGroup); found { + preferred, err := appmodel.CacheGetChannel(preferredChannelID) + if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled { + if usingGroup == "auto" { + userGroup := common.GetContextKeyString(c, appconstant.ContextKeyUserGroup) + for _, g := range service.GetUserAutoGroup(userGroup) { + if appmodel.IsChannelEnabledForGroupModel(g, modelName, preferred.Id) { + common.SetContextKey(c, appconstant.ContextKeyAutoGroup, g) + service.MarkChannelAffinityUsed(c, g, preferred.Id) + if err := middleware.SetupContextForSelectedChannel(c, preferred, modelName); err != nil { + return nil, err + } + return preferred, nil + } + } + } else if appmodel.IsChannelEnabledForGroupModel(usingGroup, modelName, preferred.Id) { + service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) + if err := middleware.SetupContextForSelectedChannel(c, preferred, modelName); err != nil { + return nil, err + } + return preferred, nil + } + } + } + } + + channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam) + if err != nil { + return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, modelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) + } + if channel == nil { + return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在（retry）", selectGroup, modelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) + } + if err := middleware.SetupContextForSelectedChannel(c, channel, modelName); err != nil { + return nil, err + } + return channel, nil +} + +func addResponsesWSUsedChannel(c *gin.Context, channelId int) { + useChannel := c.GetStringSlice("use_channel") + useChannel = append(useChannel, fmt.Sprintf("%d", channelId)) + c.Set("use_channel", useChannel) +} diff --git a/relay/responses_websocket_test.go b/relay/responses_websocket_test.go new file mode 100644 index 000000000000..38e73df130d3 --- /dev/null +++ b/relay/responses_websocket_test.go @@ -0,0 +1,350 @@ +package relay + +import ( + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + + "github.com/gorilla/websocket" +) + +func TestNormalizeResponsesWSCreateEventWrapper(t *testing.T) { + message := []byte(`{ + "type": "response.create", + "event_id": "evt_1", + "generate": false, + "response": { + "model": "gpt-5.3-codex-spark", + "input": "hi", + "store": false, + "stream": true, + "stream_options": {"include_usage": true} + } + }`) + + create, eventID, err := normalizeResponsesWSCreateEvent(message) + if err != nil { + t.Fatalf("normalizeResponsesWSCreateEvent() error = %v", err) + } + req := create.Request + if eventID != "evt_1" { + t.Fatalf("eventID = %q, want evt_1", eventID) + } + if req.Model != "gpt-5.3-codex-spark" { + t.Fatalf("model = %q", req.Model) + } + if strings.TrimSpace(string(create.Generate)) != "false" { + t.Fatalf("generate = %s, want false", create.Generate) + } + if req.Stream != nil { + t.Fatalf("stream = %v, want nil", req.Stream) + } + if req.StreamOptions != nil { + t.Fatalf("stream_options = %#v, want nil", req.StreamOptions) + } + if strings.TrimSpace(string(req.Store)) != "false" { + t.Fatalf("store = %s, want false", req.Store) + } +} + +func TestNormalizeResponsesWSCreateEventFlat(t *testing.T) { + message := []byte(`{ + "type": "response.create", + "event_id": "evt_2", + "model": "gpt-5.3-codex-spark", + "input": "hi", + "generate": false, + "stream": true, + "background": true, + "stream_options": {"include_usage": true} + }`) + + create, eventID, err := normalizeResponsesWSCreateEvent(message) + if err != nil { + t.Fatalf("normalizeResponsesWSCreateEvent() error = %v", err) + } + req := create.Request + if eventID != "evt_2" { + t.Fatalf("eventID = %q, want evt_2", eventID) + } + if req.Model != "gpt-5.3-codex-spark" { + t.Fatalf("model = %q", req.Model) + } + if strings.TrimSpace(string(create.Generate)) != "false" { + t.Fatalf("generate = %s, want false", create.Generate) + } + if req.Stream != nil { + t.Fatalf("stream = %v, want nil", req.Stream) + } + if req.StreamOptions != nil { + t.Fatalf("stream_options = %#v, want nil", req.StreamOptions) + } +} + +func TestBuildResponsesWSCreateEventIsFlat(t *testing.T) { + payload := []byte(`{ + "model": "gpt-5.3-codex-spark", + "input": "hi", + "store": false, + "event_id": "evt_upstream", + "stream": true, + "background": true, + "stream_options": {"include_usage": true} + }`) + + got, err := buildResponsesWSCreateEvent(payload, common.RawMessage(`false`)) + if err != nil { + t.Fatalf("buildResponsesWSCreateEvent() error = %v", err) + } + var data map[string]any + if err := common.Unmarshal(got, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if data["type"] != responsesWSEventTypeResponseCreate { + t.Fatalf("type = %#v", data["type"]) + } + if data["model"] != "gpt-5.3-codex-spark" || data["input"] != "hi" || data["store"] != false { + t.Fatalf("unexpected flat event fields: %s", got) + } + if data["generate"] != false { + t.Fatalf("generate = %#v, want false", data["generate"]) + } + for _, key := range []string{"response", "event_id", "stream", "background", "stream_options"} { + if _, ok := data[key]; ok { + t.Fatalf("field %q should not be present in upstream event: %s", key, got) + } + } +} + +func TestHTTPResponsesRequestDoesNotMarshalGenerate(t *testing.T) { + var req dto.OpenAIResponsesRequest + if err := common.Unmarshal([]byte(`{"model":"gpt-5.3-codex-spark","input":"hi","generate":false}`), &req); err != nil { + t.Fatalf("unmarshal request: %v", err) + } + got, err := common.Marshal(req) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + var data map[string]any + if err := common.Unmarshal(got, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if _, ok := data["generate"]; ok { + t.Fatalf("generate leaked into HTTP request JSON: %s", got) + } +} + +func TestBuildResponsesWSErrorPayloadIncludesStatus(t *testing.T) { + payload, err := buildResponsesWSErrorPayload("evt_err", types.NewErrorWithStatusCode( + errors.New("model is required"), + types.ErrorCodeInvalidRequest, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + )) + if err != nil { + t.Fatalf("buildResponsesWSErrorPayload() error = %v", err) + } + var data struct { + Type string `json:"type"` + Status int `json:"status"` + EventID string `json:"event_id"` + Error *types.OpenAIError `json:"error"` + } + if err := common.Unmarshal(payload, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if data.Type != "error" || data.Status != http.StatusBadRequest || data.EventID != "evt_err" { + t.Fatalf("unexpected error event: %s", payload) + } + if data.Error == nil || data.Error.Code != string(types.ErrorCodeInvalidRequest) { + t.Fatalf("unexpected error body: %#v", data.Error) + } +} + +func TestResponsesWSInvalidRequestErrorUsesBadRequestStatus(t *testing.T) { + payload, err := buildResponsesWSErrorPayload("", newResponsesWSInvalidRequestError(errors.New("bad event"))) + if err != nil { + t.Fatalf("buildResponsesWSErrorPayload() error = %v", err) + } + var data struct { + Status int `json:"status"` + } + if err := common.Unmarshal(payload, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if data.Status != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", data.Status, http.StatusBadRequest) + } +} + +func TestRemoveResponsesWSTransportFields(t *testing.T) { + payload := []byte(`{ + "model": "gpt-5.3-codex-spark", + "stream": true, + "background": true, + "stream_options": {"include_usage": true}, + "store": false + }`) + + got, err := removeResponsesWSTransportFields(payload) + if err != nil { + t.Fatalf("removeResponsesWSTransportFields() error = %v", err) + } + var data map[string]any + if err := common.Unmarshal(got, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + for _, key := range []string{"stream", "background", "stream_options"} { + if _, ok := data[key]; ok { + t.Fatalf("transport field %q still present in %s", key, got) + } + } + if data["store"] != false { + t.Fatalf("store = %#v, want false", data["store"]) + } +} + +func TestToWebSocketURL(t *testing.T) { + tests := map[string]string{ + "https://api.openai.com/v1/responses": "wss://api.openai.com/v1/responses", + "http://127.0.0.1:3000/v1/responses": "ws://127.0.0.1:3000/v1/responses", + "wss://chatgpt.com/backend-api/codex/responses": "wss://chatgpt.com/backend-api/codex/responses", + "ws://127.0.0.1:3000/backend-api/codex/responses": "ws://127.0.0.1:3000/backend-api/codex/responses", + } + + for input, want := range tests { + if got := toWebSocketURL(input); got != want { + t.Fatalf("toWebSocketURL(%q) = %q, want %q", input, got, want) + } + } +} + +func TestHandleTargetWriteFailureWithStateReleasesCurrentAndClearsTarget(t *testing.T) { + target, cleanup := newTestResponsesWSTarget(t) + defer cleanup() + + var committed *bool + session := &responsesWSSession{target: target} + state := &responsesWSCallState{ + info: &relaycommon.RelayInfo{}, + commitRate: func(success bool) { + committed = &success + }, + } + session.current = state + + apiErr := session.handleTargetWriteFailureWithState(state, errors.New("write failed")) + + if apiErr == nil { + t.Fatal("apiErr is nil") + } + if session.target != nil { + t.Fatal("target was not cleared") + } + if session.getCurrent() != nil { + t.Fatal("current response was not released") + } + if committed == nil || *committed { + t.Fatalf("commit success = %v, want false", committed) + } +} + +func TestHandleControlEventWriteFailureSendsResponsesError(t *testing.T) { + clientConn, serverConn, cleanupClient := newTestWebSocketPair(t) + defer cleanupClient() + target, cleanupTarget := newTestResponsesWSTarget(t) + defer cleanupTarget() + + session := &responsesWSSession{ + client: serverConn, + target: target, + } + apiErr := session.handleControlEventWriteFailure(errors.New("write failed")) + if apiErr != nil { + t.Fatalf("handleControlEventWriteFailure() error = %v", apiErr) + } + if session.target != nil { + t.Fatal("target was not cleared") + } + + if err := clientConn.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + _, payload, err := clientConn.ReadMessage() + if err != nil { + t.Fatalf("read responses error event: %v", err) + } + var data struct { + Type string `json:"type"` + Status int `json:"status"` + } + if err := common.Unmarshal(payload, &data); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if data.Type != "error" || data.Status == 0 { + t.Fatalf("unexpected error event: %s", payload) + } +} + +func TestObserveUpstreamFailedReleasesCurrent(t *testing.T) { + var committed *bool + session := &responsesWSSession{} + state := &responsesWSCallState{ + info: &relaycommon.RelayInfo{}, + commitRate: func(success bool) { + committed = &success + }, + } + session.current = state + + session.observeUpstreamMessage([]byte(`{"type":"response.failed"}`)) + + if session.getCurrent() != nil { + t.Fatal("current response was not released") + } + if committed == nil || *committed { + t.Fatalf("commit success = %v, want false", committed) + } +} + +func newTestResponsesWSTarget(t *testing.T) (*websocket.Conn, func()) { + t.Helper() + target, _, cleanup := newTestWebSocketPair(t) + return target, cleanup +} + +func newTestWebSocketPair(t *testing.T) (*websocket.Conn, *websocket.Conn, func()) { + t.Helper() + upgrader := websocket.Upgrader{} + serverConnCh := make(chan *websocket.Conn, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + serverConnCh <- conn + })) + + targetURL := "ws" + strings.TrimPrefix(server.URL, "http") + target, _, err := websocket.DefaultDialer.Dial(targetURL, nil) + if err != nil { + server.Close() + t.Fatalf("dial websocket: %v", err) + } + serverConn := <-serverConnCh + cleanup := func() { + _ = target.Close() + _ = serverConn.Close() + server.Close() + } + return target, serverConn, cleanup +} diff --git a/router/relay-router.go b/router/relay-router.go index 17a13cad7fd6..59487c7d6c64 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -71,6 +71,11 @@ func SetRelayRouter(router *gin.Engine) { relayV1Router.Use(middleware.SystemPerformanceCheck()) relayV1Router.Use(middleware.TokenAuth()) relayV1Router.Use(middleware.ModelRequestRateLimit()) + { + // Responses WebSocket route. Channel selection happens after the first + // response.create event because the model is in the WebSocket payload. + relayV1Router.GET("/responses", controller.ResponsesWebSocket) + } { // WebSocket 路由（统一到 Relay） wsRouter := relayV1Router.Group("") diff --git a/service/relay_error.go b/service/relay_error.go new file mode 100644 index 000000000000..74b7a3956547 --- /dev/null +++ b/service/relay_error.go @@ -0,0 +1,96 @@ +package service + +import ( + "fmt" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/types" + + "github.com/bytedance/gopkg/util/gopool" + "github.com/gin-gonic/gin" +) + +func ShouldRetryRelayError(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { + if openaiErr == nil { + return false + } + if ShouldSkipRetryAfterChannelAffinityFailure(c) { + return false + } + if c != nil { + if _, ok := c.Get("specific_channel_id"); ok { + return false + } + } + if types.IsChannelError(openaiErr) { + return true + } + if types.IsSkipRetryError(openaiErr) { + return false + } + if retryTimes <= 0 { + return false + } + code := openaiErr.StatusCode + if code >= 200 && code < 300 { + return false + } + if code < 100 || code > 599 { + return true + } + if operation_setting.IsAlwaysSkipRetryCode(openaiErr.GetErrorCode()) { + return false + } + return operation_setting.ShouldRetryByStatusCode(code) +} + +func ProcessChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) { + if err == nil { + return + } + logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.MaskSensitiveErrorWithStatusCode()))) + if ShouldDisableChannel(err) && channelError.AutoBan { + gopool.Go(func() { + DisableChannel(channelError, err.ErrorWithStatusCode()) + }) + } + + if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) { + userId := c.GetInt("id") + tokenName := c.GetString("token_name") + modelName := c.GetString("original_model") + tokenId := c.GetInt("token_id") + userGroup := c.GetString("group") + channelId := c.GetInt("channel_id") + other := make(map[string]interface{}) + if c.Request != nil && c.Request.URL != nil { + other["request_path"] = c.Request.URL.Path + } + other["error_type"] = err.GetErrorType() + other["error_code"] = err.GetErrorCode() + other["status_code"] = err.StatusCode + other["channel_id"] = channelId + other["channel_name"] = c.GetString("channel_name") + other["channel_type"] = c.GetInt("channel_type") + adminInfo := make(map[string]interface{}) + adminInfo["use_channel"] = c.GetStringSlice("use_channel") + isMultiKey := common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey) + if isMultiKey { + adminInfo["is_multi_key"] = true + adminInfo["multi_key_index"] = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex) + } + AppendChannelAffinityAdminInfo(c, adminInfo) + other["admin_info"] = adminInfo + startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime) + if startTime.IsZero() { + startTime = time.Now() + } + useTimeSeconds := int(time.Since(startTime).Seconds()) + model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other) + } +} diff --git a/service/relay_error_test.go b/service/relay_error_test.go new file mode 100644 index 000000000000..4ed854fa2927 --- /dev/null +++ b/service/relay_error_test.go @@ -0,0 +1,21 @@ +package service + +import ( + "errors" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" +) + +func TestShouldRetryRelayErrorSpecificChannelSkipsChannelError(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set("specific_channel_id", "1") + err := types.NewError(errors.New("channel failed"), types.ErrorCodeChannelNoAvailableKey) + + if ShouldRetryRelayError(c, err, 1) { + t.Fatal("specific channel channel error should not retry") + } +} diff --git a/service/responses_usage.go b/service/responses_usage.go new file mode 100644 index 000000000000..c241af466813 --- /dev/null +++ b/service/responses_usage.go @@ -0,0 +1,43 @@ +package service + +import "github.com/QuantumNous/new-api/dto" + +func ApplyResponsesUsage(dst *dto.Usage, src *dto.Usage) { + if dst == nil || src == nil { + return + } + if src.InputTokens != 0 { + dst.PromptTokens = src.InputTokens + dst.InputTokens = src.InputTokens + } + if src.OutputTokens != 0 { + dst.CompletionTokens = src.OutputTokens + dst.OutputTokens = src.OutputTokens + } + if src.TotalTokens != 0 { + dst.TotalTokens = src.TotalTokens + } + if src.InputTokensDetails != nil { + inputDetails := *src.InputTokensDetails + dst.InputTokensDetails = &inputDetails + dst.PromptTokensDetails = inputDetails + } + outputDetails := src.CompletionTokenDetails + if src.OutputTokensDetails != nil { + outputDetails = *src.OutputTokensDetails + } + if !isZeroOutputTokenDetails(outputDetails) { + dst.CompletionTokenDetails = outputDetails + dst.OutputTokensDetails = &outputDetails + } + dst.PromptCacheHitTokens = src.PromptCacheHitTokens + dst.UsageSemantic = src.UsageSemantic + dst.UsageSource = src.UsageSource +} + +func isZeroOutputTokenDetails(details dto.OutputTokenDetails) bool { + return details.TextTokens == 0 && + details.AudioTokens == 0 && + details.ImageTokens == 0 && + details.ReasoningTokens == 0 +} diff --git a/service/responses_usage_test.go b/service/responses_usage_test.go new file mode 100644 index 000000000000..9f7f7a5e961f --- /dev/null +++ b/service/responses_usage_test.go @@ -0,0 +1,87 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/dto" +) + +func TestApplyResponsesUsageCopiesTokenDetails(t *testing.T) { + dst := &dto.Usage{} + src := &dto.Usage{ + InputTokens: 11, + OutputTokens: 7, + TotalTokens: 18, + InputTokensDetails: &dto.InputTokenDetails{ + CachedTokens: 3, + CachedCreationTokens: 2, + TextTokens: 6, + AudioTokens: 4, + ImageTokens: 5, + }, + OutputTokensDetails: &dto.OutputTokenDetails{ + TextTokens: 1, + AudioTokens: 2, + ImageTokens: 3, + ReasoningTokens: 4, + }, + PromptCacheHitTokens: 3, + UsageSemantic: "openai", + UsageSource: "upstream", + } + + ApplyResponsesUsage(dst, src) + + if dst.PromptTokens != 11 || dst.CompletionTokens != 7 || dst.TotalTokens != 18 { + t.Fatalf("usage tokens = %#v", dst) + } + if dst.InputTokensDetails == nil { + t.Fatal("InputTokensDetails is nil") + } + if dst.PromptTokensDetails.CachedTokens != 3 || + dst.PromptTokensDetails.CachedCreationTokens != 2 || + dst.PromptTokensDetails.TextTokens != 6 || + dst.PromptTokensDetails.AudioTokens != 4 || + dst.PromptTokensDetails.ImageTokens != 5 { + t.Fatalf("prompt details = %#v", dst.PromptTokensDetails) + } + if dst.CompletionTokenDetails.TextTokens != 1 || + dst.CompletionTokenDetails.AudioTokens != 2 || + dst.CompletionTokenDetails.ImageTokens != 3 || + dst.CompletionTokenDetails.ReasoningTokens != 4 { + t.Fatalf("completion details = %#v", dst.CompletionTokenDetails) + } + if dst.OutputTokensDetails == nil { + t.Fatal("OutputTokensDetails is nil") + } + if dst.OutputTokensDetails.TextTokens != 1 || + dst.OutputTokensDetails.AudioTokens != 2 || + dst.OutputTokensDetails.ImageTokens != 3 || + dst.OutputTokensDetails.ReasoningTokens != 4 { + t.Fatalf("output details = %#v", dst.OutputTokensDetails) + } + if dst.UsageSemantic != "openai" || dst.UsageSource != "upstream" { + t.Fatalf("usage metadata = %#v", dst) + } +} + +func TestApplyResponsesUsageFallsBackToCompletionTokenDetails(t *testing.T) { + dst := &dto.Usage{} + src := &dto.Usage{ + CompletionTokenDetails: dto.OutputTokenDetails{ + ReasoningTokens: 9, + }, + } + + ApplyResponsesUsage(dst, src) + + if dst.CompletionTokenDetails.ReasoningTokens != 9 { + t.Fatalf("reasoning tokens = %d, want 9", dst.CompletionTokenDetails.ReasoningTokens) + } + if dst.OutputTokensDetails == nil { + t.Fatal("OutputTokensDetails is nil") + } + if dst.OutputTokensDetails.ReasoningTokens != 9 { + t.Fatalf("output reasoning tokens = %d, want 9", dst.OutputTokensDetails.ReasoningTokens) + } +}