Alpha - #1775
Conversation
WalkthroughPropagates request context into RelayErrorHandler across many relay handlers and tests, refactors pre-consume/return quota to use relayInfo.FinalPreConsumedQuota, adds ImageRequest.MarshalJSON to flatten extras, wraps upstream Do errors into NewError, updates NewError/NewOpenAIError behavior, and adjusts channel auto-ban goroutine gating and cached-creation-token quota math. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant Handler as Relay Handler (Text/Image/Embed...)
participant Upstream
participant Service as service.RelayErrorHandler
participant Logger
Client->>Handler: HTTP request
Handler->>Upstream: forward request
Upstream-->>Handler: HTTP response (non-OK)
Handler->>Service: RelayErrorHandler(ctx, httpResp, skipRetry)
Service->>Logger: logger.LogInfo(ctx, ...) [debug path]
Service-->>Handler: *types.NewAPIError
Handler-->>Client: return error response (status reset)
sequenceDiagram
autonumber
participant Controller
participant Gate as ShouldDisableChannel & AutoBan
participant Goroutine as background goroutine
participant Channel as ChannelService
Controller->>Gate: evaluate ShouldDisableChannel(channel) && channel.AutoBan
alt true
Controller->>Goroutine: spawn
Goroutine->>Channel: DisableChannel(channelError, msg)
else false
Controller-->>Controller: do not spawn goroutine
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Pre-merge checks (1 passed, 1 warning, 1 inconclusive)❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
Poem
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. ✨ Finishing touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
service/error.go (1)
83-91: Always close resp.Body and surface read errors.On io.ReadAll failure, the body isn’t closed and Err is empty. Defer the close and attach the read error.
-func RelayErrorHandler(ctx context.Context, 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) + defer CloseResponseBodyGracefully(resp) + responseBody, err := io.ReadAll(resp.Body) if err != nil { - return + newApiErr.Err = fmt.Errorf("read response body failed: %w", err) + return } - CloseResponseBodyGracefully(resp)
🧹 Nitpick comments (12)
types/error.go (1)
329-337: Fix debug log formatting in ErrOptionWithHideErrMsgUse %v for error and add newline; current %s is incorrect for error and lacks newline. Consider using common.SysLog for consistency.
- if common.DebugEnabled { - fmt.Printf("ErrOptionWithHideErrMsg: %s, origin error: %s", replaceStr, e.Err) - } + if common.DebugEnabled { + fmt.Printf("ErrOptionWithHideErrMsg: %s, origin error: %v\n", replaceStr, e.Err) + }relay/channel/api_request.go (3)
266-269: Wrap is good; also handle cancel/timeout and close bodies on error pathAvoid leaking resources and misclassifying user-canceled requests. Add explicit handling and close bodies when Do fails.
resp, err := client.Do(req) if err != nil { - return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed")) + _ = req.Body.Close() + _ = c.Request.Body.Close() + // classify client-side cancellations/timeouts to avoid auto-ban/retry + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, + types.ErrOptionWithHideErrMsg("client canceled or timed out"), + types.ErrOptionWithSkipRetry()) + } + return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, + types.ErrOptionWithHideErrMsg("upstream error: do request failed")) }
270-272: Return typed error for nil response for consistencyKeep error shaping consistent with other paths.
- if resp == nil { - return nil, errors.New("resp is nil") - } + if resp == nil { + return nil, types.NewError(errors.New("resp is nil"), + types.ErrorCodeDoRequestFailed, + types.ErrOptionWithHideErrMsg("upstream error: response is nil")) + }
67-71: Avoid double-wrapping typed errors from doRequestdoRequest already returns *types.NewAPIError; re-wrapping with fmt.Errorf loses shape at call sites.
- resp, err := doRequest(c, req, info) - if err != nil { - return nil, fmt.Errorf("do request failed: %w", err) - } + resp, err := doRequest(c, req, info) + if err != nil { + return nil, err + }(Apply in both DoApiRequest and DoFormRequest.)
Also applies to: 104-107
controller/channel-test.go (1)
238-243: Avoid double-wrapping RelayErrorHandler result in testsUse the returned *types.NewAPIError directly to preserve status/code.
- err := service.RelayErrorHandler(c.Request.Context(), httpResp, true) + err := service.RelayErrorHandler(c.Request.Context(), httpResp, true) return testResult{ context: c, - localErr: err, - newAPIError: types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError), + localErr: err, + newAPIError: err, }relay/claude_handler.go (1)
113-118: Treat all 2xx responses as success, not just 200Some providers legitimately return 201/204. Narrow 200-check can misclassify successes as errors.
Apply:
- if httpResp.StatusCode != http.StatusOK { + if httpResp.StatusCode/100 != 2 { newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false) // reset status code 重置状态码 service.ResetStatusCode(newAPIError, statusCodeMappingStr) return newAPIError }relay/image_handler.go (1)
93-98: Broaden success check to any 2xxPrevents false negatives for providers that return 201/204.
- if httpResp.StatusCode != http.StatusOK { + if httpResp.StatusCode/100 != 2 { newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false) // reset status code 重置状态码 service.ResetStatusCode(newAPIError, statusCodeMappingStr) return newAPIError }relay/audio_handler.go (1)
55-60: Use 2xx range for successAvoid misclassifying valid non-200 statuses.
- if httpResp.StatusCode != http.StatusOK { + if httpResp.StatusCode/100 != 2 { newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false) // reset status code 重置状态码 service.ResetStatusCode(newAPIError, statusCodeMappingStr) return newAPIError }relay/embedding_handler.go (2)
46-49: Use common.Marshal and correct error code for consistency.Other handlers use common.Marshal and return ErrorCodeJsonMarshalFailed on JSON serialization failures. Align this handler.
- jsonData, err := json.Marshal(convertedRequest) + jsonData, err := common.Marshal(convertedRequest) if err != nil { - return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + return types.NewError(err, types.ErrorCodeJsonMarshalFailed, types.ErrOptionWithSkipRetry()) }
3-16: Drop unused encoding/json import after switching to common.Marshal.import ( "bytes" - "encoding/json" "fmt" "net/http" "one-api/common"service/error.go (1)
98-102: Optional: downgrade log level under debug gate.LogInfo for expected error-path diagnostics may be noisy; consider LogDebug while DebugEnabled is true.
relay/compatible_handler.go (1)
292-318: Clamp baseTokens to non-negative to avoid undercharging when subtracted components exceed prompt tokens.Rare but possible with misreported details; clamp after all subtractions.
// 减去 Gemini audio tokens if !dAudioTokens.IsZero() { audioInputPrice = operation_setting.GetGeminiInputAudioPricePerMillionTokens(modelName) if audioInputPrice > 0 { // 重新计算 base tokens baseTokens = baseTokens.Sub(dAudioTokens) audioInputQuota = decimal.NewFromFloat(audioInputPrice).Div(decimal.NewFromInt(1000000)).Mul(dAudioTokens).Mul(dGroupRatio).Mul(dQuotaPerUnit) extraContent += fmt.Sprintf("Audio Input 花费 %s", audioInputQuota.String()) } } + // 防止出现负数 + if baseTokens.LessThan(decimal.Zero) { + baseTokens = decimal.Zero + } promptQuota := baseTokens.Add(cachedTokensWithRatio). Add(imageTokensWithRatio). Add(dCachedCreationTokensWithRatio)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (14)
controller/channel-test.go(1 hunks)controller/relay.go(1 hunks)dto/openai_image.go(1 hunks)relay/audio_handler.go(1 hunks)relay/channel/api_request.go(1 hunks)relay/claude_handler.go(1 hunks)relay/compatible_handler.go(6 hunks)relay/embedding_handler.go(1 hunks)relay/gemini_handler.go(2 hunks)relay/image_handler.go(2 hunks)relay/rerank_handler.go(1 hunks)relay/responses_handler.go(1 hunks)service/error.go(3 hunks)types/error.go(3 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-08-08T17:12:43.157Z
Learnt from: RedwindA
PR: QuantumNous/new-api#1537
File: relay/gemini_handler.go:330-342
Timestamp: 2025-08-08T17:12:43.157Z
Learning: In the new-api repository, the `GeminiEmbeddingHandler` function in `relay/gemini_handler.go` is designed specifically for native Gemini embedding requests and therefore does not require the `ConvertGeminiRequest` step that is used in the chat handler. The embedding requests are already in the native Gemini format and don't need conversion.
Applied to files:
relay/gemini_handler.gorelay/embedding_handler.go
🧬 Code graph analysis (14)
dto/openai_image.go (1)
common/json.go (2)
Marshal(20-22)Unmarshal(8-10)
relay/rerank_handler.go (1)
service/error.go (1)
RelayErrorHandler(83-112)
relay/channel/api_request.go (1)
types/error.go (3)
NewError(187-207)ErrorCodeDoRequestFailed(46-46)ErrOptionWithHideErrMsg(329-336)
relay/audio_handler.go (1)
service/error.go (1)
RelayErrorHandler(83-112)
controller/relay.go (1)
service/channel.go (2)
ShouldDisableChannel(46-96)DisableChannel(20-35)
service/error.go (2)
types/error.go (1)
NewAPIError(82-90)logger/logger.go (1)
LogInfo(55-57)
relay/claude_handler.go (1)
service/error.go (1)
RelayErrorHandler(83-112)
relay/gemini_handler.go (1)
service/error.go (1)
RelayErrorHandler(83-112)
controller/channel-test.go (1)
service/error.go (1)
RelayErrorHandler(83-112)
relay/embedding_handler.go (1)
service/error.go (1)
RelayErrorHandler(83-112)
relay/responses_handler.go (1)
service/error.go (1)
RelayErrorHandler(83-112)
types/error.go (2)
dto/error.go (1)
OpenAIError(5-10)common/constants.go (1)
DebugEnabled(70-70)
relay/image_handler.go (1)
service/error.go (1)
RelayErrorHandler(83-112)
relay/compatible_handler.go (2)
service/error.go (1)
RelayErrorHandler(83-112)types/price_data.go (1)
PriceData(11-21)
🔇 Additional comments (15)
controller/relay.go (1)
282-286: LGTM: Gate goroutine creation on the conditionMoving the condition outside avoids spawning unnecessary goroutines.
types/error.go (1)
187-195: LGTM: Preserve inner NewAPIError and apply optionsAvoids double-wrapping and keeps original metadata.
relay/rerank_handler.go (1)
84-88: Context-aware RelayErrorHandler — LGTMPassing request context aligns with the new signature and enables contextual logging.
relay/gemini_handler.go (1)
155-159: Context propagation to RelayErrorHandler — LGTMConsistent with service/error.go; no further changes needed here.
Also applies to: 252-255
relay/responses_handler.go (1)
85-89: Context-aware RelayErrorHandler — LGTMMatches updated signature; error remapping remains intact.
relay/claude_handler.go (2)
114-114: Good: context propagated into RelayErrorHandlerPassing c.Request.Context() enables contextual logging/tracing in service/error.go. Looks consistent with the new signature.
114-114: All RelayErrorHandler calls now pass a context parameter
Every invocation was found using c.Request.Context() as its first argument.relay/image_handler.go (2)
94-94: Good: context-aware error handlingRelayErrorHandler now receives request context; aligns with service/error.go logging.
122-124: Logging tweak LGTMIncluding 张数 improves observability for quota/audit logs.
relay/audio_handler.go (1)
56-56: Good: pass ctx to error handlerConsistent with the new RelayErrorHandler signature.
relay/embedding_handler.go (2)
61-65: Context propagation to RelayErrorHandler looks correct.Passing c.Request.Context() aligns with the new signature and enables contextual logging.
57-60: Guard the type assertion on resp to avoid a panic if an adaptor ever changes return type.Low risk today, but cheap to harden.
var httpResp *http.Response if resp != nil { - httpResp = resp.(*http.Response) + var ok bool + httpResp, ok = resp.(*http.Response) + if !ok { + return types.NewError(fmt.Errorf("unexpected response type %T", resp), types.ErrorCodeDoRequestFailed, types.ErrOptionWithSkipRetry()) + } if httpResp.StatusCode != http.StatusOK {relay/compatible_handler.go (3)
161-165: Context passed to RelayErrorHandler — good update.
198-225: Incorporating cachedCreationTokens and cacheCreationRatio is correct and aligns with PriceData.
410-413: Telemetry additions for cache_creation_ look good.*Helps observability of cached creation economics.
| service.DisableChannel(channelError, err.Error()) | ||
| } | ||
| }) | ||
| }) |
There was a problem hiding this comment.
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.
| 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 | ||
| } |
There was a problem hiding this comment.
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.
| 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.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
controller/relay.go (2)
282-286: Pass channel type to ShouldDisableChannel and mask disable reason.Use
channelError.ChannelType(not id), and avoid leaking secrets in disable notifications.- if service.ShouldDisableChannel(channelError.ChannelId, err) && channelError.AutoBan { - gopool.Go(func() { - service.DisableChannel(channelError, err.Error()) - }) - } + if service.ShouldDisableChannel(channelError.ChannelType, err) && channelError.AutoBan { + reason := err.MaskSensitiveError() + gopool.Go(func() { + service.DisableChannel(channelError, reason) + }) + }
279-279: Mask sensitive details in error logs.
err.Error()may include secrets; prefer the masked message.- logger.LogError(c, fmt.Sprintf("relay error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, err.Error())) + logger.LogError(c, fmt.Sprintf("relay error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, err.MaskSensitiveError()))
🧹 Nitpick comments (1)
service/pre_consume_quota.go (1)
16-28: Avoid capturing mutable state in the goroutine.Copy
relayInfobefore spawning to eliminate any race with later mutations.- gopool.Go(func() { - relayInfoCopy := *relayInfo - - err := PostConsumeQuota(&relayInfoCopy, -relayInfo.FinalPreConsumedQuota, 0, false) - if err != nil { - common.SysLog("error return pre-consumed quota: " + err.Error()) - } - }) + relayInfoCopy := *relayInfo + gopool.Go(func() { + if err := PostConsumeQuota(&relayInfoCopy, -relayInfo.FinalPreConsumedQuota, 0, false); err != nil { + common.SysLog("error return pre-consumed quota: " + err.Error()) + } + })
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
controller/relay.go(2 hunks)service/pre_consume_quota.go(3 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
service/pre_consume_quota.go (5)
relay/common/relay_info.go (1)
RelayInfo(74-120)logger/logger.go (2)
LogInfo(55-57)FormatQuota(102-108)service/quota.go (1)
PostConsumeQuota(490-519)types/error.go (5)
NewAPIError(82-90)NewError(187-207)ErrOptionWithSkipRetry(317-321)NewErrorWithStatusCode(242-258)ErrOptionWithNoRecordErrorLog(323-327)model/user.go (2)
GetUserQuota(656-681)DecreaseUserQuota(777-792)
controller/relay.go (2)
service/pre_consume_quota.go (2)
PreConsumeQuota(32-78)ReturnPreConsumedQuota(16-28)service/channel.go (2)
ShouldDisableChannel(46-96)DisableChannel(20-35)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Push Docker image to multiple registries
🔇 Additional comments (2)
controller/relay.go (2)
142-146: API alignment looks good.Updated call to
service.PreConsumeQuotamatches the new signature and flow.
147-152: Quota return guard is correct.Double-check is fine;
ReturnPreConsumedQuotaalso guards internally. No action needed.
| 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 |
There was a problem hiding this comment.
🛠️ 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.
| 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.
3 participants
Summary by CodeRabbit
New Features
Refactor