Skip to content
Merged

Alpha #1775

Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion controller/channel-test.go
Original file line number Diff line number Diff line change
Expand Up @@ -235,7 +235,7 @@ func testChannel(channel *model.Channel, testModel string) testResult {
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp.StatusCode != http.StatusOK {
err := service.RelayErrorHandler(httpResp, true)
err := service.RelayErrorHandler(c.Request.Context(), httpResp, true)
return testResult{
context: c,
localErr: err,
Expand Down
19 changes: 9 additions & 10 deletions controller/relay.go
Original file line number Diff line number Diff line change
Expand Up @@ -139,15 +139,15 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {

// common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)

preConsumedQuota, newAPIError := service.PreConsumeQuota(c, priceData.ShouldPreConsumedQuota, relayInfo)
newAPIError = service.PreConsumeQuota(c, priceData.ShouldPreConsumedQuota, relayInfo)
if newAPIError != nil {
return
}

defer func() {
// Only return quota if downstream failed and quota was actually pre-consumed
if newAPIError != nil && preConsumedQuota != 0 {
service.ReturnPreConsumedQuota(c, relayInfo, preConsumedQuota)
if newAPIError != nil && relayInfo.FinalPreConsumedQuota != 0 {
service.ReturnPreConsumedQuota(c, relayInfo)
}
}()

Expand Down Expand Up @@ -277,14 +277,13 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b

func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) {
logger.LogError(c, fmt.Sprintf("relay error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, err.Error()))

gopool.Go(func() {
// 不要使用context获取渠道信息,异步处理时可能会出现渠道信息不一致的情况
// do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
if service.ShouldDisableChannel(channelError.ChannelId, err) && channelError.AutoBan {
// 不要使用context获取渠道信息,异步处理时可能会出现渠道信息不一致的情况
// do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
if service.ShouldDisableChannel(channelError.ChannelId, err) && channelError.AutoBan {
gopool.Go(func() {
service.DisableChannel(channelError, err.Error())
}
})
})
Comment on lines 284 to +285

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Mask sensitive error details when disabling a channel

Reason text is logged and notified; err.Error() may leak secrets. Use the masked message.

-   service.DisableChannel(channelError, err.Error())
+   service.DisableChannel(channelError, err.MaskSensitiveError())
🤖 Prompt for AI Agents
In controller/relay.go around lines 284-285, the call
service.DisableChannel(channelError, err.Error()) exposes sensitive data;
replace err.Error() with a redacted/masked reason string. Use the project’s
existing error-masking utility (e.g., maskError(err) or sanitizeError(err)) to
produce a non-sensitive message and pass that to service.DisableChannel; if no
utility exists, create a small redact function that strips tokens/PII and use it
here so notifications and logs never include raw err.Error() contents.

}
Comment thread
creamlike1024 marked this conversation as resolved.

if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) {
// 保存错误日志到mysql中
Expand Down
25 changes: 25 additions & 0 deletions dto/openai_image.go
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,31 @@ func (i *ImageRequest) UnmarshalJSON(data []byte) error {
return nil
}

// 序列化时需要重新把字段平铺
func (r ImageRequest) MarshalJSON() ([]byte, error) {
// 将已定义字段转为 map
type Alias ImageRequest
alias := Alias(r)
base, err := common.Marshal(alias)
if err != nil {
return nil, err
}

var baseMap map[string]json.RawMessage
if err := common.Unmarshal(base, &baseMap); err != nil {
return nil, err
}

// 合并 ExtraFields
for k, v := range r.Extra {
if _, exists := baseMap[k]; !exists {
baseMap[k] = v
}
}

return json.Marshal(baseMap)
}

func GetJSONFieldNames(t reflect.Type) map[string]struct{} {
fields := make(map[string]struct{})
for i := 0; i < t.NumField(); i++ {
Expand Down
2 changes: 1 addition & 1 deletion relay/audio_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ func AudioHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(httpResp, false)
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
Expand Down
3 changes: 1 addition & 2 deletions relay/channel/api_request.go
Original file line number Diff line number Diff line change
Expand Up @@ -264,9 +264,8 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http
}

resp, err := client.Do(req)

if err != nil {
return nil, err
return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed"))
}
if resp == nil {
return nil, errors.New("resp is nil")
Expand Down
2 changes: 1 addition & 1 deletion relay/claude_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
httpResp = resp.(*http.Response)
info.IsStream = info.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(httpResp, false)
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
Expand Down
20 changes: 18 additions & 2 deletions relay/compatible_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
httpResp = resp.(*http.Response)
info.IsStream = info.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
if httpResp.StatusCode != http.StatusOK {
newApiErr := service.RelayErrorHandler(httpResp, false)
newApiErr := service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
return newApiErr
Expand Down Expand Up @@ -195,6 +195,8 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
imageTokens := usage.PromptTokensDetails.ImageTokens
audioTokens := usage.PromptTokensDetails.AudioTokens
completionTokens := usage.CompletionTokens
cachedCreationTokens := usage.PromptTokensDetails.CachedCreationTokens

modelName := relayInfo.OriginModelName

tokenName := ctx.GetString("token_name")
Expand All @@ -204,19 +206,22 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
modelRatio := relayInfo.PriceData.ModelRatio
groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
modelPrice := relayInfo.PriceData.ModelPrice
cachedCreationRatio := relayInfo.PriceData.CacheCreationRatio

// Convert values to decimal for precise calculation
dPromptTokens := decimal.NewFromInt(int64(promptTokens))
dCacheTokens := decimal.NewFromInt(int64(cacheTokens))
dImageTokens := decimal.NewFromInt(int64(imageTokens))
dAudioTokens := decimal.NewFromInt(int64(audioTokens))
dCompletionTokens := decimal.NewFromInt(int64(completionTokens))
dCachedCreationTokens := decimal.NewFromInt(int64(cachedCreationTokens))
dCompletionRatio := decimal.NewFromFloat(completionRatio)
dCacheRatio := decimal.NewFromFloat(cacheRatio)
dImageRatio := decimal.NewFromFloat(imageRatio)
dModelRatio := decimal.NewFromFloat(modelRatio)
dGroupRatio := decimal.NewFromFloat(groupRatio)
dModelPrice := decimal.NewFromFloat(modelPrice)
dCachedCreationRatio := decimal.NewFromFloat(cachedCreationRatio)
dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit)

ratio := dModelRatio.Mul(dGroupRatio)
Expand Down Expand Up @@ -284,6 +289,11 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
baseTokens = baseTokens.Sub(dCacheTokens)
cachedTokensWithRatio = dCacheTokens.Mul(dCacheRatio)
}
var dCachedCreationTokensWithRatio decimal.Decimal
if !dCachedCreationTokens.IsZero() {
baseTokens = baseTokens.Sub(dCachedCreationTokens)
dCachedCreationTokensWithRatio = dCachedCreationTokens.Mul(dCachedCreationRatio)
}

// 减去 image tokens
var imageTokensWithRatio decimal.Decimal
Expand All @@ -302,7 +312,9 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
extraContent += fmt.Sprintf("Audio Input 花费 %s", audioInputQuota.String())
}
}
promptQuota := baseTokens.Add(cachedTokensWithRatio).Add(imageTokensWithRatio)
promptQuota := baseTokens.Add(cachedTokensWithRatio).
Add(imageTokensWithRatio).
Add(dCachedCreationTokensWithRatio)

completionQuota := dCompletionTokens.Mul(dCompletionRatio)

Expand Down Expand Up @@ -395,6 +407,10 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage
other["image_ratio"] = imageRatio
other["image_output"] = imageTokens
}
if cachedCreationTokens != 0 {
other["cache_creation_tokens"] = cachedCreationTokens
other["cache_creation_ratio"] = cachedCreationRatio
}
if !dWebSearchQuota.IsZero() {
if relayInfo.ResponsesUsageInfo != nil {
if webSearchTool, exists := relayInfo.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview]; exists {
Expand Down
2 changes: 1 addition & 1 deletion relay/embedding_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(httpResp, false)
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
Expand Down
4 changes: 2 additions & 2 deletions relay/gemini_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
httpResp = resp.(*http.Response)
info.IsStream = info.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(httpResp, false)
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
Expand Down Expand Up @@ -249,7 +249,7 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(httpResp, false)
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
}
Expand Down
4 changes: 2 additions & 2 deletions relay/image_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
httpResp = resp.(*http.Response)
info.IsStream = info.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(httpResp, false)
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
Expand Down Expand Up @@ -120,7 +120,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
var logContent string

if len(request.Size) > 0 {
logContent = fmt.Sprintf("大小 %s, 品质 %s", request.Size, quality)
logContent = fmt.Sprintf("大小 %s, 品质 %s, 张数 %d", request.Size, quality, request.N)
}

postConsumeQuota(c, info, usage.(*dto.Usage), logContent)
Expand Down
2 changes: 1 addition & 1 deletion relay/rerank_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(httpResp, false)
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
Expand Down
2 changes: 1 addition & 1 deletion relay/responses_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
httpResp = resp.(*http.Response)

if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(httpResp, false)
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
return newAPIError
Expand Down
6 changes: 4 additions & 2 deletions service/error.go
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
package service

import (
"context"
"errors"
"fmt"
"io"
"net/http"
"one-api/common"
"one-api/dto"
"one-api/logger"
"one-api/types"
"strconv"
"strings"
Expand Down Expand Up @@ -78,7 +80,7 @@ func ClaudeErrorWrapperLocal(err error, code string, statusCode int) *dto.Claude
return claudeErr
}

func RelayErrorHandler(resp *http.Response, showBodyWhenFail bool) (newApiErr *types.NewAPIError) {
func RelayErrorHandler(ctx context.Context, resp *http.Response, showBodyWhenFail bool) (newApiErr *types.NewAPIError) {
newApiErr = types.InitOpenAIError(types.ErrorCodeBadResponseStatusCode, resp.StatusCode)

responseBody, err := io.ReadAll(resp.Body)
Expand All @@ -94,7 +96,7 @@ func RelayErrorHandler(resp *http.Response, showBodyWhenFail bool) (newApiErr *t
newApiErr.Err = fmt.Errorf("bad response status code %d, body: %s", resp.StatusCode, string(responseBody))
} else {
if common.DebugEnabled {
println(fmt.Sprintf("bad response status code %d, body: %s", resp.StatusCode, string(responseBody)))
logger.LogInfo(ctx, fmt.Sprintf("bad response status code %d, body: %s", resp.StatusCode, string(responseBody)))
}
newApiErr.Err = fmt.Errorf("bad response status code %d", resp.StatusCode)
}
Expand Down
22 changes: 11 additions & 11 deletions service/pre_consume_quota.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,13 @@ import (
"github.com/gin-gonic/gin"
)

func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) {
if preConsumedQuota != 0 {
logger.LogInfo(c, fmt.Sprintf("用户 %d 请求失败, 返还预扣费额度 %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota)))
func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo) {
if relayInfo.FinalPreConsumedQuota != 0 {
logger.LogInfo(c, fmt.Sprintf("用户 %d 请求失败, 返还预扣费额度 %s", relayInfo.UserId, logger.FormatQuota(relayInfo.FinalPreConsumedQuota)))
gopool.Go(func() {
relayInfoCopy := *relayInfo

err := PostConsumeQuota(&relayInfoCopy, -preConsumedQuota, 0, false)
err := PostConsumeQuota(&relayInfoCopy, -relayInfo.FinalPreConsumedQuota, 0, false)
if err != nil {
common.SysLog("error return pre-consumed quota: " + err.Error())
}
Expand All @@ -29,16 +29,16 @@ func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, pr

// PreConsumeQuota checks if the user has enough quota to pre-consume.
// It returns the pre-consumed quota if successful, or an error if not.
func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) (int, *types.NewAPIError) {
func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.NewAPIError {
userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
if err != nil {
return 0, types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
return types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
}
if userQuota <= 0 {
return 0, types.NewErrorWithStatusCode(fmt.Errorf("用户额度不足, 剩余额度: %s", logger.FormatQuota(userQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
return types.NewErrorWithStatusCode(fmt.Errorf("用户额度不足, 剩余额度: %s", logger.FormatQuota(userQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
if userQuota-preConsumedQuota < 0 {
return 0, types.NewErrorWithStatusCode(fmt.Errorf("预扣费额度失败, 用户剩余额度: %s, 需要预扣费额度: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
return types.NewErrorWithStatusCode(fmt.Errorf("预扣费额度失败, 用户剩余额度: %s, 需要预扣费额度: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}

trustQuota := common.GetTrustQuota()
Expand All @@ -65,14 +65,14 @@ func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo
if preConsumedQuota > 0 {
err := PreConsumeTokenQuota(relayInfo, preConsumedQuota)
if err != nil {
return 0, types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
err = model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota)
if err != nil {
return 0, types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
}
logger.LogInfo(c, fmt.Sprintf("用户 %d 预扣费 %s, 预扣费后剩余额度: %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota), logger.FormatQuota(userQuota-preConsumedQuota)))
}
relayInfo.FinalPreConsumedQuota = preConsumedQuota
return preConsumedQuota, nil
return nil
Comment on lines 65 to +77

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Prevent quota inconsistency: roll back on partial pre-consume failure (or reorder ops).

If PreConsumeTokenQuota succeeds and DecreaseUserQuota fails, token quota remains pre-consumed with no rollback (the outer defer to ReturnPreConsumedQuota isn’t set yet). Reorder to decrease user quota first and roll back on token pre-consume failure.

Apply:

- if preConsumedQuota > 0 {
-   err := PreConsumeTokenQuota(relayInfo, preConsumedQuota)
-   if err != nil {
-     return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
-   }
-   err = model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota)
-   if err != nil {
-     return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
-   }
-   logger.LogInfo(c, fmt.Sprintf("用户 %d 预扣费 %s, 预扣费后剩余额度: %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota), logger.FormatQuota(userQuota-preConsumedQuota)))
- }
- relayInfo.FinalPreConsumedQuota = preConsumedQuota
+ if preConsumedQuota > 0 {
+   // 先扣用户额度;若令牌预扣失败,则回滚用户额度
+   if err := model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota); err != nil {
+     return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
+   }
+   if err := PreConsumeTokenQuota(relayInfo, preConsumedQuota); err != nil {
+     // 回滚用户额度(最佳努力)
+     if rerr := model.IncreaseUserQuota(relayInfo.UserId, preConsumedQuota, false); rerr != nil {
+       common.SysLog("rollback user quota failed after token pre-consume error: " + rerr.Error())
+     }
+     return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
+   }
+   relayInfo.FinalPreConsumedQuota = preConsumedQuota
+   logger.LogInfo(c, fmt.Sprintf("用户 %d 预扣费 %s, 预扣费后剩余额度: %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota), logger.FormatQuota(userQuota-preConsumedQuota)))
+ } else {
+   relayInfo.FinalPreConsumedQuota = 0
+ }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if preConsumedQuota > 0 {
err := PreConsumeTokenQuota(relayInfo, preConsumedQuota)
if err != nil {
return 0, types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
err = model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota)
if err != nil {
return 0, types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
}
logger.LogInfo(c, fmt.Sprintf("用户 %d 预扣费 %s, 预扣费后剩余额度: %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota), logger.FormatQuota(userQuota-preConsumedQuota)))
}
relayInfo.FinalPreConsumedQuota = preConsumedQuota
return preConsumedQuota, nil
return nil
if preConsumedQuota > 0 {
// 先扣用户额度;若令牌预扣失败,则回滚用户额度
if err := model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota); err != nil {
return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
}
if err := PreConsumeTokenQuota(relayInfo, preConsumedQuota); err != nil {
// 回滚用户额度(最佳努力)
if rerr := model.IncreaseUserQuota(relayInfo.UserId, preConsumedQuota, false); rerr != nil {
common.SysLog("rollback user quota failed after token pre-consume error: " + rerr.Error())
}
return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
relayInfo.FinalPreConsumedQuota = preConsumedQuota
logger.LogInfo(c, fmt.Sprintf("用户 %d 预扣费 %s, 预扣费后剩余额度: %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota), logger.FormatQuota(userQuota-preConsumedQuota)))
} else {
relayInfo.FinalPreConsumedQuota = 0
}
return nil
🤖 Prompt for AI Agents
In service/pre_consume_quota.go around lines 65 to 77, the current sequence
pre-consumes token quota first then decreases the user's stored quota, which
causes quota inconsistency if DecreaseUserQuota fails; either reorder to call
model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota) first and only call
PreConsumeTokenQuota(relayInfo, preConsumedQuota) after that succeeds (and only
log and set relayInfo.FinalPreConsumedQuota after both succeed), or keep the
current order but add a rollback on DecreaseUserQuota failure by calling
ReturnPreConsumedQuota(relayInfo, preConsumedQuota) (handling and logging any
rollback error) before returning the wrapped error; ensure error wrapping and
status codes remain the same and that FinalPreConsumedQuota is only set when the
net state is consistent.

}
34 changes: 32 additions & 2 deletions types/error.go
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,14 @@ func (e *NewAPIError) ToClaudeError() ClaudeError {
type NewAPIErrorOptions func(*NewAPIError)

func NewError(err error, errorCode ErrorCode, ops ...NewAPIErrorOptions) *NewAPIError {
var newErr *NewAPIError
// 保留深层传递的 new err
if errors.As(err, &newErr) {
for _, op := range ops {
op(newErr)
}
return newErr
}
e := &NewAPIError{
Err: err,
RelayError: nil,
Expand All @@ -199,8 +207,21 @@ func NewError(err error, errorCode ErrorCode, ops ...NewAPIErrorOptions) *NewAPI
}

func NewOpenAIError(err error, errorCode ErrorCode, statusCode int, ops ...NewAPIErrorOptions) *NewAPIError {
if errorCode == ErrorCodeDoRequestFailed {
err = errors.New("upstream error: do request failed")
var newErr *NewAPIError
// 保留深层传递的 new err
if errors.As(err, &newErr) {
if newErr.RelayError == nil {
openaiError := OpenAIError{
Message: newErr.Error(),
Type: string(errorCode),
Code: errorCode,
}
newErr.RelayError = openaiError
}
for _, op := range ops {
op(newErr)
}
return newErr
}
Comment on lines 209 to 225

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Fix: NewOpenAIError should set type, status, and code on existing NewAPIError

Without these, responses may carry wrong error type/status and ToOpenAIError won’t use RelayError.

   if errors.As(err, &newErr) {
     if newErr.RelayError == nil {
       openaiError := OpenAIError{
         Message: newErr.Error(),
         Type:    string(errorCode),
         Code:    errorCode,
       }
       newErr.RelayError = openaiError
     }
+    // ensure correct classification and HTTP status on passthrough
+    newErr.errorType = ErrorTypeOpenAIError
+    newErr.StatusCode = statusCode
+    newErr.errorCode = errorCode
     for _, op := range ops {
       op(newErr)
     }
     return newErr
   }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
func NewOpenAIError(err error, errorCode ErrorCode, statusCode int, ops ...NewAPIErrorOptions) *NewAPIError {
if errorCode == ErrorCodeDoRequestFailed {
err = errors.New("upstream error: do request failed")
var newErr *NewAPIError
// 保留深层传递的 new err
if errors.As(err, &newErr) {
if newErr.RelayError == nil {
openaiError := OpenAIError{
Message: newErr.Error(),
Type: string(errorCode),
Code: errorCode,
}
newErr.RelayError = openaiError
}
for _, op := range ops {
op(newErr)
}
return newErr
}
func NewOpenAIError(err error, errorCode ErrorCode, statusCode int, ops ...NewAPIErrorOptions) *NewAPIError {
var newErr *NewAPIError
// 保留深层传递的 new err
if errors.As(err, &newErr) {
if newErr.RelayError == nil {
openaiError := OpenAIError{
Message: newErr.Error(),
Type: string(errorCode),
Code: errorCode,
}
newErr.RelayError = openaiError
}
// ensure correct classification and HTTP status on passthrough
newErr.errorType = ErrorTypeOpenAIError
newErr.StatusCode = statusCode
newErr.errorCode = errorCode
for _, op := range ops {
op(newErr)
}
return newErr
}
// ... rest of NewOpenAIError for non-NewAPIError cases ...
}
🤖 Prompt for AI Agents
In types/error.go around lines 209 to 225, when NewOpenAIError detects an
existing *NewAPIError via errors.As it currently only sets RelayError when nil
but does not update the NewAPIError's Type, Status, and Code fields; update the
function to assign newErr.Type = string(errorCode), newErr.Status = statusCode,
and newErr.Code = errorCode (and still ensure RelayError is set if nil), then
apply the ops and return newErr so ToOpenAIError and response handling see the
correct type/status/code.

openaiError := OpenAIError{
Message: err.Error(),
Expand Down Expand Up @@ -305,6 +326,15 @@ func ErrOptionWithNoRecordErrorLog() NewAPIErrorOptions {
}
}

func ErrOptionWithHideErrMsg(replaceStr string) NewAPIErrorOptions {
return func(e *NewAPIError) {
if common.DebugEnabled {
fmt.Printf("ErrOptionWithHideErrMsg: %s, origin error: %s", replaceStr, e.Err)
}
e.Err = errors.New(replaceStr)
}
}

func IsRecordErrorLog(e *NewAPIError) bool {
if e == nil {
return false
Expand Down