refactor: Introduce pre-consume quota and unify relay handlers - #1594
Conversation
This commit introduces a major architectural refactoring to improve quota management, centralize logging, and streamline the relay handling logic. Key changes: - **Pre-consume Quota:** Implements a new mechanism to check and reserve user quota *before* making the request to the upstream provider. This ensures more accurate quota deduction and prevents users from exceeding their limits due to concurrent requests. - **Unified Relay Handlers:** Refactors the relay logic to use generic handlers (e.g., `ChatHandler`, `ImageHandler`) instead of provider-specific implementations. This significantly reduces code duplication and simplifies adding new channels. - **Centralized Logger:** A new dedicated `logger` package is introduced, and all system logging calls are migrated to use it, moving this responsibility out of the `common` package. - **Code Reorganization:** DTOs are generalized (e.g., `dalle.go` -> `openai_image.go`) and utility code is moved to more appropriate packages (e.g., `common/http.go` -> `service/http.go`) for better code structure.
This commit refactors the logging mechanism across the application by replacing direct logger calls with a centralized logging approach using the `common` package. Key changes include: - Replaced instances of `logger.SysLog` and `logger.FatalLog` with `common.SysLog` and `common.FatalLog` for consistent logging practices. - Updated resource initialization error handling to utilize the new logging structure, enhancing maintainability and readability. - Minor adjustments to improve code clarity and organization throughout various modules. This change aims to streamline logging and improve the overall architecture of the codebase.
|
Caution Review failedThe pull request is closed. WalkthroughRefactors relay into a format-driven engine with a typed RelayInfo/ChannelMeta; adds token-count metadata across DTOs; centralizes pricing and pre-consume quota logic (new pre_consume_quota service and common.GetTrustQuota); migrates logging to a new logger package; switches many adaptors to ChannelBaseUrl; and updates many handlers/controllers to operate on RelayInfo. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Controller as Relay Controller
participant Gen as GenRelayInfo
participant Info as RelayInfo
participant Helper as ModelMappedHelper
participant Price as ModelPriceHelper
participant PreCons as PreConsumeQuota
participant Adaptor
participant Upstream
participant Service as PostConsumeQuota
Client->>Controller: request (format, payload)
Controller->>Gen: GenRelayInfo(format, parsedRequest, ws?)
Gen-->>Info: RelayInfo (with ChannelMeta, Request)
Controller->>Helper: ModelMappedHelper(info, request)
Helper-->>Info: UpstreamModelName set
Controller->>Price: ModelPriceHelper(c, info, promptTokens, meta)
Price-->>Info: PriceData
Controller->>PreCons: PreConsumeQuota(c, preConsume, info)
PreCons-->>Info: FinalPreConsumedQuota
Controller->>Adaptor: adaptor.Init(info) / adaptor.ConvertRequest(info, request)
Adaptor->>Upstream: DoRequest (HTTP/WSS)
Upstream-->>Adaptor: Response
Adaptor-->>Controller: DoResponse -> usage/data
Controller->>Service: PostConsumeQuota(c, info, usage)
Service-->>Controller: updated quotas/logs
Controller-->>Client: stream or JSON response
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120+ minutes Possibly related PRs
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ 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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 48
🔭 Outside diff range comments (36)
relay/common_handler/rerank.go (2)
18-23: Close the response body on all paths; use defer to avoid leaksCurrently, the body is only closed after a successful ReadAll. If io.ReadAll fails (Line 19), the function returns without closing, leaking the connection.
Apply this diff to ensure the body is closed regardless of the read outcome:
func RerankHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { - responseBody, err := io.ReadAll(resp.Body) + defer service.CloseResponseBodyGracefully(resp) + responseBody, err := io.ReadAll(resp.Body) if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) } - service.CloseResponseBodyGracefully(resp)
39-47: Guard against out-of-bounds access when backfilling empty document stringsAccessing info.Documents[result.Index] (Line 44) will panic if index is out of range. Add a bounds check before indexing.
Apply this diff:
- if doc, ok := result.Document.(string); ok { - if doc == "" { - document = info.Documents[result.Index] - } else { - document = doc - } - } else { + if doc, ok := result.Document.(string); ok { + if doc == "" { + if result.Index >= 0 && result.Index < len(info.Documents) { + document = info.Documents[result.Index] + } else { + // index out of range; keep empty document or handle as needed + } + } else { + document = doc + } + } else { document = result.Document }controller/oidc.go (3)
64-74: Harden token exchange error handling; include status/body and keep logs structuredCurrently you decode regardless of HTTP status and only test for an empty access_token, which loses actionable context. Check HTTP status first and log a short, non-sensitive snippet of the error body.
Apply this diff:
- defer res.Body.Close() - var oidcResponse OidcResponse - err = json.NewDecoder(res.Body).Decode(&oidcResponse) - if err != nil { - return nil, err - } - - if oidcResponse.AccessToken == "" { - common.SysLog("OIDC 获取 Token 失败，请检查设置！") - return nil, errors.New("OIDC 获取 Token 失败，请检查设置！") - } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + b, _ := io.ReadAll(io.LimitReader(res.Body, 512)) + common.SysLog(fmt.Sprintf("OIDC 获取 Token 失败(status=%d): %s", res.StatusCode, strings.TrimSpace(string(b)))) + return nil, errors.New("OIDC 获取 Token 失败，请检查设置！") + } + var oidcResponse OidcResponse + if err := json.NewDecoder(res.Body).Decode(&oidcResponse); err != nil { + return nil, err + } + if oidcResponse.AccessToken == "" { + common.SysLog("OIDC 获取 Token 响应缺少 access_token") + return nil, errors.New("OIDC 获取 Token 失败，请检查设置！") + }Add missing import:
// add to the imports "io"
106-113: Guard against panic on session state type assertion
session.Get("oauth_state").(string)will panic if the session value is missing or not a string. Use a safe assertion.- state := c.Query("state") - if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) { + state := c.Query("state") + stateInSessionIface := session.Get("oauth_state") + stateInSession, ok := stateInSessionIface.(string) + if state == "" || !ok || stateInSession == "" || state != stateInSession { c.JSON(http.StatusForbidden, gin.H{ "success": false, "message": "state is empty or not same", }) return }
209-213: Avoid potential panic onid.(int); validate and handle missing/invalid session idDirect type assertion can panic if the stored type differs (or is missing). Validate before use.
- id := session.Get("id") - // id := c.GetInt("id") // critical bug! - user.Id = id.(int) + sid := session.Get("id") + idInt, ok := sid.(int) + if !ok || idInt <= 0 { + common.ApiError(c, errors.New("invalid session id")) + return + } + user.Id = idIntrelay/channel/xunfei/relay-xunfei.go (2)
205-208: Critical: possible nil pointer panic when websocket handshake fails (status != 101).If
err == nilbutresp.StatusCode != 101, this returnsnil, nil, nil. The caller then dereferencesconn(nil) at WriteJSON, causing a panic. Also,resp.Bodyisn’t closed in the non-101 case.- if err != nil || resp.StatusCode != 101 { - return nil, nil, err - } + if err != nil { + return nil, nil, err + } + if resp == nil || resp.StatusCode != 101 { + if conn != nil { + _ = conn.Close() + } + var statusErr error + if resp != nil { + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + statusErr = fmt.Errorf("websocket handshake failed: status=%d, body=%s", resp.StatusCode, string(body)) + } else { + statusErr = fmt.Errorf("websocket handshake failed: nil response") + } + return nil, nil, statusErr + }
215-217: Prevent potential goroutine block on stop signaling.
stopChanis unbuffered. If the consumer stops reading (client disconnects, context canceled), thestopChan <- truesend after the loop can block indefinitely, leaking the goroutine.Change declaration to buffered:
stopChan := make(chan bool, 1)No other changes are required since you already send only a single value.
relay/channel/siliconflow/relay-siliconflow.go (1)
16-21: Always close the upstream response body even on read errorsCurrently, the body is only closed after a successful ReadAll, leaking the connection on early-return paths and hurting keep-alive reuse. Defer the close immediately and remove the later explicit close.
Apply this diff:
func siliconflowRerankHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { - responseBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) - } - service.CloseResponseBodyGracefully(resp) + defer service.CloseResponseBodyGracefully(resp) + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) + }relay/channel/coze/relay-coze.go (3)
48-52: Defer response body close to avoid leak on read errorIf
io.ReadAllfails,resp.Bodyis never closed. Defer the close before reading to guarantee cleanup on all paths.func cozeChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { - responseBody, err := io.ReadAll(resp.Body) + defer service.CloseResponseBodyGracefully(resp) + responseBody, err := io.ReadAll(resp.Body) if err != nil { return nil, types.NewError(err, types.ErrorCodeBadResponseBody) } - service.CloseResponseBodyGracefully(resp)
98-101: Close streaming resp.Body and increase scanner buffer to avoid token-size limit
resp.Bodyis not closed in the stream handler, leaking resources. Also,bufio.Scannerdefaults to a 64KB token limit, which can truncate SSE lines with large payloads.func cozeChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { - scanner := bufio.NewScanner(resp.Body) - scanner.Split(bufio.ScanLines) + defer service.CloseResponseBodyGracefully(resp) + scanner := bufio.NewScanner(resp.Body) + // Allow larger SSE data lines (default 64KB may be too small) + scanner.Buffer(make([]byte, 64*1024), 2*1024*1024) + scanner.Split(bufio.ScanLines) helper.SetEventStreamHeaders(c)
237-246: Handle non-2xx HTTP status codes before JSON unmarshalWhen Coze returns an error status, attempting to unmarshal into a success DTO can mask the real error. Check
StatusCodeand surface the upstream body for diagnostics.responseBody, err := io.ReadAll(resp.Body) if err != nil { return fmt.Errorf("read response body failed: %w", err), false } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return fmt.Errorf("unexpected status %d: %s", resp.StatusCode, string(responseBody)), false + } err = json.Unmarshal(responseBody, &cozeResponse)relay/channel/volcengine/adaptor.go (4)
187-205: Harden URL construction: trim trailing slashes and validate base to avoid // and empty hostSwitching to ChannelBaseUrl is correct, but guard against empty or trailing-slash bases to avoid malformed URLs or sending relative paths.
Apply:
-func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { - switch info.RelayMode { +func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { + base := strings.TrimRight(info.ChannelBaseUrl, "/") + if base == "" { + return "", fmt.Errorf("channel base URL is empty") + } + switch info.RelayMode { @@ - if strings.HasPrefix(info.UpstreamModelName, "bot") { - return fmt.Sprintf("%s/api/v3/bots/chat/completions", info.ChannelBaseUrl), nil - } - return fmt.Sprintf("%s/api/v3/chat/completions", info.ChannelBaseUrl), nil + if strings.HasPrefix(info.UpstreamModelName, "bot") { + return fmt.Sprintf("%s/api/v3/bots/chat/completions", base), nil + } + return fmt.Sprintf("%s/api/v3/chat/completions", base), nil @@ - return fmt.Sprintf("%s/api/v3/embeddings", info.ChannelBaseUrl), nil + return fmt.Sprintf("%s/api/v3/embeddings", base), nil @@ - return fmt.Sprintf("%s/api/v3/images/generations", info.ChannelBaseUrl), nil + return fmt.Sprintf("%s/api/v3/images/generations", base), nil @@ - return fmt.Sprintf("%s/api/v3/images/edits", info.ChannelBaseUrl), nil + return fmt.Sprintf("%s/api/v3/images/edits", base), nil @@ - return fmt.Sprintf("%s/api/v3/rerank", info.ChannelBaseUrl), nil + return fmt.Sprintf("%s/api/v3/rerank", base), nilNote: If you’re on Go 1.19+, url.JoinPath can be used for safer joining (optional).
49-65: Bug: Non-file form fields read before parsing; use MultipartForm.Value after ParseMultipartFormUsing PostForm here is unreliable for multipart requests (it’s empty until parsed). Parse first, then read from MultipartForm.Value to ensure all text fields are preserved.
- // 获取所有表单字段 - formData := c.Request.PostForm - // 遍历表单字段并打印输出 - for key, values := range formData { - if key == "model" { - continue - } - for _, value := range values { - writer.WriteField(key, value) - } - } - - // Parse the multipart form to handle both single image and multiple images - if err := c.Request.ParseMultipartForm(32 << 20); err != nil { // 32MB max memory - return nil, errors.New("failed to parse multipart form") - } + // Parse multipart form first, then copy non-file fields into the new multipart + if err := c.Request.ParseMultipartForm(32 << 20); err != nil { // 32MB max memory + return nil, errors.New("failed to parse multipart form") + } + if c.Request.MultipartForm != nil { + for key, values := range c.Request.MultipartForm.Value { + // Skip fields we handle separately + if key == "model" || key == "image" || key == "image[]" || strings.HasPrefix(key, "image[") || key == "mask" { + continue + } + for _, value := range values { + writer.WriteField(key, value) + } + } + }
154-157: Handle writer.Close error and stabilize Content-Type propagationIgnoring writer.Close errors risks truncated bodies. Also capture FormDataContentType before closing for safety.
- // 关闭 multipart 编写器以设置分界线 - writer.Close() - c.Request.Header.Set("Content-Type", writer.FormDataContentType()) + // 关闭 multipart 编写器以设置分界线 + contentType := writer.FormDataContentType() + if err := writer.Close(); err != nil { + return nil, fmt.Errorf("close multipart writer failed: %w", err) + } + c.Request.Header.Set("Content-Type", contentType)
220-222: Fix: ConvertRerankRequest returns nil, nil (likely breaks downstream)Returning (nil, nil) can cause unexpected nil dereferences or no-op calls. If no transformation is needed, pass the request through; otherwise, return a clear error.
-func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) { - return nil, nil -} +func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) { + // Volcengine rerank is pass-through; adjust here if vendor-specific mapping is required. + return request, nil +}relay/channel/xai/text.go (1)
62-69: Fix nil-deref + stop scanning on client write failure in xAI stream handlerstreamResponseXAI2OpenAI can return nil (when xAIResp == nil) and the code currently dereferences openaiResponse before checking — this risks a panic. Also, if helper.ObjectData fails (e.g., client disconnected) we should stop scanning to avoid wasted upstream reads and logging spam.
- Location to fix:
- relay/channel/xai/text.go — xAIStreamHandler callback (current lines ~62–69).
- Verified behavior:
- streamResponseXAI2OpenAI returns nil when xAIResp == nil.
- StreamScannerHandler contract: dataHandler returns true = continue, false = stop (scanner goroutine does
if !success { return }).Suggested patch:
- openaiResponse := streamResponseXAI2OpenAI(xAIResp, usage) - _ = openai.ProcessStreamResponse(*openaiResponse, &responseTextBuilder, &toolCount) - err = helper.ObjectData(c, openaiResponse) - if err != nil { - common.SysLog(err.Error()) - } - return true + openaiResponse := streamResponseXAI2OpenAI(xAIResp, usage) + if openaiResponse == nil { + return true + } + _ = openai.ProcessStreamResponse(*openaiResponse, &responseTextBuilder, &toolCount) + if err := helper.ObjectData(c, openaiResponse); err != nil { + common.SysLog(err.Error()) + // stop scanning if client write fails to avoid wasted upstream reads + return false + } + return truerelay/channel/jimeng/image.go (1)
52-58: Close the response body on all paths (use defer).If io.ReadAll fails, resp.Body isn’t closed. Move to a deferred close to guarantee cleanup.
Apply this diff:
func jimengImageHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) { - var jimengResponse ImageResponse - responseBody, err := io.ReadAll(resp.Body) + defer service.CloseResponseBodyGracefully(resp) + var jimengResponse ImageResponse + responseBody, err := io.ReadAll(resp.Body) if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) } - service.CloseResponseBodyGracefully(resp)model/token.go (1)
79-85: Guard key slicing in model/token.go — add safe maskKey and use it in both placesVerified with ripgrep: token Key is defined as
char(48)(model/token.go) and redemption Key aschar(32)(model/redemption.go) — keys are normally long, but malformed/short inputs can still occur, so defend the slices.Files/locations to change:
- model/token.go — replace raw slices at ~lines 79–85 and 108–111.
- Add a helper maskKey in model/token.go (outside the changed ranges).
Apply these diffs:
- keyPrefix := key[:3] - keySuffix := key[len(key)-3:] - return token, errors.New("该令牌额度已用尽 TokenStatusExhausted[sk-" + keyPrefix + "***" + keySuffix + "]") + kp, ks := maskKey(key) + return token, errors.New("该令牌额度已用尽 TokenStatusExhausted[sk-" + kp + "***" + ks + "]")- keyPrefix := key[:3] - keySuffix := key[len(key)-3:] - return token, errors.New(fmt.Sprintf("[sk-%s***%s] 该令牌额度已用尽 !token.UnlimitedQuota && token.RemainQuota = %d", keyPrefix, keySuffix, token.RemainQuota)) + kp, ks := maskKey(key) + return token, errors.New(fmt.Sprintf("[sk-%s***%s] 该令牌额度已用尽 !token.UnlimitedQuota && token.RemainQuota = %d", kp, ks, token.RemainQuota))Add this helper in model/token.go (outside the selected ranges):
// maskKey returns safe 3-char prefix/suffix for display; falls back gracefully for short inputs. func maskKey(key string) (string, string) { if len(key) >= 6 { return key[:3], key[len(key)-3:] } if len(key) >= 3 { return key[:3], "" } return key, "" }relay/channel/gemini/adaptor.go (2)
73-86: Default SampleCount to 1 when request.N is zeroUpstream may reject a zero sample count; also, our ImageRequest.N defaults to zero when unset. Ensure a sane default.
-// build gemini imagen request -geminiRequest := dto.GeminiImageRequest{ +// build gemini imagen request +sampleCount := int(request.N) +if sampleCount <= 0 { + sampleCount = 1 +} +geminiRequest := dto.GeminiImageRequest{ @@ - Parameters: dto.GeminiImageParameters{ - SampleCount: int(request.N), + Parameters: dto.GeminiImageParameters{ + SampleCount: sampleCount, AspectRatio: aspectRatio, PersonGeneration: "allow_adult", // default allow adult },
153-156: Return an explicit error for unsupported rerank conversionReturning (nil, nil) can lead to downstream nil dereferences. Return a clear error.
-func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) { - return nil, nil -} +func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) { + return nil, errors.New("gemini: rerank is not supported") +}relay/channel/mokaai/adaptor.go (1)
53-58: Avoid double slashes in generated channel URLs — trim base and remove trailing slash from suffixShort: The repo shows ChannelBaseUrl is taken directly from context (relay/common/relay_info.go) and many adaptors build URLs with "%s/…" or concatenation — a trailing slash in ChannelBaseUrl will produce double slashes. Fix the mokaai adaptor minimally and prefer central normalization.
Files needing attention (non-exhaustive)
- relay/channel/mokaai/adaptor.go — immediate fix (shown below).
- relay/common/relay_info.go — InitChannelMeta sets ChannelBaseUrl from context without trimming; normalize here to prevent widespread issues.
- Other adaptors using "%s/…" or concatenation (examples found in search): relay/channel/siliconflow/adaptor.go, relay/channel/cloudflare/adaptor.go, relay/channel/volcengine/adaptor.go, relay/channel/tencent/adaptor.go, relay/channel/ollama/adaptor.go, relay/channel/openai/adaptor.go, relay/channel/coze/adaptor.go.
Minimal/local fix for mokaai (recommended if you want a tiny change)
- suffix := "chat/" + suffix := "chat" if strings.HasPrefix(info.UpstreamModelName, "m3e") { suffix = "embeddings" } - fullRequestURL := fmt.Sprintf("%s/%s", info.ChannelBaseUrl, suffix) + base := strings.TrimRight(info.ChannelBaseUrl, "/") + fullRequestURL := fmt.Sprintf("%s/%s", base, suffix) return fullRequestURL, nilPreferred/global fix (recommended)
- Normalize ChannelBaseUrl once in InitChannelMeta so all adaptors get a stable base:
- ChannelBaseUrl: common.GetContextKeyString(c, constant.ContextKeyChannelBaseUrl), + ChannelBaseUrl: strings.TrimRight(common.GetContextKeyString(c, constant.ContextKeyChannelBaseUrl), "/"),(add
import "strings"to relay/common/relay_info.go if not present)Alternative: use url.JoinPath (Go 1.19+) or an existing helper (relaycommon.GetFullRequestURL is already used by some adaptors) to construct paths safely.
relay/channel/jimeng/sign.go (1)
47-51: Fix logger context type and avoid logging full request payload (PII leak).
- logger.LogInfo expects context.Context, but a *gin.Context is passed — this won’t compile.
- Logging the entire request body can leak sensitive info. Log only safe metadata (length and hash).
Apply this diff:
- logger.LogInfo(c, fmt.Sprintf("SetPayloadHash body: %s", body)) - payloadHash := sha256.Sum256(body) - hexPayloadHash := hex.EncodeToString(payloadHash[:]) + payloadHash := sha256.Sum256(body) + hexPayloadHash := hex.EncodeToString(payloadHash[:]) + // Consider LogDebug if available instead of LogInfo + logger.LogInfo(c.Request.Context(), fmt.Sprintf("SetPayloadHash: len=%d sha256=%s", len(body), hexPayloadHash))relay/channel/mokaai/relay-mokaai.go (1)
54-83: Avoid writing headers twice; IOCopyBytesGracefully already sets headers and status.You set Content-Type and call WriteHeader before service.IOCopyBytesGracefully, but IOCopyBytesGracefully copies headers and writes the status itself. This can cause “superfluous response.WriteHeader call” and header inconsistencies.
- c.Writer.Header().Set("Content-Type", "application/json") - c.Writer.WriteHeader(resp.StatusCode) service.IOCopyBytesGracefully(c, resp, jsonResponse)Note: Since IOCopyBytesGracefully copies headers from resp, ensure the upstream response includes a correct Content-Type (typically application/json). If you must force application/json regardless of upstream, consider passing src=nil to IOCopyBytesGracefully and setting headers yourself, e.g.:
// Alternative approach if you need to force content-type: // c.Writer.Header().Set("Content-Type", "application/json") // service.IOCopyBytesGracefully(c, nil, jsonResponse)But in that case you’ll also need to ensure the proper status code is written manually before copying.
relay/helper/common.go (1)
102-122: Use correct logger context type and websocket.TextMessage constant.
- logger.LogError expects context.Context, but *gin.Context is passed — compilation issue.
- Replace magic number 1 with websocket.TextMessage for clarity.
func WssString(c *gin.Context, ws *websocket.Conn, str string) error { if ws == nil { - logger.LogError(c, "websocket connection is nil") + logger.LogError(c.Request.Context(), "websocket connection is nil") return errors.New("websocket connection is nil") } - //common.LogInfo(c, fmt.Sprintf("sending message: %s", str)) - return ws.WriteMessage(1, []byte(str)) + // common.LogInfo(c, fmt.Sprintf("sending message: %s", str)) + return ws.WriteMessage(websocket.TextMessage, []byte(str)) } func WssObject(c *gin.Context, ws *websocket.Conn, object interface{}) error { jsonData, err := json.Marshal(object) if err != nil { return fmt.Errorf("error marshalling object: %w", err) } if ws == nil { - logger.LogError(c, "websocket connection is nil") + logger.LogError(c.Request.Context(), "websocket connection is nil") return errors.New("websocket connection is nil") } - //common.LogInfo(c, fmt.Sprintf("sending message: %s", jsonData)) - return ws.WriteMessage(1, jsonData) + // common.LogInfo(c, fmt.Sprintf("sending message: %s", jsonData)) + return ws.WriteMessage(websocket.TextMessage, jsonData) }relay/channel/ali/rerank.go (1)
34-41: Ensure response body is closed on all pathsIf io.ReadAll fails, resp.Body is never closed (leak). Defer the close before reading and remove the later explicit close.
Apply this diff:
func RerankHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*types.NewAPIError, *dto.Usage) { - responseBody, err := io.ReadAll(resp.Body) + defer service.CloseResponseBodyGracefully(resp) + responseBody, err := io.ReadAll(resp.Body) if err != nil { return types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError), nil } - service.CloseResponseBodyGracefully(resp)relay/channel/cohere/relay-cohere.go (1)
165-170: Response body is never closed after streaming (resource leak).zhipuStreamHandler already closes with service.CloseResponseBodyGracefully(resp). Mirror that here to avoid leaking connections.
Apply this diff right after c.Stream finishes:
- }) + }) + service.CloseResponseBodyGracefully(resp)relay/channel/cloudflare/relay_cloudflare.go (2)
57-60: Bug: modifying choice by value has no effect; iterate by indexfor _, choice := range response.Choices iterates by value, so setting choice.Delta.Role doesn’t modify response.Choices. This also risks responseText not reflecting the intended mutation path.
Apply this diff:
- for _, choice := range response.Choices { - choice.Delta.Role = "assistant" - responseText += choice.Delta.GetContentString() - } + for i := range response.Choices { + response.Choices[i].Delta.Role = "assistant" + responseText += response.Choices[i].Delta.GetContentString() + }
41-46: Guard against non-data SSE linesSSE streams can include fields like "event:" or "id:". Trimming without verifying the prefix may attempt to unmarshal non-data lines.
Apply this diff:
- if len(data) < len("data: ") { + if len(data) < len("data: ") || !strings.HasPrefix(data, "data: ") { continue } data = strings.TrimPrefix(data, "data: ") data = strings.TrimSuffix(data, "\r")relay/channel/moonshot/adaptor.go (1)
90-101: Add default in DoResponse to avoid returning (nil, nil) for unsupported formatsIf RelayFormat is neither OpenAI nor Claude, the function falls through and returns (nil, nil), which is ambiguous and can cause upstream confusion.
Apply this diff:
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) { switch info.RelayFormat { - case types.RelayFormatOpenAI: + case types.RelayFormatOpenAI: adaptor := openai.Adaptor{} return adaptor.DoResponse(c, resp, info) case types.RelayFormatClaude: if info.IsStream { err, usage = claude.ClaudeStreamHandler(c, resp, info, claude.RequestModeMessage) } else { err, usage = claude.ClaudeHandler(c, resp, info, claude.RequestModeMessage) } + default: + // Fallback to OpenAI-compatible handling + adaptor := openai.Adaptor{} + return adaptor.DoResponse(c, resp, info) } return }relay/channel/dify/relay-dify.go (1)
153-160: Critical: nil pointer dereference for remote images; also incorrect Type valueWhen media.IsRemoteImage() is true, file is nil and fields are set on a nil pointer. Additionally, Type should likely be "image" (semantic type), not the MIME string (e.g., "image/jpeg").
Apply this diff:
- var file *DifyFile - if media.IsRemoteImage() { - file.Type = media.MimeType - file.TransferMode = "remote_url" - file.URL = media.Url - } else { - file = uploadDifyFile(c, info, difyReq.User, mediaContent) - } + var file *DifyFile + if media.IsRemoteImage() { + file = &DifyFile{ + Type: "image", + TransferMode: "remote_url", + URL: media.Url, + } + } else { + file = uploadDifyFile(c, info, difyReq.User, mediaContent) + }controller/task.go (2)
138-144: Wrong printf verb and returning a possibly-nil err when IsSuccess is false
- The log uses
%dbut passes a string.- Returning
errhere can return nil because unmarshal succeeded; explicitly return a new error.- if !responseItems.IsSuccess() { - common.SysLog(fmt.Sprintf("渠道 #%d 未完成的任务有: %d, 成功获取到任务数: %d", channelId, len(taskIds), string(responseBody))) - return err - } + if !responseItems.IsSuccess() { + logger.LogError(ctx, fmt.Sprintf("渠道 #%d 获取任务失败，响应: %s", channelId, string(responseBody))) + return fmt.Errorf("Get Task response unsuccessful") + }Also applies to: 141-144
157-174: Remove stale error check; it references an unrelated err and hides quota compensation errors
if err != nilchecks a stale variable. This silently skips quota compensation or logs the wrong thing. Flatten the logic and handle the IncreaseUserQuota error locally.- if responseItem.FailReason != "" || task.Status == model.TaskStatusFailure { - logger.LogInfo(ctx, task.TaskID+" 构建失败，"+task.FailReason) - task.Progress = "100%" - //err = model.CacheUpdateUserQuota(task.UserId) ? - if err != nil { - logger.LogError(ctx, "error update user quota cache: "+err.Error()) - } else { - quota := task.Quota - if quota != 0 { - err = model.IncreaseUserQuota(task.UserId, quota, false) - if err != nil { - logger.LogError(ctx, "fail to increase user quota: "+err.Error()) - } - logContent := fmt.Sprintf("异步任务执行失败 %s，补偿 %s", task.TaskID, logger.LogQuota(quota)) - model.RecordLog(task.UserId, model.LogTypeSystem, logContent) - } - } - } + if responseItem.FailReason != "" || task.Status == model.TaskStatusFailure { + logger.LogInfo(ctx, task.TaskID+" 构建失败，"+task.FailReason) + task.Progress = "100%" + quota := task.Quota + if quota != 0 { + if incErr := model.IncreaseUserQuota(task.UserId, quota, false); incErr != nil { + logger.LogError(ctx, "fail to increase user quota: "+incErr.Error()) + } + logContent := fmt.Sprintf("异步任务执行失败 %s，补偿 %s", task.TaskID, logger.LogQuota(quota)) + model.RecordLog(task.UserId, model.LogTypeSystem, logContent) + } + }model/user.go (1)
281-290: Transaction rollback happens unconditionally; use a panic-only rollback and proper SELECT FOR UPDATE
defer tx.Rollback()will attempt a rollback even after a successful commit, causing “transaction has already been committed/rolled back” errors in logs.- GORM v2 doesn’t respect
Set("gorm:query_option", "FOR UPDATE"). UseClauses(clause.Locking{Strength: "UPDATE"}).- // 开始数据库事务 - tx := DB.Begin() - if tx.Error != nil { - return tx.Error - } - defer tx.Rollback() // 确保在函数退出时事务能回滚 + // 开始数据库事务 + tx := DB.Begin() + if tx.Error != nil { + return tx.Error + } + // 仅在 panic 时回滚；正常路径在末尾 Commit + defer func() { + if r := recover(); r != nil { + _ = tx.Rollback() + panic(r) + } + }() @@ - // 加锁查询用户以确保数据一致性 - err := tx.Set("gorm:query_option", "FOR UPDATE").First(&user, user.Id).Error + // 加锁查询用户以确保数据一致性 + err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&user, user.Id).ErrorYou’ll need this import:
-import ( +import ( ... - "gorm.io/gorm" + "gorm.io/gorm" + "gorm.io/gorm/clause" )controller/midjourney.go (1)
80-99: Always cancel contexts and close bodies; otherwise you’ll leak goroutines/connections
- Add
defer cancel()immediately after creating the timeout context.- Close
req.Bodyandresp.Bodyvia defers; also close on non-200 paths.req, err := http.NewRequest("POST", requestUrl, bytes.NewBuffer(body)) if err != nil { logger.LogError(ctx, fmt.Sprintf("Get Task error: %v", err)) continue } + defer req.Body.Close() // 设置超时时间 timeout := time.Second * 15 ctx, cancel := context.WithTimeout(context.Background(), timeout) // 使用带有超时的 context 创建新的请求 req = req.WithContext(ctx) + defer cancel() @@ resp, err := service.GetHttpClient().Do(req) if err != nil { logger.LogError(ctx, fmt.Sprintf("Get Task Do req error: %v", err)) continue } + defer service.CloseResponseBodyGracefully(resp) if resp.StatusCode != http.StatusOK { logger.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode)) continue }controller/channel-test.go (2)
245-253: Avoid panic on type assertion for usageYou check usageA != nil, but still blindly cast to *dto.Usage. If DoResponse returns a different usage type or a non-nil interface with nil underlying value, this will panic.
Apply this diff:
- usage := usageA.(*dto.Usage) + usage, ok := usageA.(*dto.Usage) + if !ok || usage == nil { + return testResult{ + context: c, + localErr: errors.New("usage type assertion failed"), + newAPIError: types.NewOpenAIError(errors.New("usage type assertion failed"), types.ErrorCodeBadResponseBody, http.StatusInternalServerError), + } + }
264-273: Final quota calculation ignores GroupRatio; may under/over chargeModelPriceHelper applies GroupRatio in pre-consumed quota but priceData.ModelRatio does not include group ratio. Your final quota path multiplies only by ModelRatio, skipping group ratio. Include group ratio to keep consistency with pre-consume logic.
Apply this diff:
- quota = usage.PromptTokens + int(math.Round(float64(usage.CompletionTokens)*priceData.CompletionRatio)) - quota = int(math.Round(float64(quota) * priceData.ModelRatio)) + quota = usage.PromptTokens + int(math.Round(float64(usage.CompletionTokens)*priceData.CompletionRatio)) + totalRatio := priceData.ModelRatio * priceData.GroupRatioInfo.GroupRatio + quota = int(math.Round(float64(quota) * totalRatio)) if priceData.ModelRatio != 0 && quota <= 0 { quota = 1 }Optionally, consider a single helper for quota computation to avoid drift between pre-consume and final consume logic.
♻️ Duplicate comments (1)
relay/gemini_handler.go (1)
168-168: Verify the postConsumeQuota function signatureBoth
GeminiHelperandGeminiEmbeddingHandlercallpostConsumeQuotawith the same signature. As mentioned in the embedding_handler review, please ensure this function exists with the correct implementation.Also applies to: 264-264
| // 读取全部 option | ||
| opts, err := model.AllOption() | ||
| if err != nil { | ||
| c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) | ||
| return | ||
| } | ||
| // 建立 map | ||
| valMap := map[string]string{} | ||
| for _, o := range opts { | ||
| valMap[o.Key] = o.Value | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Preallocate the map and prepare to track migrated keys.
Micro-optimization and forward-prep for safe deletion later. Also, a dedicated slice for migrated keys will help avoid deleting non-migrated data.
Apply this diff:
- // 建立 map
- valMap := map[string]string{}
+ // 建立 map
+ valMap := make(map[string]string, len(opts))
for _, o := range opts {
valMap[o.Key] = o.Value
}
+
+ // 收集成功迁移的旧键（用于最后仅删除已成功迁移的键，避免数据丢失）
+ migratedKeys := make([]string, 0, 5)📝 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.
| // 读取全部 option | |
| opts, err := model.AllOption() | |
| if err != nil { | |
| c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) | |
| return | |
| } | |
| // 建立 map | |
| valMap := map[string]string{} | |
| for _, o := range opts { | |
| valMap[o.Key] = o.Value | |
| } | |
| // 读取全部 option | |
| opts, err := model.AllOption() | |
| if err != nil { | |
| c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) | |
| return | |
| } | |
| // 建立 map | |
| valMap := make(map[string]string, len(opts)) | |
| for _, o := range opts { | |
| valMap[o.Key] = o.Value | |
| } | |
| // 收集成功迁移的旧键（用于最后仅删除已成功迁移的键，避免数据丢失） | |
| migratedKeys := make([]string, 0, 5) |
🤖 Prompt for AI Agents
In controller/console_migrate.go around lines 16 to 26, the options map is
created without preallocation and there is no collection of migrated keys for
safe deletion; change creation of valMap to preallocate capacity using len(opts)
(e.g., make(map[string]string, len(opts))) and add a migratedKeys slice (e.g.,
var migratedKeys = make([]string, 0, len(opts))) then, as you iterate opts,
populate valMap[o.Key]=o.Value and append o.Key to migratedKeys so you can later
delete only those migrated keys safely.
| // 处理 APIInfo | ||
| if v := valMap["ApiInfo"]; v != "" { | ||
| var arr []map[string]interface{} | ||
| if err := json.Unmarshal([]byte(v), &arr); err == nil { | ||
| if len(arr) > 50 { | ||
| arr = arr[:50] | ||
| } | ||
| bytes, _ := json.Marshal(arr) | ||
| model.UpdateOption("console_setting.api_info", string(bytes)) | ||
| } | ||
| model.UpdateOption("ApiInfo", "") | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Avoid data loss when ApiInfo JSON is invalid; handle errors from UpdateOption and json.Marshal.
Currently the code clears the old key even if JSON parsing/marshalling fails, causing irreversible data loss. Also, UpdateOption errors are ignored.
Apply this diff:
// 处理 APIInfo
if v := valMap["ApiInfo"]; v != "" {
var arr []map[string]interface{}
- if err := json.Unmarshal([]byte(v), &arr); err == nil {
- if len(arr) > 50 {
- arr = arr[:50]
- }
- bytes, _ := json.Marshal(arr)
- model.UpdateOption("console_setting.api_info", string(bytes))
- }
- model.UpdateOption("ApiInfo", "")
+ if err := json.Unmarshal([]byte(v), &arr); err == nil {
+ if len(arr) > 50 {
+ arr = arr[:50]
+ }
+ if data, err := json.Marshal(arr); err == nil {
+ if err := model.UpdateOption("console_setting.api_info", string(data)); err != nil {
+ c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()})
+ return
+ }
+ if err := model.UpdateOption("ApiInfo", ""); err != nil {
+ c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()})
+ return
+ }
+ migratedKeys = append(migratedKeys, "ApiInfo")
+ } // marshalling failed: 保留旧值，不清空
+ } // parsing failed: 保留旧值，不清空
}📝 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.
| // 处理 APIInfo | |
| if v := valMap["ApiInfo"]; v != "" { | |
| var arr []map[string]interface{} | |
| if err := json.Unmarshal([]byte(v), &arr); err == nil { | |
| if len(arr) > 50 { | |
| arr = arr[:50] | |
| } | |
| bytes, _ := json.Marshal(arr) | |
| model.UpdateOption("console_setting.api_info", string(bytes)) | |
| } | |
| model.UpdateOption("ApiInfo", "") | |
| } | |
| // 处理 APIInfo | |
| if v := valMap["ApiInfo"]; v != "" { | |
| var arr []map[string]interface{} | |
| if err := json.Unmarshal([]byte(v), &arr); err == nil { | |
| if len(arr) > 50 { | |
| arr = arr[:50] | |
| } | |
| if data, err := json.Marshal(arr); err == nil { | |
| if err := model.UpdateOption("console_setting.api_info", string(data)); err != nil { | |
| c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) | |
| return | |
| } | |
| if err := model.UpdateOption("ApiInfo", ""); err != nil { | |
| c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) | |
| return | |
| } | |
| migratedKeys = append(migratedKeys, "ApiInfo") | |
| } // marshalling failed: 保留旧值，不清空 | |
| } // parsing failed: 保留旧值，不清空 | |
| } |
🤖 Prompt for AI Agents
In controller/console_migrate.go around lines 28-39, the code currently clears
the original "ApiInfo" key even when json.Unmarshal or json.Marshal fails and
ignores errors returned by model.UpdateOption; change the flow so that you only
update "console_setting.api_info" and clear "ApiInfo" after successful
json.Unmarshal and json.Marshal, keep the original data if parsing fails (do not
blank it), check and handle errors returned by json.Unmarshal, json.Marshal and
model.UpdateOption (log and return or propagate the error as appropriate), and
ensure the trimming to 50 items happens before marshalling and only proceed to
call UpdateOption when marshalling succeeds.
|
|
||
| func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { | ||
| baseUrl := fmt.Sprintf("%s/api/paas/v4", info.BaseUrl) | ||
| baseUrl := fmt.Sprintf("%s/api/paas/v4", info.ChannelBaseUrl) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Guard against empty or trailing-slash ChannelBaseUrl
If ChannelBaseUrl is empty, the constructed URL becomes a relative path, breaking http.NewRequest. Trailing slashes also produce double slashes. Trim and validate.
Apply these diffs:
@@
-import (
+import (
"errors"
"fmt"
"io"
"net/http"
"one-api/dto"
"one-api/relay/channel"
"one-api/relay/channel/openai"
relaycommon "one-api/relay/common"
relayconstant "one-api/relay/constant"
"one-api/types"
"github.com/gin-gonic/gin"
+ "strings"
) func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
- baseUrl := fmt.Sprintf("%s/api/paas/v4", info.ChannelBaseUrl)
+ base := strings.TrimRight(info.ChannelBaseUrl, "/")
+ if base == "" {
+ return "", fmt.Errorf("empty ChannelBaseUrl")
+ }
+ baseUrl := fmt.Sprintf("%s/api/paas/v4", base)
switch info.RelayMode {
case relayconstant.RelayModeEmbeddings:
return fmt.Sprintf("%s/embeddings", baseUrl), nil
default:
return fmt.Sprintf("%s/chat/completions", baseUrl), nil
}
}📝 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.
| baseUrl := fmt.Sprintf("%s/api/paas/v4", info.ChannelBaseUrl) | |
| import ( | |
| "errors" | |
| "fmt" | |
| "io" | |
| "net/http" | |
| "one-api/dto" | |
| "one-api/relay/channel" | |
| "one-api/relay/channel/openai" | |
| relaycommon "one-api/relay/common" | |
| relayconstant "one-api/relay/constant" | |
| "one-api/types" | |
| "github.com/gin-gonic/gin" | |
| "strings" | |
| ) | |
| func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { | |
| base := strings.TrimRight(info.ChannelBaseUrl, "/") | |
| if base == "" { | |
| return "", fmt.Errorf("empty ChannelBaseUrl") | |
| } | |
| baseUrl := fmt.Sprintf("%s/api/paas/v4", base) | |
| switch info.RelayMode { | |
| case relayconstant.RelayModeEmbeddings: | |
| return fmt.Sprintf("%s/embeddings", baseUrl), nil | |
| default: | |
| return fmt.Sprintf("%s/chat/completions", baseUrl), nil | |
| } | |
| } |
| common.SysLog("invalid zhipu key: " + apikey) | ||
| return "" | ||
| } |
There was a problem hiding this comment.
Avoid logging raw API keys (secrets leak).
This logs the full API key on invalid format, leaking credentials into logs. Log a generic message instead.
Apply this diff:
- common.SysLog("invalid zhipu key: " + apikey)
+ common.SysLog("invalid zhipu key format")📝 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.
| common.SysLog("invalid zhipu key: " + apikey) | |
| return "" | |
| } | |
| common.SysLog("invalid zhipu key format") | |
| return "" | |
| } |
🤖 Prompt for AI Agents
In relay/channel/zhipu/relay-zhipu.go around lines 42 to 44, the code currently
logs the full API key on invalid format which leaks secrets; change the log to
avoid printing the raw apikey (use a generic message like "invalid zhipu key
format" or log only a masked/hashed representation), and ensure no sensitive
data is concatenated into logs; update the SysLog call accordingly so it does
not include apikey.
| if relayFormat == types.RelayFormatOpenAIRealtime { | ||
| var err error | ||
| ws, err = upgrader.Upgrade(c.Writer, c.Request, nil) | ||
| if err != nil { | ||
| common.LogError(c, err.Error()) | ||
| newAPIError = err | ||
| break | ||
| helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError()) | ||
| return | ||
| } | ||
| defer ws.Close() | ||
| } |
There was a problem hiding this comment.
WebSocket upgrade error handler can dereference a nil ws; respond over HTTP instead.
When Upgrade fails, ws is nil. Calling helper.WssError(c, ws, ...) risks a nil deref or a no-op. Return an HTTP error (JSON) in this branch.
- ws, err = upgrader.Upgrade(c.Writer, c.Request, nil)
- if err != nil {
- helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError())
- return
- }
- defer ws.Close()
+ ws, err = upgrader.Upgrade(c.Writer, c.Request, nil)
+ if err != nil {
+ // Cannot write via WS when upgrade fails; respond HTTP JSON
+ openaiErr := types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError()
+ c.JSON(http.StatusBadRequest, gin.H{"error": openaiErr})
+ return
+ }
+ defer func() { _ = ws.Close() }()📝 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 relayFormat == types.RelayFormatOpenAIRealtime { | |
| var err error | |
| ws, err = upgrader.Upgrade(c.Writer, c.Request, nil) | |
| if err != nil { | |
| common.LogError(c, err.Error()) | |
| newAPIError = err | |
| break | |
| helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError()) | |
| return | |
| } | |
| defer ws.Close() | |
| } | |
| if relayFormat == types.RelayFormatOpenAIRealtime { | |
| var err error | |
| ws, err = upgrader.Upgrade(c.Writer, c.Request, nil) | |
| if err != nil { | |
| // Cannot write via WS when upgrade fails; respond HTTP JSON | |
| openaiErr := types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError() | |
| c.JSON(http.StatusBadRequest, gin.H{"error": openaiErr}) | |
| return | |
| } | |
| defer func() { _ = ws.Close() }() | |
| } |
🤖 Prompt for AI Agents
In controller/relay.go around lines 72 to 80, the WebSocket upgrade error branch
currently calls helper.WssError with a nil ws which can cause a nil deref;
instead, when upgrader.Upgrade returns an error, respond over HTTP (JSON) using
the Gin context and the same error payload (e.g.
types.NewError(...).ToOpenAIError()) with an appropriate HTTP status (500 or
mapped status), return immediately, and only call defer ws.Close() after a
successful upgrade; remove the helper.WssError(c, ws, ...) call from the error
path and ensure ws is non-nil before using or deferring Close().
| var upgrader = websocket.Upgrader{ | ||
| Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol | ||
| CheckOrigin: func(r *http.Request) bool { | ||
| return true // 允许跨域 | ||
| }, | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
WS Upgrader allows any origin; restrict to same-origin or an allowlist.
CheckOrigin: return true enables cross-site WebSocket hijacking. At minimum, restrict to same-host; ideally, implement an allowlist from config.
var upgrader = websocket.Upgrader{
- Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol
- CheckOrigin: func(r *http.Request) bool {
- return true // 允许跨域
- },
+ Subprotocols: []string{"realtime"}, // TODO: add other supported protocols if needed
+ CheckOrigin: func(r *http.Request) bool {
+ origin := r.Header.Get("Origin")
+ if origin == "" {
+ return true
+ }
+ // Allow same-host requests. Replace with a proper allowlist from settings if available.
+ return strings.Contains(origin, r.Host)
+ },
}Optionally, expand Subprotocols to include common variants ("oai-realtime", "oai-realtime-v1", "v1.realtime") once validated.
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In controller/relay.go around lines 214 to 219, the websocket.Upgrader currently
uses CheckOrigin: return true which permits any origin; change it to validate
origins by either enforcing same-origin or consulting a configurable allowlist:
read allowed origins from configuration (env or config struct), parse the
request's Origin header, and return true only if the Origin matches the request
host (same-origin) or is present in the allowlist; ensure nil/missing Origin is
handled safely. Replace the unrestricted CheckOrigin with a function that
performs this check and logs/returns false for disallowed origins. Optionally,
expand Subprotocols to include validated variants like "oai-realtime",
"oai-realtime-v1", "v1.realtime" after confirming compatibility.
There was a problem hiding this comment.
Actionable comments posted: 3
🔭 Outside diff range comments (2)
model/channel.go (1)
405-414: Harden GetBaseURL: avoid panic and ensure consistent fallback when BaseURL is unsetTwo issues:
- Potential panic: indexing ChannelBaseURLs with an out-of-range channel.Type.
- Inconsistent fallback: nil BaseURL currently returns "", instead of falling back to the default for channel.Type.
Apply this fix to safely guard the index and fall back when BaseURL is nil or empty.
func (channel *Channel) GetBaseURL() string { - if channel.BaseURL == nil { - return "" - } - url := *channel.BaseURL - if url == "" { - url = constant.ChannelBaseURLs[channel.Type] - } - return url + // Prefer explicit BaseURL if set and non-empty; otherwise fallback by type + if channel.BaseURL != nil && *channel.BaseURL != "" { + return *channel.BaseURL + } + // Guard against out-of-range type values + if channel.Type >= 0 && channel.Type < len(constant.ChannelBaseURLs) { + return constant.ChannelBaseURLs[channel.Type] + } + return "" }Optional nit: consider normalizing the returned URL (e.g., trimming trailing slashes) to reduce downstream URL-join edge cases.
controller/channel-test.go (1)
314-315: Action required: make ParseInput accept []string (or normalize inputs before use)Verified: GeneralOpenAIRequest.ParseInput (dto/openai_request.go) and EmbeddingRequest.ParseInput (dto/embedding.go) only handle string and []any — they do not handle []string. That explains the test using []any as a workaround, but it’s fragile.
Files to fix:
- dto/openai_request.go — func (r *GeneralOpenAIRequest) ParseInput() []string
- dto/embedding.go — func (r *EmbeddingRequest) ParseInput() []string
- controller/channel-test.go — buildTestRequest currently uses testRequest.Input = []any{"hello world"} (workaround)
Suggested minimal change (add this case to both ParseInput implementations):
case []string: arr := r.Input.([]string) input = make([]string, 0, len(arr)) for _, s := range arr { input = append(input, s) }After adding the above, you can revert the test to use a natural type:
testRequest.Input = []string{"hello world"} // ParseInput now supports []string
♻️ Duplicate comments (6)
controller/channel-test.go (1)
137-145: Incorrect RelayFormat may cause issues for embedding modelsWhen testing an embeddings model (based on the requestPath being "/v1/embeddings"), you're still passing
types.RelayFormatOpenAItoGenRelayInfo. This may lead to incorrect relay mode detection and processing. Consider using the appropriateRelayFormatEmbeddingfor embedding requests.Apply this diff to use the correct relay format:
+ relayFormat := types.RelayFormatOpenAI + if requestPath == "/v1/embeddings" { + relayFormat = types.RelayFormatEmbedding + } request := buildTestRequest(testModel) - info, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAI, request, nil) + info, err := relaycommon.GenRelayInfo(c, relayFormat, request, nil)controller/relay.go (4)
72-80: WebSocket error handling could fail when upgrade failsWhen the WebSocket upgrade fails,
wsis nil but you're still passing it tohelper.WssError. This could cause issues ifWssErrordoesn't handle nil WebSocket connections properly.
115-122: Incorrect sensitive words detection logicThe sensitive words check is inverted. According to
service.CheckSensitiveTextsignature, it returns(words []string, err error)where detection is indicated by non-empty words with nil error, not by an error.
136-145: Missing error propagation for pre-consume quota failureWhen
PreConsumeQuotafails, the function returns without setting a response, leaving the client hanging. The defer for returning quota also won't execute since it's registered after the early return.
214-219: WebSocket upgrader allows any originThe
CheckOriginfunction returnstruefor all origins, which could enable cross-site WebSocket hijacking attacks.relay/common/relay_info.go (1)
297-297: Compile error: SupportStreamOptions field no longer exists on RelayInfoThe code assigns
info.SupportStreamOptions = falsebut this field has been moved toChannelMeta. This will cause a compilation error.
🧹 Nitpick comments (1)
service/pre_consume_quota.go (1)
64-77: Good transaction ordering but consider rollback scenariosThe code correctly decreases token quota before user quota, which is good for consistency. However, if
DecreaseUserQuotafails afterPreConsumeTokenQuotasucceeds, the token quota won't be rolled back.Consider implementing a rollback mechanism:
if preConsumedQuota > 0 { err := PreConsumeTokenQuota(relayInfo, preConsumedQuota) if err != nil { return 0, types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) } err = model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota) if err != nil { + // Attempt to rollback token quota + rollbackErr := model.IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, preConsumedQuota) + if rollbackErr != nil { + logger.LogError(c, fmt.Sprintf("Failed to rollback token quota after user quota decrease failed: %v", rollbackErr)) + } return 0, types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry()) } logger.LogInfo(c, fmt.Sprintf("用户 %d 预扣费 %s, 预扣费后剩余额度: %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota), logger.FormatQuota(userQuota-preConsumedQuota))) }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (5)
controller/channel-test.go(2 hunks)controller/relay.go(6 hunks)model/channel.go(14 hunks)relay/common/relay_info.go(6 hunks)service/pre_consume_quota.go(1 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:
controller/channel-test.go
🧬 Code Graph Analysis (5)
service/pre_consume_quota.go (7)
relay/common/relay_info.go (1)
RelayInfo(73-120)service/quota.go (2)
PostConsumeQuota(481-510)PreConsumeTokenQuota(457-479)common/logger.go (1)
SysLog(10-13)types/error.go (9)
NewAPIError(80-88)NewError(179-191)ErrorCodeQueryDataError(72-72)ErrOptionWithSkipRetry(288-292)NewErrorWithStatusCode(213-229)ErrorCodeInsufficientUserQuota(76-76)ErrOptionWithNoRecordErrorLog(294-298)ErrorCodePreConsumeTokenQuotaFailed(77-77)ErrorCodeUpdateDataError(73-73)model/user.go (2)
GetUserQuota(570-595)DecreaseUserQuota(691-706)logger/logger.go (1)
FormatQuota(99-105)common/quota.go (1)
GetTrustQuota(3-5)
controller/relay.go (18)
relay/common/relay_info.go (2)
RelayInfo(73-120)GenRelayInfo(404-437)relay/embedding_handler.go (1)
EmbeddingHelper(18-72)relay/relay-text.go (1)
TextHelper(27-184)dto/request_common.go (1)
Request(8-11)relay/gemini_handler.go (1)
GeminiHelper(53-170)types/relay_format.go (5)
RelayFormat(3-3)RelayFormatOpenAIRealtime(12-12)RelayFormatClaude(7-7)RelayFormatGemini(8-8)RelayFormatMjProxy(17-17)common/utils.go (1)
MessageWithRequestId(209-211)logger/logger.go (3)
LogWarn(56-58)LogError(60-62)LogInfo(52-54)service/token_counter.go (1)
CountRequestToken(246-288)relay/helper/price.go (1)
ModelPriceHelper(44-102)service/pre_consume_quota.go (1)
PreConsumeQuota(31-77)common/gin.go (2)
GetRequestBody(15-27)GetContextKeyBool(68-70)relay/websocket.go (1)
WssHelper(14-45)relay/claude_handler.go (1)
ClaudeHelper(20-131)constant/env.go (1)
ErrorLogEnabled(15-15)constant/context_key.go (2)
ContextKeyChannelIsMultiKey(33-33)ContextKeyChannelMultiKeyIndex(34-34)model/log.go (1)
RecordErrorLog(95-135)service/error.go (1)
TaskErrorWrapperLocal(131-135)
model/channel.go (2)
common/logger.go (1)
SysLog(10-13)constant/channel.go (1)
ChannelBaseURLs(57-111)
controller/channel-test.go (6)
relay/common/relay_info.go (1)
GenRelayInfo(404-437)types/relay_format.go (1)
RelayFormatOpenAI(6-6)types/error.go (3)
NewError(179-191)ErrorCodeGenRelayInfoFailed(48-48)ErrorCodeChannelModelMappedError(53-53)relay/helper/model_mapped.go (1)
ModelMappedHelper(15-93)common/logger.go (1)
SysLog(10-13)relay/helper/price.go (1)
ModelPriceHelper(44-102)
relay/common/relay_info.go (11)
dto/channel_settings.go (2)
ChannelSettings(3-10)ChannelOtherSettings(12-14)common/constants.go (1)
StartTime(12-12)dto/realtime.go (1)
RealTimeTool(66-71)types/relay_format.go (12)
RelayFormat(3-3)RelayFormatOpenAIRealtime(12-12)RelayFormatClaude(7-7)RelayFormatRerank(13-13)RelayFormatOpenAIAudio(10-10)RelayFormatEmbedding(14-14)RelayFormatOpenAIResponses(9-9)RelayFormatGemini(8-8)RelayFormatOpenAIImage(11-11)RelayFormatOpenAI(6-6)RelayFormatTask(16-16)RelayFormatMjProxy(17-17)types/price_data.go (1)
PriceData(11-21)dto/request_common.go (1)
Request(8-11)common/gin.go (6)
GetContextKeyInt(64-66)GetContextKeyStringMap(76-78)GetContextKeyBool(68-70)GetContextKeyString(60-62)GetContextKeyType(84-92)GetContextKeyTime(80-82)constant/context_key.go (20)
ContextKeyChannelType(25-25)ContextKeyChannelParamOverride(28-28)ContextKeyChannelId(21-21)ContextKeyChannelIsMultiKey(33-33)ContextKeyChannelMultiKeyIndex(34-34)ContextKeyChannelBaseUrl(24-24)ContextKeyChannelKey(35-35)ContextKeyOriginalModel(8-8)ContextKeyChannelSetting(26-26)ContextKeyChannelOtherSetting(27-27)ContextKeyRequestStartTime(9-9)ContextKeyUserId(38-38)ContextKeyUsingGroup(44-44)ContextKeyUserGroup(43-43)ContextKeyUserQuota(40-40)ContextKeyUserEmail(42-42)ContextKeyPromptTokens(6-6)ContextKeyTokenId(14-14)ContextKeyTokenKey(13-13)ContextKeyTokenUnlimited(12-12)dto/rerank.go (1)
RerankRequest(10-18)relay/constant/relay_mode.go (3)
RelayModeRerank(46-46)RelayModeResponses(48-48)Path2RelayMode(55-89)dto/openai_request.go (1)
OpenAIResponsesRequest(754-776)
🔇 Additional comments (8)
controller/channel-test.go (1)
189-200: Good separation of conversion logic based on RelayModeThe code correctly handles different conversion paths for embeddings vs other request types, which aligns with the new relay architecture.
service/pre_consume_quota.go (2)
36-41: Clear and informative error messagesThe error messages properly distinguish between zero quota and insufficient quota scenarios, providing clear feedback to users with formatted quota values.
43-62: Well-designed trust quota logicThe trust quota mechanism elegantly handles the balance between user experience and system protection. When users have sufficient quota history, it avoids pre-consumption to reduce latency, while still protecting against abuse for newer or lower-quota users.
controller/relay.go (1)
159-168: Clean relay format dispatchThe switch statement properly routes different relay formats to their respective handlers, maintaining good separation of concerns.
relay/common/relay_info.go (4)
54-71: Well-structured ChannelMeta extractionThe extraction of channel-related metadata into a dedicated
ChannelMetastruct improves code organization and makes the relay flow more maintainable. The fields are comprehensive and properly typed.
122-157: Robust channel metadata initializationThe
InitChannelMetamethod properly extracts all channel-related information from the context with safe type assertions and default values. The stream support detection is cleanly integrated.
159-232: Excellent ToString implementation with proper secret maskingThe
ToStringmethod provides comprehensive debugging information while properly masking sensitive data (API keys, token keys). The structured output format makes it easy to parse logs.
404-437: Clean relay format routing with proper error handlingThe
GenRelayInfofunction provides a centralized dispatcher for different relay formats with appropriate type assertions and error messages for invalid inputs.
| func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) { | ||
| if preConsumedQuota != 0 { | ||
| gopool.Go(func() { | ||
| relayInfoCopy := *relayInfo | ||
|
|
||
| err := PostConsumeQuota(&relayInfoCopy, -preConsumedQuota, 0, false) | ||
| if err != nil { | ||
| common.SysLog("error return pre-consumed quota: " + err.Error()) | ||
| } | ||
| }) | ||
| } | ||
| } |
There was a problem hiding this comment.
Potential race condition in async quota return
The goroutine copies *relayInfo but this shallow copy might not be safe if RelayInfo contains pointers to mutable data that could be modified after the goroutine starts. The copied RelayInfo shares the same pointer fields (like ChannelMeta, ClaudeConvertInfo, etc.) with the original.
Consider deep copying or ensuring the RelayInfo data is immutable after this point:
func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) {
if preConsumedQuota != 0 {
gopool.Go(func() {
- relayInfoCopy := *relayInfo
+ // Create a defensive copy to avoid race conditions
+ // Only copy the fields needed for PostConsumeQuota
+ relayInfoCopy := relaycommon.RelayInfo{
+ UserId: relayInfo.UserId,
+ TokenId: relayInfo.TokenId,
+ TokenKey: relayInfo.TokenKey,
+ IsPlayground: relayInfo.IsPlayground,
+ }
err := PostConsumeQuota(&relayInfoCopy, -preConsumedQuota, 0, false)
if err != nil {
common.SysLog("error return pre-consumed quota: " + err.Error())
}
})
}
}📝 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 ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) { | |
| if preConsumedQuota != 0 { | |
| gopool.Go(func() { | |
| relayInfoCopy := *relayInfo | |
| err := PostConsumeQuota(&relayInfoCopy, -preConsumedQuota, 0, false) | |
| if err != nil { | |
| common.SysLog("error return pre-consumed quota: " + err.Error()) | |
| } | |
| }) | |
| } | |
| } | |
| func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) { | |
| if preConsumedQuota != 0 { | |
| gopool.Go(func() { | |
| // Create a defensive copy to avoid race conditions | |
| // Only copy the fields needed for PostConsumeQuota | |
| relayInfoCopy := relaycommon.RelayInfo{ | |
| UserId: relayInfo.UserId, | |
| TokenId: relayInfo.TokenId, | |
| TokenKey: relayInfo.TokenKey, | |
| IsPlayground: relayInfo.IsPlayground, | |
| } | |
| err := PostConsumeQuota(&relayInfoCopy, -preConsumedQuota, 0, false) | |
| if err != nil { | |
| common.SysLog("error return pre-consumed quota: " + err.Error()) | |
| } | |
| }) | |
| } | |
| } |
🤖 Prompt for AI Agents
In service/pre_consume_quota.go around lines 16 to 27, the goroutine makes a
shallow copy of *relayInfo which can still share pointer fields and cause races
when the original is mutated; replace the shallow copy with a safe deep copy or
snapshot of the data before spawning the goroutine (for example implement and
call a RelayInfo.Clone() that deep-copies all pointer/slice/map fields or
marshal/unmarshal via protobuf/JSON to produce an independent instance), or
alternatively capture only the primitive fields needed by PostConsumeQuota and
pass those instead; ensure the copy is created synchronously before gopool.Go so
the goroutine has an isolated, immutable relayInfo.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
common/str.go (1)
190-214: Domain-only masking: remove redundant URL check and align ccTLD handling with URL logic.
- The "skip if it's part of a URL" check will almost never trigger, because URL masking has already replaced hosts inside URLs (e.g., http://openai.com → http://***.com), so "://"+domain won’t exist anymore. It can be dropped.
- Consider aligning country-code TLD handling with the URL branch to keep behavior consistent between "http://sub.domain.co.uk" and "sub.domain.co.uk". Currently, URL masking preserves "co.uk" (.co.uk), but domain-only masking reduces it to ".***.uk".
Apply this diff to simplify and align behaviors:
- // Mask domain names without protocol (like openai.com, www.openai.com) - domainPattern := regexp.MustCompile(`\b(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}\b`) - str = domainPattern.ReplaceAllStringFunc(str, func(domain string) string { - // Skip if it's already been processed as part of a URL - if strings.Contains(str, "://"+domain) { - return domain - } - - parts := strings.Split(domain, ".") - if len(parts) < 2 { - return domain - } - - // Handle different domain patterns - if len(parts) == 2 { - // openai.com -> ***.com - return "***." + parts[1] - } else { - // www.openai.com -> ***.***.com - // api.openai.com -> ***.***.com - lastPart := parts[len(parts)-1] - return "***.***." + lastPart - } - }) + // Mask domain names without protocol (like openai.com, www.openai.com) + domainPattern := regexp.MustCompile(`\b(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}\b`) + str = domainPattern.ReplaceAllStringFunc(str, func(domain string) string { + parts := strings.Split(domain, ".") + if len(parts) < 2 { + return domain + } + if len(parts) == 2 { + // openai.com -> ***.com + return "***." + parts[1] + } + // For multi-part domains, align with the URL masking ccTLD logic + last := parts[len(parts)-1] + secondLast := parts[len(parts)-2] + if len(last) == 2 && len(secondLast) <= 3 { + // Likely ccTLD like co.uk / com.cn + return "***." + secondLast + "." + last + } + return "***." + last + })Additional optional improvements (outside the selected range):
- Prefer url.URL.Hostname() and Port() when masking URLs to avoid leaking and mis-parsing ports (e.g., example.com:8443). For example:
// inside the URL masking block host := u.Hostname() port := u.Port() // compute maskedHost from `host` only, then: if port != "" { maskedHost += ":" + port }
- Precompile regexes at package scope to avoid per-call compilation overhead, and use a simpler/more robust URL pattern:
- urlPattern:
https?://\S+- Also consider adding IPv6 masking. If regex becomes unwieldy, a replace-by-callback loop using net.ParseIP can be more robust than regex for IPs.
types/error.go (1)
123-134: Improve message normalization for 503 to cover non-Chinese messages and rely on code first.
- Since you now pass ErrorCodeModelNotFound explicitly, prefer checking the error code first, and only fall back to substring heuristics.
- Expand substring heuristics to include English to avoid inconsistent user-facing behavior in mixed-language deployments.
Apply this diff within MaskSensitiveError:
- errStr := e.Err.Error() - if e.StatusCode == http.StatusServiceUnavailable { - if e.errorCode == ErrorCodeModelNotFound { - errStr = "上游分组模型服务不可用，请稍后再试" - } else { - if strings.Contains(errStr, "分组") || strings.Contains(errStr, "渠道") { - errStr = "上游分组模型服务不可用，请稍后再试" - } - } - } - return common.MaskSensitiveInfo(errStr) + errStr := e.Err.Error() + if e.StatusCode == http.StatusServiceUnavailable { + if e.errorCode == ErrorCodeModelNotFound { + errStr = "上游分组模型服务不可用，请稍后再试" + } else { + // Heuristic fallback: cover both CN and EN keywords + lower := strings.ToLower(errStr) + if strings.Contains(errStr, "分组") || strings.Contains(errStr, "渠道") || + strings.Contains(lower, "group") || strings.Contains(lower, "channel") { + errStr = "上游分组模型服务不可用，请稍后再试" + } + } + } + return common.MaskSensitiveInfo(errStr)Optional: To apply this normalization consistently across OpenAI/Claude conversions, consider using e.MaskSensitiveError() for the message in ToOpenAIError/ToClaudeError instead of masking e.Error() locally.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (4)
common/str.go(2 hunks)middleware/distributor.go(1 hunks)middleware/utils.go(2 hunks)types/error.go(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- middleware/utils.go
🧰 Additional context used
🧬 Code Graph Analysis (2)
middleware/distributor.go (1)
types/error.go (1)
ErrorCodeModelNotFound(70-70)
types/error.go (2)
relay/channel/baidu/dto.go (1)
Error(26-29)common/str.go (1)
MaskSensitiveInfo(111-220)
🔇 Additional comments (3)
common/str.go (1)
102-111: Doc updates look good and clearer masking examples help.Examples now explicitly cover domain-only inputs and align with the implemented behavior.
middleware/distributor.go (1)
110-115: Good: explicit error code propagation verified (model_not_found) — no action requiredabortWithOpenAiMessage has an optional code param and the distributor changes use it correctly.
- middleware/utils.go:10 — func abortWithOpenAiMessage(c *gin.Context, statusCode int, message string, code ...string)
- middleware/distributor.go:110–115 — calls pass string(types.ErrorCodeModelNotFound)
types/error.go (1)
41-49: New error codes are appropriate and consistent with existing taxonomy.
- ErrorCodeGenRelayInfoFailed in the "new api error" group
- ErrorCodeModelNotFound in the "response error" group
Both are clear and integrate well with the broader refactor.
Also applies to: 70-70
This commit introduces a new function, MaskEmail, to mask user email addresses in logs, preventing PII leakage. Additionally, the RelayInfo logging has been updated to utilize this new masking function, ensuring sensitive information is properly handled. The channel test logic has also been improved to dynamically determine the relay format based on the request path.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (3)
relay/common/relay_info.go (2)
179-183: PII masking in ToString is correct nowUserEmail is masked via common.MaskEmail; Token key and ApiKey are redacted. This resolves prior PII leakage concerns.
292-320: Compile-time error: assigning removed field (SupportStreamOptions) on RelayInfoRelayInfo no longer has SupportStreamOptions; it lives under ChannelMeta. This line won’t compile and must be removed. If you need to disallow streaming options for Responses, enforce it in adaptor/validation, not on RelayInfo.
func GenRelayInfoResponses(c *gin.Context, request *dto.OpenAIResponsesRequest) *RelayInfo { info := genBaseRelayInfo(c, request) info.RelayMode = relayconstant.RelayModeResponses info.RelayFormat = types.RelayFormatOpenAIResponses - info.SupportStreamOptions = false - info.ResponsesUsageInfo = &ResponsesUsageInfo{ BuiltInTools: make(map[string]*BuildInToolInfo), }controller/channel-test.go (1)
137-141: Correctly select RelayFormat for embeddings (addresses prior feedback)Switching to RelayFormatEmbedding when the request path is /v1/embeddings ensures the generator builds the right RelayInfo.
🧹 Nitpick comments (9)
common/str.go (5)
104-118: Harden email masking: trim input and use LastIndex for '@' to avoid odd casesUsing strings.Index can leak unexpected characters if the input contains multiple '@'s; trimming helps with stray whitespace. Extract the domain using the last '@' and validate both sides exist.
func MaskEmail(email string) string { - if email == "" { + email = strings.TrimSpace(email) + if email == "" { return "***masked***" } - // Find the @ symbol - atIndex := strings.Index(email, "@") - if atIndex == -1 { - // No @ symbol found, return masked + // Find the last @ symbol to handle odd inputs like "a@b@c.com" + atIndex := strings.LastIndex(email, "@") + if atIndex <= 0 || atIndex == len(email)-1 { + // Invalid or missing @, return masked return "***masked***" } // Return only the domain part with @ symbol return "***@" + email[atIndex+1:] }
120-129: Doc examples look good; consider clarifying ccTLD behavior for bare domainsFor consistency with URL masking examples that preserve ccTLD patterns (e.g., co.uk), consider adding a bare-domain example like "sub.domain.co.uk -> .co.uk" (or explicitly documenting the intended ".***.uk" behavior).
226-230: Align bare-domain ccTLD handling with URL masking (optional)For domains with multi-part ccTLDs (e.g., co.uk, com.cn), URL masking preserves the last two parts. Mirror that here for consistency.
- } else { - // www.openai.com -> ***.***.com - // api.openai.com -> ***.***.com - lastPart := parts[len(parts)-1] - return "***.***." + lastPart - } + } else { + // Preserve common ccTLD second-level patterns (e.g., co.uk, com.cn) + last := parts[len(parts)-1] + secondLast := parts[len(parts)-2] + if len(last) == 2 && len(secondLast) <= 3 { + // sub.domain.co.uk -> ***.co.uk + return "***." + secondLast + "." + last + } + // www.openai.com -> ***.***.com + // api.openai.com -> ***.***.com + return "***.***." + last + }
209-209: Precompile regexes to avoid per-call allocationsurlPattern, domainPattern, and ipPattern are compiled on every call. Precompile once at package scope to reduce GC pressure and improve throughput.
Add at package scope:
var ( urlPattern = regexp.MustCompile(`(http|https)://[^\s/$.?#].[^\s]*`) domainPattern = regexp.MustCompile(`\b(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}\b`) ipPattern = regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}\b`) )Then inside MaskSensitiveInfo, use these precompiled variables and remove local MustCompile calls.
120-238: Add unit tests for new masking paths (bare domains and email)The added logic is easy to regress. Recommend table-driven tests covering:
- Bare domains: openai.com, www.openai.com, api.openai.com, sub.domain.co.uk
- URLs with/without ports and ccTLDs
- Emails: valid, empty, whitespace, invalid (no '@', multiple '@')
I can draft a go test file (common/str_test.go) with table-driven cases to lock behavior. Want me to provide it?
relay/common/relay_info.go (2)
159-232: Clamp negative latency in ToString to avoid confusing logsFirstResponseTime is initialized to StartTime-1s, which makes latency negative until the first response. Clamp to >= 0 for clearer logs.
- latencyMs := info.FirstResponseTime.Sub(info.StartTime).Milliseconds() + latencyMs := info.FirstResponseTime.Sub(info.StartTime).Milliseconds() + if latencyMs < 0 { + latencyMs = 0 + }
342-402: Centralize ChannelMeta initialization to reduce call-site burdenToday, callers must remember to invoke info.InitChannelMeta(c) (e.g., controller/channel-test.go). Consider initializing ChannelMeta inside genBaseRelayInfo so every RelayInfo is fully hydrated when returned.
func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { @@ info := &RelayInfo{ @@ } + // Ensure ChannelMeta is always initialized for downstream consumers. + info.InitChannelMeta(c) + return info }Also, consider removing the commented-out context-getter code to keep genBaseRelayInfo tight.
controller/channel-test.go (2)
234-244: Close HTTP response body to avoid resource leaksEnsure response bodies are closed even on success or error paths.
- if resp != nil { - httpResp = resp.(*http.Response) - if httpResp.StatusCode != http.StatusOK { + if resp != nil { + httpResp = resp.(*http.Response) + defer httpResp.Body.Close() + if httpResp.StatusCode != http.StatusOK { err := service.RelayErrorHandler(httpResp, true) return testResult{ context: c, localErr: err, newAPIError: types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError), } } }
261-269: Also close the httptest recorder response bodyMinor hygiene to prevent descriptor leaks in tests.
- result := w.Result() - respBody, err := io.ReadAll(result.Body) + result := w.Result() + defer result.Body.Close() + respBody, err := io.ReadAll(result.Body)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (4)
common/str.go(2 hunks)controller/channel-test.go(2 hunks)relay/common/relay_info.go(6 hunks)service/pre_consume_quota.go(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- service/pre_consume_quota.go
🧰 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:
controller/channel-test.go
🧬 Code Graph Analysis (2)
relay/common/relay_info.go (11)
dto/channel_settings.go (2)
ChannelSettings(3-10)ChannelOtherSettings(12-14)dto/realtime.go (1)
RealTimeTool(66-71)types/relay_format.go (12)
RelayFormat(3-3)RelayFormatOpenAIRealtime(12-12)RelayFormatClaude(7-7)RelayFormatRerank(13-13)RelayFormatOpenAIAudio(10-10)RelayFormatEmbedding(14-14)RelayFormatOpenAIResponses(9-9)RelayFormatGemini(8-8)RelayFormatOpenAIImage(11-11)RelayFormatOpenAI(6-6)RelayFormatTask(16-16)RelayFormatMjProxy(17-17)types/price_data.go (1)
PriceData(11-21)dto/request_common.go (1)
Request(8-11)common/gin.go (6)
GetContextKeyInt(64-66)GetContextKeyStringMap(76-78)GetContextKeyBool(68-70)GetContextKeyString(60-62)GetContextKeyType(84-92)GetContextKeyTime(80-82)common/api_type.go (1)
ChannelType2APIType(5-75)common/str.go (1)
MaskEmail(104-118)dto/rerank.go (1)
RerankRequest(10-18)relay/constant/relay_mode.go (2)
RelayModeRerank(46-46)Path2RelayMode(55-89)dto/openai_request.go (1)
OpenAIResponsesRequest(754-776)
controller/channel-test.go (6)
types/relay_format.go (2)
RelayFormatOpenAI(6-6)RelayFormatEmbedding(14-14)dto/request_common.go (1)
Request(8-11)relay/common/relay_info.go (1)
GenRelayInfo(404-437)types/error.go (3)
NewError(190-202)ErrorCodeGenRelayInfoFailed(48-48)ErrorCodeChannelModelMappedError(53-53)relay/helper/model_mapped.go (1)
ModelMappedHelper(15-93)common/logger.go (1)
SysLog(10-13)
🔇 Additional comments (17)
relay/common/relay_info.go (13)
54-71: Good extraction of channel-scoped fields into ChannelMetaConsolidating channel/config fields into ChannelMeta improves cohesion and makes RelayInfo leaner. The SupportStreamOptions flag belongs here.
122-157: InitChannelMeta initialization looks solidReads from gin.Context using typed getters and derives ApiType from ChannelType. Populating SupportStreamOptions via streamSupportedChannels is appropriate.
249-257: LGTM: Realtime info generationUsing base relay info and setting OpenAIRealtime format plus audio defaults is clean and consistent.
259-267: LGTM: Claude relay info generationSets format and disables usage accumulation as expected. The ClaudeConvertInfo bootstrap is fine.
269-277: LGTM: Rerank info generationSets RelayMode and RelayFormat appropriately and captures request-specific knobs (documents, return flags).
280-284: LGTM: OpenAI Audio relay info generation
286-290: LGTM: Embedding relay info generation
322-328: LGTM: Gemini relay info generationFormat and usage toggling align with expectations for Gemini handlers.
330-334: LGTM: Image relay info generation
336-340: LGTM: OpenAI relay info generation
404-437: Dispatcher reads clean and type-safeSwitch-based dispatch with type assertions for specialized requests is straightforward. Error messages for mismatched request types are helpful.
439-441: Minor helper addition is fineSetPromptTokens is small but useful for post-processing updates.
235-247: Incorrect — ChannelCloudflare constant exists; no change neededSearch shows the constant is already defined:
- constant/channel.go:
ChannelCloudflare = 39The proposed change to
ChannelTypeCloudflareis incorrect — ignore the suggested diff.Likely an incorrect or invalid review comment.
controller/channel-test.go (4)
145-151: Good: handle GenRelayInfo errors with a specific error codeReturning ErrorCodeGenRelayInfoFailed on generator failure improves debuggability.
153-156: Hydrating ChannelMeta early is the right moveExplicitly calling InitChannelMeta before model mapping aligns downstream behavior (adaptor selection, stream options) with channel config.
177-181: Logs now rely on ToString() which masks secrets/PIIUsing info.ToString() avoids leaking ApiKey, TokenKey, and masks user emails. Good correction over manual redaction.
193-206: Embedding conversion path is appropriately specializedUsing ConvertEmbeddingRequest for embeddings and ConvertOpenAIRequest otherwise matches adaptor expectations.
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
relay/gemini_handler.go (1)
178-183: Remove unused inputTexts and loops to avoid compile-time errors.inputTexts is never used; the accumulation loops cause “declared and not used” compilation errors.
- var req any - var err error - var inputTexts []string + var req any + var err error- for _, r := range batchRequest.Requests { - for _, part := range r.Content.Parts { - if part.Text != "" { - inputTexts = append(inputTexts, part.Text) - } - } - } + // parsing of input texts removed (unused)- for _, part := range singleRequest.Content.Parts { - if part.Text != "" { - inputTexts = append(inputTexts, part.Text) - } - } + // parsing of input texts removed (unused)Also applies to: 188-195, 203-207
♻️ Duplicate comments (5)
controller/relay.go (4)
215-219: WS Upgrader allows any origin; restrict to same-origin or an allowlist.Returning true invites cross-site WS hijacking. At minimum, allow only same-host requests (or consult a configured allowlist).
var upgrader = websocket.Upgrader{ - Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol - CheckOrigin: func(r *http.Request) bool { - return true // 允许跨域 - }, + Subprotocols: []string{"realtime"}, // TODO: add other supported protocols if needed + CheckOrigin: func(r *http.Request) bool { + origin := r.Header.Get("Origin") + if origin == "" { + return true + } + // Allow same-host requests. Replace with a proper allowlist from settings if available. + return strings.Contains(origin, r.Host) + }, }
72-80: Fix WS upgrade error path: don't write via nil ws; respond HTTP JSON and safely close.When Upgrade fails, ws is nil. Writing over WS and deferring Close risks nil deref and lost response. Respond via HTTP and only defer Close after a successful upgrade.
- ws, err = upgrader.Upgrade(c.Writer, c.Request, nil) - if err != nil { - helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError()) - return - } - defer ws.Close() + ws, err = upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + // Cannot write via WS when upgrade fails; respond HTTP JSON + openaiErr := types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError() + c.JSON(http.StatusBadRequest, gin.H{"error": openaiErr}) + return + } + defer func() { _ = ws.Close() }()
115-122: Sensitive-words check: wrong signature usage and nil err in NewError.
- The check ignores errors and uses an unrelated err (likely nil), producing an empty NewAPIError and potentially no message.
- Detection should hinge on non-empty words, not an error condition.
- if setting.ShouldCheckPromptSensitive() { - contains, words := service.CheckSensitiveText(meta.CombineText) - if contains { - logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", "))) - newAPIError = types.NewError(err, types.ErrorCodeSensitiveWordsDetected) - return - } - } + if setting.ShouldCheckPromptSensitive() { + words, checkErr := service.CheckSensitiveText(meta.CombineText) + if checkErr != nil { + logger.LogError(c, fmt.Sprintf("sensitive check failed: %v", checkErr)) + } else if len(words) > 0 { + logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", "))) + newAPIError = types.NewError(fmt.Errorf("sensitive words detected"), types.ErrorCodeSensitiveWordsDetected) + return + } + }
136-146: Pre-consume quota: shadowing newAPIError causes silent failure; propagate correctly and guard defer.The short variable declaration shadows newAPIError, so the deferred error responder never runs. Also guard ReturnPreConsumedQuota by downstream failure and actual pre-consumption.
- preConsumedQuota, 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) - } - }() + preConsumedQuota, preConsumeErr := service.PreConsumeQuota(c, priceData.ShouldPreConsumedQuota, relayInfo) + if preConsumeErr != nil { + newAPIError = preConsumeErr + return + } + relayInfo.FinalPreConsumedQuota = preConsumedQuota + defer func() { + // Only return quota if downstream failed and quota was actually pre-consumed + if newAPIError != nil && preConsumedQuota > 0 { + service.ReturnPreConsumedQuota(c, relayInfo, preConsumedQuota) + } + }()relay/common/relay_info.go (1)
295-297: Compile error: RelayInfo has no field SupportStreamOptions; it's in ChannelMeta now.This assignment will not compile. Default (false) is already set in ChannelMeta.Init; override it there if needed.
- info.SupportStreamOptions = false
🧹 Nitpick comments (1)
controller/relay.go (1)
333-334: Use unified logger instead of stdlib log for Midjourney.Aligns with the new centralized logger and avoids mixed logging styles.
- //err = relayMidjourneySubmit(c, relayMode) - log.Println(mjErr) + //err = relayMidjourneySubmit(c, relayMode) + if mjErr != nil { + logger.LogDebug(c, fmt.Sprintf("midjourney response: %+v", mjErr)) + }Note: remove the now-unused "log" import accordingly.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (11)
common/str.go(3 hunks)controller/relay.go(6 hunks)logger/logger.go(1 hunks)relay/channel/xunfei/relay-xunfei.go(3 hunks)relay/claude_handler.go(5 hunks)relay/common/relay_info.go(6 hunks)relay/embedding_handler.go(3 hunks)relay/gemini_handler.go(10 hunks)relay/helper/common.go(3 hunks)service/sensitive.go(1 hunks)types/error.go(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (6)
- types/error.go
- relay/helper/common.go
- relay/claude_handler.go
- logger/logger.go
- common/str.go
- relay/channel/xunfei/relay-xunfei.go
🧰 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 (4)
relay/gemini_handler.go (15)
relay/common/relay_info.go (1)
RelayInfo(73-118)types/error.go (3)
NewAPIError(81-89)NewError(181-193)NewErrorWithStatusCode(215-231)dto/gemini.go (1)
GeminiChatRequest(12-18)common/logger.go (1)
FatalLog(20-24)relay/helper/model_mapped.go (1)
ModelMappedHelper(15-93)setting/model_setting/gemini.go (1)
GetGeminiSettings(43-45)relay/channel/gemini/relay-gemini.go (1)
ThinkingAdaptor(120-178)relay/relay_adaptor.go (1)
GetAdaptor(42-106)setting/model_setting/global.go (1)
GetGlobalSettings(24-26)common/gin.go (1)
GetRequestBody(15-27)common/json.go (1)
Unmarshal(8-10)logger/logger.go (2)
LogDebug(67-71)LogError(63-65)relay/channel/api_request.go (1)
DoRequest(207-209)service/error.go (1)
ResetStatusCode(112-129)dto/openai_response.go (1)
Usage(217-230)
controller/relay.go (19)
relay/common/relay_info.go (2)
RelayInfo(73-118)GenRelayInfo(402-435)relay/image_handler.go (1)
ImageHelper(22-129)relay/audio_handler.go (1)
AudioHelper(16-67)relay/embedding_handler.go (1)
EmbeddingHelper(18-72)relay/relay-text.go (1)
TextHelper(27-184)dto/request_common.go (1)
Request(8-11)relay/gemini_handler.go (2)
GeminiEmbeddingHandler(172-266)GeminiHelper(53-170)types/relay_format.go (5)
RelayFormat(3-3)RelayFormatOpenAIRealtime(12-12)RelayFormatClaude(7-7)RelayFormatGemini(8-8)RelayFormatMjProxy(17-17)common/constants.go (1)
RequestIdKey(129-129)relay/helper/valid_request.go (1)
GetAndValidateRequest(17-44)logger/logger.go (3)
LogWarn(59-61)LogError(63-65)LogInfo(55-57)service/token_counter.go (1)
CountRequestToken(246-288)relay/helper/price.go (1)
ModelPriceHelper(44-102)service/pre_consume_quota.go (1)
PreConsumeQuota(32-78)common/gin.go (1)
GetRequestBody(15-27)relay/claude_handler.go (1)
ClaudeHelper(20-131)constant/env.go (1)
ErrorLogEnabled(15-15)model/log.go (1)
RecordErrorLog(95-135)service/channel.go (2)
ShouldDisableChannel(38-88)DisableChannel(20-27)
relay/embedding_handler.go (9)
relay/common/relay_info.go (1)
RelayInfo(73-118)types/error.go (2)
NewAPIError(81-89)NewError(181-193)dto/request_common.go (1)
Request(8-11)dto/embedding.go (1)
EmbeddingRequest(21-32)common/logger.go (1)
FatalLog(20-24)relay/helper/model_mapped.go (1)
ModelMappedHelper(15-93)relay/relay_adaptor.go (1)
GetAdaptor(42-106)relay/channel/api_request.go (1)
DoRequest(207-209)service/error.go (1)
ResetStatusCode(112-129)
relay/common/relay_info.go (12)
dto/channel_settings.go (2)
ChannelSettings(3-10)ChannelOtherSettings(12-14)dto/realtime.go (1)
RealTimeTool(66-71)types/relay_format.go (12)
RelayFormat(3-3)RelayFormatOpenAIRealtime(12-12)RelayFormatClaude(7-7)RelayFormatRerank(13-13)RelayFormatOpenAIAudio(10-10)RelayFormatEmbedding(14-14)RelayFormatOpenAIResponses(9-9)RelayFormatGemini(8-8)RelayFormatOpenAIImage(11-11)RelayFormatOpenAI(6-6)RelayFormatTask(16-16)RelayFormatMjProxy(17-17)types/price_data.go (1)
PriceData(11-21)dto/request_common.go (1)
Request(8-11)common/gin.go (6)
GetContextKeyInt(64-66)GetContextKeyStringMap(76-78)GetContextKeyBool(68-70)GetContextKeyString(60-62)GetContextKeyType(84-92)GetContextKeyTime(80-82)constant/context_key.go (20)
ContextKeyChannelType(25-25)ContextKeyChannelParamOverride(28-28)ContextKeyChannelId(21-21)ContextKeyChannelIsMultiKey(33-33)ContextKeyChannelMultiKeyIndex(34-34)ContextKeyChannelBaseUrl(24-24)ContextKeyChannelKey(35-35)ContextKeyOriginalModel(8-8)ContextKeyChannelSetting(26-26)ContextKeyChannelOtherSetting(27-27)ContextKeyRequestStartTime(9-9)ContextKeyUserId(38-38)ContextKeyUsingGroup(44-44)ContextKeyUserGroup(43-43)ContextKeyUserQuota(40-40)ContextKeyUserEmail(42-42)ContextKeyPromptTokens(6-6)ContextKeyTokenId(14-14)ContextKeyTokenKey(13-13)ContextKeyTokenUnlimited(12-12)common/api_type.go (1)
ChannelType2APIType(5-75)common/str.go (1)
MaskEmail(104-118)dto/rerank.go (1)
RerankRequest(10-18)relay/constant/relay_mode.go (3)
RelayModeRerank(46-46)RelayModeResponses(48-48)Path2RelayMode(55-89)dto/openai_request.go (1)
OpenAIResponsesRequest(754-776)
🔇 Additional comments (5)
relay/embedding_handler.go (1)
22-25: LGTM: Corrected type-assertion error message.The fatal log now reflects the expected pointer type (*dto.EmbeddingRequest).
service/sensitive.go (4)
34-36: LGTM: Function signature update aligns with new API design.The updated
CheckSensitiveTextfunction now returns(bool, []string)instead of([]string, error), which provides a cleaner API by separating the detection result from the matched words list. The delegation toSensitiveWordContainsmaintains consistency across the sensitive word detection logic.
38-48: LGTM: Well-designed centralized detection function.The
SensitiveWordContainsfunction provides a solid foundation for sensitive word detection with proper edge case handling (empty settings and text) and consistent lowercase normalization before matching.
50-76: LGTM: Robust replacement implementation with efficient string building.The
SensitiveWordReplacefunction demonstrates good practices:
- Proper capacity pre-allocation for the words slice
- Efficient string building with
strings.BuilderandGrow()- Correct position tracking to preserve original text structure
- Clear separation of detection and replacement logic
26-28: All callers updated for the new SensitiveWordContains API — no action requiredVerified the codebase: CheckSensitiveText now delegates to SensitiveWordContains (bool, []string) and all callers handle the new signature. No usages of the removed CheckSensitiveInput remain.
Files checked:
- service/sensitive.go
- lines ~26-28: if ok, words := SensitiveWordContains(m.Text); ok { return words, errors.New("sensitive words detected") }
- lines ~34-36: func CheckSensitiveText(text string) (bool, []string) { return SensitiveWordContains(text) }
- controller/relay.go
- lines ~116-119: contains, words := service.CheckSensitiveText(meta.CombineText) — caller already destructures (bool, []string)
| usage, openaiErr := adaptor.DoResponse(c, resp.(*http.Response), info) | ||
| if openaiErr != nil { | ||
| service.ResetStatusCode(openaiErr, statusCodeMappingStr) | ||
| return openaiErr | ||
| } | ||
|
|
||
| postConsumeQuota(c, relayInfo, usage.(*dto.Usage), preConsumedQuota, userQuota, priceData, "") | ||
| postConsumeQuota(c, info, usage.(*dto.Usage), "") | ||
| return nil |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Avoid potential panic: don't cast resp without nil-check; use httpResp variable.
If adaptor returns a nil response with nil error, resp.(*http.Response) will panic. Use the httpResp var you already set (or guard nil).
- usage, openaiErr := adaptor.DoResponse(c, resp.(*http.Response), info)
+ usage, openaiErr := adaptor.DoResponse(c, httpResp, info)📝 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.
| usage, openaiErr := adaptor.DoResponse(c, resp.(*http.Response), info) | |
| if openaiErr != nil { | |
| service.ResetStatusCode(openaiErr, statusCodeMappingStr) | |
| return openaiErr | |
| } | |
| postConsumeQuota(c, relayInfo, usage.(*dto.Usage), preConsumedQuota, userQuota, priceData, "") | |
| postConsumeQuota(c, info, usage.(*dto.Usage), "") | |
| return nil | |
| usage, openaiErr := adaptor.DoResponse(c, httpResp, info) | |
| if openaiErr != nil { | |
| service.ResetStatusCode(openaiErr, statusCodeMappingStr) | |
| return openaiErr | |
| } | |
| postConsumeQuota(c, info, usage.(*dto.Usage), "") | |
| return nil |
🤖 Prompt for AI Agents
In relay/gemini_handler.go around lines 162 to 169, the code type-asserts resp
to *http.Response without a nil-check which can panic if resp is nil or not the
expected type; use the already-populated httpResp variable (or perform a safe
type-assert with the ok idiom) when calling adaptor.DoResponse, and add a guard
to return an error if httpResp is nil before proceeding; also safely assert
usage to *dto.Usage (check for nil/type-assertion success) before passing it to
postConsumeQuota to avoid further panics.
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (6)
common/sys_log.go (4)
10-13: Trim trailing spaces before newline in logs.There is an extra space before the newline in both
SysLogandSysError. Removing it avoids visual artifacts and is consistent with typical log formats.func SysLog(s string) { t := time.Now() - _, _ = fmt.Fprintf(gin.DefaultWriter, "[SYS] %v | %s \n", t.Format("2006/01/02 - 15:04:05"), s) + _, _ = fmt.Fprintf(gin.DefaultWriter, "[SYS] %v | %s\n", t.Format("2006/01/02 - 15:04:05"), s) } func SysError(s string) { t := time.Now() - _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[SYS] %v | %s \n", t.Format("2006/01/02 - 15:04:05"), s) + _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[SYS] %v | %s\n", t.Format("2006/01/02 - 15:04:05"), s) }Also applies to: 15-18
15-18: Consider adding severity to SysError prefix.Right now both
SysLogandSysErroremit[SYS]. For log parsing/alerting, consider a distinct severity marker like[SYS][ERROR](or align with your new logger’s levels).
11-12: Extract timestamp layout to a const and (optionally) use UTC.The layout string is duplicated. A module-level
const logTimeLayout = "2006/01/02 - 15:04:05"improves consistency. Also considertime.Now().UTC()for consistent cross-node timestamps if logs are aggregated.Also applies to: 16-17, 21-22
3-8: Align with the new logger package for consistency and future extensibility.Given the PR’s logger migration, consider making these helpers thin wrappers around the new
one-api/logger(with a “system” channel) instead of writing directly to gin’s writers. That centralizes formatting, levels, destinations, and JSON/structured output.Also applies to: 10-24
dto/claude.go (1)
324-326: Mark unused parameter in IsStream to appease lintersThe gin.Context is required by the interface but unused. Rename to underscore.
-func (claudeRequest *ClaudeRequest) IsStream(c *gin.Context) bool { +func (claudeRequest *ClaudeRequest) IsStream(_ *gin.Context) bool { return claudeRequest.Stream }service/token_counter.go (1)
79-81: Clarify error message for nil image metaThe function consumes FileMeta rather than a URL; make the error accurate.
- if fileMeta == nil { - return 0, fmt.Errorf("image_url_is_nil") - } + if fileMeta == nil { + return 0, fmt.Errorf("image_meta_is_nil") + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
common/sys_log.go(1 hunks)dto/claude.go(3 hunks)service/token_counter.go(3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (2)
dto/claude.go (3)
common/json.go (1)
Marshal(20-22)types/request_meta.go (6)
TokenCountMeta(20-31)TokenType(12-12)TokenTypeTokenizer(16-16)FileMeta(33-38)FileType(3-3)FileTypeImage(6-6)common/utils.go (1)
Interface2String(118-128)
service/token_counter.go (9)
types/request_meta.go (4)
FileMeta(33-38)TokenCountMeta(20-31)TokenType(12-12)FileType(3-3)constant/env.go (2)
GetMediaTokenNotStream(8-8)GetMediaToken(7-7)service/image.go (2)
DecodeUrlImageData(113-153)DecodeBase64ImageData(18-35)common/sys_log.go (1)
SysLog(10-13)common/constants.go (1)
DebugEnabled(70-70)relay/common/relay_info.go (1)
RelayInfo(73-118)types/relay_format.go (4)
RelayFormat(3-3)RelayFormatOpenAIRealtime(12-12)RelayFormatOpenAI(6-6)RelayFormatGemini(8-8)common/gin.go (2)
GetContextKeyString(60-62)SetContextKey(52-54)constant/context_key.go (2)
ContextKeyOriginalModel(8-8)ContextKeyPromptTokens(6-6)
🔇 Additional comments (2)
dto/claude.go (1)
87-89: Good switch to common.Marshal for consistencyUsing common.Marshal centralizes JSON handling and keeps behavior aligned across DTOs.
service/token_counter.go (1)
246-294: Meta-based counting flow looks solid; context propagation is a plusThe CountRequestToken surface correctly consumes TokenCountMeta, adds OpenAI framing costs, handles files, and writes prompt tokens to context. This aligns with the PR’s centralized counting objective.
| func FatalLog(v ...any) { | ||
| t := time.Now() | ||
| _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %v \n", t.Format("2006/01/02 - 15:04:05"), v) | ||
| os.Exit(1) | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Fix FatalLog argument formatting (variadic slice is logged as a slice).
Passing v ...any into %v prints the slice (e.g., [a 1]) instead of a space-separated message. Use fmt.Sprint(v...) (or provide a FatalLogf) and drop the extra space before newline.
func FatalLog(v ...any) {
t := time.Now()
- _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %v \n", t.Format("2006/01/02 - 15:04:05"), v)
+ _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %s\n", t.Format("2006/01/02 - 15:04:05"), fmt.Sprint(v...))
os.Exit(1)
}📝 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 FatalLog(v ...any) { | |
| t := time.Now() | |
| _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %v \n", t.Format("2006/01/02 - 15:04:05"), v) | |
| os.Exit(1) | |
| } | |
| func FatalLog(v ...any) { | |
| t := time.Now() | |
| _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %s\n", t.Format("2006/01/02 - 15:04:05"), fmt.Sprint(v...)) | |
| os.Exit(1) | |
| } |
🤖 Prompt for AI Agents
In common/sys_log.go around lines 20 to 24, FatalLog currently prints the
variadic argument slice as a slice because it passes v directly to %v and leaves
an extra space before the newline; change the log formatting to convert the
variadic args into a single string (e.g., use fmt.Sprint(v...) or equivalent)
and remove the extra space before the newline so the message is space-separated
and the line ends cleanly.
💡 Verification agent
🧩 Analysis chain
Validate FatalLog usage; os.Exit in library code can kill the whole process unexpectedly.
os.Exit(1) bypasses defers and will terminate the server if called in a request path. Ensure it’s only used during startup or unrecoverable init failures.
Run this script to find call sites and check for risky usage within HTTP handlers:
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== FatalLog call sites (with context) =="
rg -nC3 --type=go '\bFatalLog\s*\(' || true
echo
echo "== Files containing FatalLog and potential HTTP handler hints (gin.Context, ServeHTTP, HandleFunc) =="
files=$(rg -n --type=go '\bFatalLog\s*\(' | cut -d: -f1 | sort -u)
for f in $files; do
echo "---- $f"
rg -n -C2 --type=go -e '\bgin\.Context\b' -e '\bServeHTTP\b' -e '\bHandleFunc\b' "$f" || true
doneLength of output: 11071
Do not call os.Exit from FatalLog — it will kill the whole server when used in request handlers
os.Exit(1) in common/sys_log.go will bypass defers and terminate the process. The repo shows FatalLog is invoked from request-handling code, so this will bring down the server on malformed requests.
Call sites that need immediate fixes (called from HTTP handlers / helper funcs with *gin.Context):
- relay/claude_handler.go:27 — common.FatalLog(...) inside ClaudeHelper(c *gin.Context)
- relay/embedding_handler.go:24 — common.FatalLog(...) inside EmbeddingHelper(c *gin.Context)
- relay/gemini_handler.go:58 — common.FatalLog(...) inside Gemini helper(s) with *gin.Context
- relay/image_handler.go:29 — common.FatalLog(...) inside ImageHelper(c *gin.Context)
- relay/relay-text.go:35 — common.FatalLog(...) inside TextHelper(c *gin.Context)
- relay/rerank_handler.go:32 — common.FatalLog(...) inside RerankHelper(c *gin.Context)
- relay/responses_handler.go:26 — common.FatalLog(...) inside ResponsesHelper(c *gin.Context)
Other usages (likely startup/init and safer to exit from main if desired):
- common/sys_log.go:20-23 — FatalLog implementation uses os.Exit(1)
- common/redis.go:37,47,59 — FatalLog used during Redis init
- main.go:39,54,74,93,100,159,187 — FatalLog used in init/cleanup/startup paths
- model/main.go:223,260 — FatalLog used during migrations/init
Recommendation (concise):
- Remove os.Exit from FatalLog. Make FatalLog only log (and optionally panic), and handle termination explicitly from main() during startup.
- For request handlers, replace FatalLog calls with proper error handling: return an API error, log, and use gin.Context to abort/respond (do not exit).
- Alternatively, create two functions: FatalInitExit(...) used only in main/init (calls os.Exit) and FatalLogRuntime(...) that logs without exiting.
Please address the relay/* call sites first — they are running in request paths and currently can terminate the whole server.
🤖 Prompt for AI Agents
In common/sys_log.go around lines 20-24, FatalLog currently calls os.Exit(1)
which will terminate the whole server when used in request handlers; change
FatalLog to only log (no os.Exit) and add a new function (e.g., FatalInitExit or
FatalLogExit) used exclusively during startup/init that logs then calls
os.Exit(1). Update the relay/* request-handler call sites listed in the review
to use the non-exiting logger (or return/abort with gin.Context and log the
error) instead of FatalLog, and update startup/init call sites to use the
exiting variant where termination is appropriate.
| func (c *ClaudeRequest) GetTokenCountMeta() *types.TokenCountMeta { | ||
| var tokenCountMeta = types.TokenCountMeta{ | ||
| TokenType: types.TokenTypeTokenizer, | ||
| MaxTokens: int(c.MaxTokens), | ||
| } | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Respect both MaxTokens and MaxTokensToSample when populating meta
Claude requests may use either field; currently only MaxTokens is propagated, potentially undercounting.
Apply this diff to pick the non-zero value from both fields:
-func (c *ClaudeRequest) GetTokenCountMeta() *types.TokenCountMeta {
- var tokenCountMeta = types.TokenCountMeta{
- TokenType: types.TokenTypeTokenizer,
- MaxTokens: int(c.MaxTokens),
- }
+func (c *ClaudeRequest) GetTokenCountMeta() *types.TokenCountMeta {
+ maxTokens := int(c.MaxTokens)
+ if maxTokens == 0 && c.MaxTokensToSample != 0 {
+ maxTokens = int(c.MaxTokensToSample)
+ }
+ var tokenCountMeta = types.TokenCountMeta{
+ TokenType: types.TokenTypeTokenizer,
+ MaxTokens: maxTokens,
+ }📝 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 (c *ClaudeRequest) GetTokenCountMeta() *types.TokenCountMeta { | |
| var tokenCountMeta = types.TokenCountMeta{ | |
| TokenType: types.TokenTypeTokenizer, | |
| MaxTokens: int(c.MaxTokens), | |
| } | |
| func (c *ClaudeRequest) GetTokenCountMeta() *types.TokenCountMeta { | |
| maxTokens := int(c.MaxTokens) | |
| if maxTokens == 0 && c.MaxTokensToSample != 0 { | |
| maxTokens = int(c.MaxTokensToSample) | |
| } | |
| var tokenCountMeta = types.TokenCountMeta{ | |
| TokenType: types.TokenTypeTokenizer, | |
| MaxTokens: maxTokens, | |
| } |
🤖 Prompt for AI Agents
In dto/claude.go around lines 205 to 210, the TokenCountMeta.MaxTokens is
currently set only from c.MaxTokens which ignores c.MaxTokensToSample; update
the assignment to use the non-zero value between the two fields (e.g., if
c.MaxTokensToSample != 0 use that, otherwise use c.MaxTokens), converting to
int, so the meta reflects the actual sampling limit.
| data := media.Source.Url | ||
| if data == "" { | ||
| data = common.Interface2String(media.Source.Data) | ||
| } | ||
| if data != "" { | ||
| fileMeta = append(fileMeta, &types.FileMeta{FileType: types.FileTypeImage, Data: data}) | ||
| } | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Avoid “Not Implemented” sentinel from Interface2String contaminating Files; assert string type explicitly
When Source.Data is not a string, common.Interface2String returns "Not Implemented", which would be treated as a valid file data and skew image-token counting. Prefer a type assertion to string and skip otherwise.
Apply this diff at both occurrences:
- if data == "" {
- data = common.Interface2String(media.Source.Data)
- }
+ if data == "" {
+ if s, ok := media.Source.Data.(string); ok {
+ data = s
+ }
+ }Also applies to: 262-268
🤖 Prompt for AI Agents
In dto/claude.go around lines 229-236 (and also apply same change at lines
262-268), the code uses common.Interface2String which returns the sentinel "Not
Implemented" for non-string values and can contaminate file metadata; instead
perform a direct type assertion: check whether media.Source.Data is a string (s,
ok := media.Source.Data.(string)) and only assign/use s when ok is true,
otherwise skip adding the FileMeta entry; remove use of Interface2String in
these locations so non-string Data values are ignored.
|
|
||
| // messages | ||
| for _, message := range c.Messages { | ||
| tokenCountMeta.MessagesCount++ |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Include Prompt in combined text for completion-style requests
Prompt content is currently ignored, undercounting tokens for legacy/completion usage.
Insert this snippet before processing messages:
}
- // messages
+ // prompt (completion-style)
+ if c.Prompt != "" {
+ texts = append(texts, c.Prompt)
+ }
+
+ // messages📝 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.
| // messages | |
| for _, message := range c.Messages { | |
| tokenCountMeta.MessagesCount++ | |
| // prompt (completion-style) | |
| if c.Prompt != "" { | |
| texts = append(texts, c.Prompt) | |
| } | |
| // messages | |
| for _, message := range c.Messages { | |
| tokenCountMeta.MessagesCount++ |
🤖 Prompt for AI Agents
In dto/claude.go around lines 241 to 244, the code iterates only over c.Messages
and misses including the top-level Prompt for legacy/completion-style requests,
undercounting tokens; modify the code to check if c.Prompt (or the field name
used) is non-empty and append it to the combined text (and account for it in
token count) before looping messages so the prompt text is included in the
tokenization and counting logic.
| func getImageToken(fileMeta *types.FileMeta, model string, stream bool) (int, error) { | ||
| if fileMeta == nil { | ||
| return 0, fmt.Errorf("image_url_is_nil") | ||
| } | ||
|
|
||
| // Defaults for 4o/4.1/4.5 family unless overridden below | ||
| baseTokens := 85 | ||
| if model == "glm-4v" { | ||
| tileTokens := 170 | ||
|
|
||
| // Model classification | ||
| lowerModel := strings.ToLower(model) | ||
|
|
||
| // Special cases from existing behavior | ||
| if strings.HasPrefix(lowerModel, "glm-4") { | ||
| return 1047, nil | ||
| } | ||
| if imageUrl.Detail == "low" { | ||
|
|
||
| // Patch-based models (32x32 patches, capped at 1536, with multiplier) | ||
| isPatchBased := false | ||
| multiplier := 1.0 | ||
| switch { | ||
| case strings.Contains(lowerModel, "gpt-4.1-mini"): | ||
| isPatchBased = true | ||
| multiplier = 1.62 | ||
| case strings.Contains(lowerModel, "gpt-4.1-nano"): | ||
| isPatchBased = true | ||
| multiplier = 2.46 | ||
| case strings.HasPrefix(lowerModel, "o4-mini"): | ||
| isPatchBased = true | ||
| multiplier = 1.72 | ||
| case strings.HasPrefix(lowerModel, "gpt-5-mini"): | ||
| isPatchBased = true | ||
| multiplier = 1.62 | ||
| case strings.HasPrefix(lowerModel, "gpt-5-nano"): | ||
| isPatchBased = true | ||
| multiplier = 2.46 | ||
| } | ||
|
|
||
| // Tile-based model tokens and bases per doc | ||
| if !isPatchBased { | ||
| if strings.HasPrefix(lowerModel, "gpt-4o-mini") { | ||
| baseTokens = 2833 | ||
| tileTokens = 5667 | ||
| } else if strings.HasPrefix(lowerModel, "gpt-5-chat-latest") || (strings.HasPrefix(lowerModel, "gpt-5") && !strings.Contains(lowerModel, "mini") && !strings.Contains(lowerModel, "nano")) { | ||
| baseTokens = 70 | ||
| tileTokens = 140 | ||
| } else if strings.HasPrefix(lowerModel, "o1") || strings.HasPrefix(lowerModel, "o3") || strings.HasPrefix(lowerModel, "o1-pro") { | ||
| baseTokens = 75 | ||
| tileTokens = 150 | ||
| } else if strings.Contains(lowerModel, "computer-use-preview") { | ||
| baseTokens = 65 | ||
| tileTokens = 129 | ||
| } else if strings.Contains(lowerModel, "4.1") || strings.Contains(lowerModel, "4o") || strings.Contains(lowerModel, "4.5") { | ||
| baseTokens = 85 | ||
| tileTokens = 170 | ||
| } | ||
| } | ||
|
|
||
| // Respect existing feature flags/short-circuits | ||
| if fileMeta.Detail == "low" && !isPatchBased { | ||
| return baseTokens, nil | ||
| } | ||
| if !constant.GetMediaTokenNotStream && !stream { | ||
| return 3 * baseTokens, nil | ||
| } | ||
|
|
||
| // 同步One API的图片计费逻辑 | ||
| if imageUrl.Detail == "auto" || imageUrl.Detail == "" { | ||
| imageUrl.Detail = "high" | ||
| // Normalize detail | ||
| if fileMeta.Detail == "auto" || fileMeta.Detail == "" { | ||
| fileMeta.Detail = "high" | ||
| } | ||
|
|
||
| tileTokens := 170 | ||
| if strings.HasPrefix(model, "gpt-4o-mini") { | ||
| tileTokens = 5667 | ||
| baseTokens = 2833 | ||
| } | ||
| // 是否统计图片token | ||
| // Whether to count image tokens at all | ||
| if !constant.GetMediaToken { | ||
| return 3 * baseTokens, nil | ||
| } | ||
| if info.ChannelType == constant.ChannelTypeGemini || info.ChannelType == constant.ChannelTypeVertexAi || info.ChannelType == constant.ChannelTypeAnthropic { | ||
| return 3 * baseTokens, nil | ||
| } | ||
|
|
||
| // Decode image to get dimensions | ||
| var config image.Config | ||
| var err error | ||
| var format string | ||
| var b64str string | ||
| if strings.HasPrefix(imageUrl.Url, "http") { | ||
| config, format, err = DecodeUrlImageData(imageUrl.Url) | ||
| if strings.HasPrefix(fileMeta.Data, "http") { | ||
| config, format, err = DecodeUrlImageData(fileMeta.Data) | ||
| } else { | ||
| common.SysLog(fmt.Sprintf("decoding image")) | ||
| config, format, b64str, err = DecodeBase64ImageData(imageUrl.Url) | ||
| config, format, b64str, err = DecodeBase64ImageData(fileMeta.Data) | ||
| } | ||
| if err != nil { | ||
| return 0, err | ||
| } | ||
| imageUrl.MimeType = format | ||
| fileMeta.MimeType = format | ||
|
|
||
| if config.Width == 0 || config.Height == 0 { | ||
| // not an image | ||
| if format != "" && b64str != "" { | ||
| // file type | ||
| return 3 * baseTokens, nil | ||
| } | ||
| return 0, errors.New(fmt.Sprintf("fail to decode base64 config: %s", imageUrl.Url)) | ||
| } | ||
|
|
||
| shortSide := config.Width | ||
| otherSide := config.Height | ||
| log.Printf("format: %s, width: %d, height: %d", format, config.Width, config.Height) | ||
| // 缩放倍数 | ||
| scale := 1.0 | ||
| if config.Height < shortSide { | ||
| shortSide = config.Height | ||
| otherSide = config.Width | ||
| } | ||
|
|
||
| // 将最小变的尺寸缩小到768以下，如果大于768，则缩放到768 | ||
| if shortSide > 768 { | ||
| scale = float64(shortSide) / 768 | ||
| shortSide = 768 | ||
| } | ||
| // 将另一边按照相同的比例缩小，向上取整 | ||
| otherSide = int(math.Ceil(float64(otherSide) / scale)) | ||
| log.Printf("shortSide: %d, otherSide: %d, scale: %f", shortSide, otherSide, scale) | ||
| // 计算图片的token数量(边的长度除以512，向上取整) | ||
| tiles := (shortSide + 511) / 512 * ((otherSide + 511) / 512) | ||
| log.Printf("tiles: %d", tiles) | ||
| return 0, errors.New(fmt.Sprintf("fail to decode base64 config: %s", fileMeta.Data)) | ||
| } | ||
|
|
||
| width := config.Width | ||
| height := config.Height | ||
| log.Printf("format: %s, width: %d, height: %d", format, width, height) | ||
|
|
||
| if isPatchBased { | ||
| // 32x32 patch-based calculation with 1536 cap and model multiplier | ||
| ceilDiv := func(a, b int) int { return (a + b - 1) / b } | ||
| rawPatchesW := ceilDiv(width, 32) | ||
| rawPatchesH := ceilDiv(height, 32) | ||
| rawPatches := rawPatchesW * rawPatchesH | ||
| if rawPatches > 1536 { | ||
| // scale down | ||
| area := float64(width * height) | ||
| r := math.Sqrt(float64(32*32*1536) / area) | ||
| wScaled := float64(width) * r | ||
| hScaled := float64(height) * r | ||
| // adjust to fit whole number of patches after scaling | ||
| adjW := math.Floor(wScaled/32.0) / (wScaled / 32.0) | ||
| adjH := math.Floor(hScaled/32.0) / (hScaled / 32.0) | ||
| adj := math.Min(adjW, adjH) | ||
| if !math.IsNaN(adj) && adj > 0 { | ||
| r = r * adj | ||
| } | ||
| wScaled = float64(width) * r | ||
| hScaled = float64(height) * r | ||
| patchesW := math.Ceil(wScaled / 32.0) | ||
| patchesH := math.Ceil(hScaled / 32.0) | ||
| imageTokens := int(patchesW * patchesH) | ||
| if imageTokens > 1536 { | ||
| imageTokens = 1536 | ||
| } | ||
| return int(math.Round(float64(imageTokens) * multiplier)), nil | ||
| } | ||
| // below cap | ||
| imageTokens := rawPatches | ||
| return int(math.Round(float64(imageTokens) * multiplier)), nil | ||
| } | ||
|
|
||
| // Tile-based calculation for 4o/4.1/4.5/o1/o3/etc. | ||
| // Step 1: fit within 2048x2048 square | ||
| maxSide := math.Max(float64(width), float64(height)) | ||
| fitScale := 1.0 | ||
| if maxSide > 2048 { | ||
| fitScale = maxSide / 2048.0 | ||
| } | ||
| fitW := int(math.Round(float64(width) / fitScale)) | ||
| fitH := int(math.Round(float64(height) / fitScale)) | ||
|
|
||
| // Step 2: scale so that shortest side is exactly 768 | ||
| minSide := math.Min(float64(fitW), float64(fitH)) | ||
| if minSide == 0 { | ||
| return baseTokens, nil | ||
| } | ||
| shortScale := 768.0 / minSide | ||
| finalW := int(math.Round(float64(fitW) * shortScale)) | ||
| finalH := int(math.Round(float64(fitH) * shortScale)) | ||
|
|
||
| // Count 512px tiles | ||
| tilesW := (finalW + 512 - 1) / 512 | ||
| tilesH := (finalH + 512 - 1) / 512 | ||
| tiles := tilesW * tilesH | ||
|
|
||
| if common.DebugEnabled { | ||
| log.Printf("scaled to: %dx%d, tiles: %d", finalW, finalH, tiles) | ||
| } | ||
|
|
||
| return tiles*tileTokens + baseTokens, nil | ||
| } |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Double-check tile/patch constants and model classification
The base/tile constants and patch multipliers (e.g., 85/170, 2833/5667, multipliers for 4.1-mini/nano, o4-mini, gpt-5-*) are critical. Please confirm they align with the latest vendor guidance, and consider centralizing them in a lookup to avoid drift.
If confirmed, consider refactoring model constants into a map or small struct to document the source and simplify maintenance.
🌐 Web query:
What are the current image token calculation rules for OpenAI GPT-4o/4.1/4.5/o-series (tile size, base tokens, tiles per 512px, scaling to 2048px and 768px shortest side)? Include references for mini/nano variants if available.
💡 Result:
Short answer (current as of 2025-08-15)
- Two different image-token schemes are used across OpenAI models:
- GPT‑4o / o‑series (gpt-4o, o-series): tile-based 512×512 tiles, plus a small fixed base. For detail:high images the model: (a) first scales to fit inside a 2048×2048 square (keep aspect ratio), (b) then scales so the shortest side is 768px (if applicable), (c) counts how many 512×512 tiles cover the resulting image, (d) charges one “base” token amount + a per‑tile token amount × number_of_tiles. Images with detail:low are a fixed small cost. (openai-hd4n6.mintlify.app, openai.com)
- GPT‑4.1 family (and variants): patch‑based. The image is treated as 32×32 px patches (patch_count = ceil(w/32) * ceil(h/32)), tokens = #patches (capped at 1536). Mini/nano variants apply multipliers to the image‑token count (see below). (openai-hd4n6.mintlify.app)
Details, formulas and examples
- gpt‑4o / o‑series (tile rules)
- Tile size: 512 × 512 px. (openai-hd4n6.mintlify.app)
- Resize steps (detail: "high"):
- If either dimension > 2048, scale the image to fit inside a 2048×2048 square (keep aspect ratio). (openai-hd4n6.mintlify.app)
- Then (if applicable) scale the image so its shortest side = 768 px. (Effectively the model will see an image whose shortest side is ≤ 768 after this step.) (openai-hd4n6.mintlify.app)
- Tile count = ceil(resized_width / 512) × ceil(resized_height / 512).
- Token cost (typical/generic formula historically shown in docs): tokens = base_tokens + (tile_token_cost × tile_count). Historically base_tokens ≈ 85 and tile_token_cost ≈ 170 for many of the large gpt‑4o / o‑series examples in the docs, so e.g. a 1024×1024 (high) → resized → 768×768 → 2×2 tiles → tokens = 85 + 170×4 = 765. However, the per‑model numbers shown by the pricing calculator vary by model (see “per‑model” note and examples below). (openai-hd4n6.mintlify.app, openai.com)
- gpt‑4.1 family (patch rules)
- Patch size: 32 × 32 px. Tokenized image tokens = number_of_patches = ceil(width/32) × ceil(height/32). There is a hard cap at 1536 image tokens (i.e., if computed patches > 1536 the image is scaled down to fit). Example: 1024×1024 → (1024/32)=32 → 32×32 = 1024 tokens. (openai-hd4n6.mintlify.app)
- Mini / nano variants (model-dependent multipliers or different base/tile values)
- OpenAI’s pricing and docs note that some “mini” / “nano” variants convert images differently and therefore show different base/tile token equivalents in the pricing calculator. The pricing page calls this out explicitly (models like gpt-4.1-mini, gpt-4.1-nano and o4-mini convert images into tokens differently). Use the pricing calculator for exact per‑model numbers. (openai.com)
- Example / community & Azure reports (illustrative):
- gpt‑4.1-mini / gpt‑4.1-nano: the docs show multipliers applied to the raw patch count (for gpt‑4.1-mini multiply image tokens by ~1.62; for gpt‑4.1-nano multiply by ~2.46). (This is called out in the docs' 4.1 section.) (openai-hd4n6.mintlify.app)
- gpt‑4o‑mini has been observed (and is shown in some vendor docs/examples) to present much larger “base tokens” and “tile tokens” values in the pricing calculator (e.g., examples reporting base ≈ 2,833 and tile ≈ 5,667 for that model), which makes the billed token count look large even though the final dollar cost on the pricing page equals the same image dollar amount as other models (because the per‑token price for that model is different). If you care about exact billed token numbers (quota, TPM), check the pricing calculator or your usage dashboard for the specific model. (community.openai.com, learn.microsoft.com)
Practical (compact) formulas you can implement
-
For gpt‑4o / o‑series (detail: high):
- If width>2048 or height>2048 → scale to fit inside 2048×2048 (preserve aspect ratio).
- If both sides > 768 (or shortest side > 768) → scale so shortest side = 768 (preserve aspect ratio).
- tiles_x = ceil(width / 512) ; tiles_y = ceil(height / 512) ; tiles = tiles_x * tiles_y.
- tokens = base_tokens + tile_tokens_per_tile * tiles. (Docs examples: base_tokens = 85, tile_tokens_per_tile = 170 for the canonical large models; per‑model values shown on pricing page may differ.) (openai-hd4n6.mintlify.app, openai.com)
-
For gpt‑4.1:
- patches_x = ceil(width / 32); patches_y = ceil(height / 32).
- image_tokens = patches_x * patches_y (cap at 1536).
- If using 4.1-mini / 4.1-nano, apply the documented multiplier (e.g., ×1.62 for mini, ×2.46 for nano) then bill at the model’s text token rate. (openai-hd4n6.mintlify.app)
Where to get the exact per‑model numbers (recommended)
- The OpenAI API Pricing page includes a per‑model image cost calculator and explicitly states that some mini/nano variants convert images differently — use that interactive calculator or your account usage dashboard for exact base/tile token numbers and dollar estimates per model. (The site also shows worked examples for particular models.) (openai.com)
Sources (official / authoritative)
- OpenAI docs: “Images and Vision” / calculating costs (document describing the 2048→768→512‑tiling rules for o‑series and the 32×32 patch rules for gpt‑4.1). (mirror / doc copy used in my lookup). (openai-hd4n6.mintlify.app)
- OpenAI API Pricing (pricing page + image cost calculator) — shows per‑model differences and calls out that mini/nano/or other specific models may convert images differently (use the calculator for exact base/tile numbers). (openai.com)
- Microsoft / Azure notes and community reports — corroborate the per‑model oddities for some mini models (examples that gpt‑4o‑mini shows large base/tile token numbers in the calculator). Use these if you need confirmation for Azure OpenAI customers. (learn.microsoft.com)
If you want, I can:
- Give you a tiny utility (Python or JS) that implements the gpt‑4o / o‑series tiling rules and prints token counts for any input width/height (and shows how the pricing calculator would present base/tile tokens for selected models), or
- Query the pricing calculator (for a specified model name) and copy the calculator’s current base/tile token numbers for that exact model (so you get the model’s live base/tile token values). Which would help you handle mini/nano differences precisely?
Citations:
- 1: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 2: https://openai.com/api/pricing
- 3: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 4: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 5: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 6: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 7: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 8: https://openai.com/api/pricing
- 9: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 10: https://openai.com/api/pricing
- 11: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 12: https://community.openai.com/t/token-usage-for-images-remains-constant-regardless-of-size-is-this-a-bug/951871?utm_source=chatgpt.com
- 13: https://learn.microsoft.com/en-sg/answers/questions/2126187/azure-open-ai-gpt-4o-mini-image-input-token-limita?utm_source=chatgpt.com
- 14: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 15: https://openai.com/api/pricing
- 16: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 17: https://openai.com/api/pricing
- 18: https://openai-hd4n6.mintlify.app/docs/guides/images?utm_source=chatgpt.com
- 19: https://openai.com/api/pricing
- 20: https://learn.microsoft.com/sr-latn-rs/azure/ai-services/openai/overview?utm_source=chatgpt.com
Fix image-scaling bug and centralize model/token constants
Short: the web-docs confirm the 512‑tile and 32×32‑patch rules, but I found a definite bug (unconditional upscaling to shortest side = 768) and a fragile model-detection/constant layout. Please fix the bug and centralize model constants.
Points to address (service/token_counter.go — getImageToken, ~lines 78–244)
- Incorrect scaling (tile path):
- Problem: code always scales the shortest side to 768 (it upscales small images). Per docs you should only downscale so shortest side = 768 when the shortest side > 768.
- Fix: apply the 768 scaling only when minSide > 768. Example replacement:
minSide := math.Min(float64(fitW), float64(fitH)) if minSide == 0 { return baseTokens, nil } shortScale := 1.0 if minSide > 768.0 { shortScale = 768.0 / minSide } finalW := int(math.Round(float64(fitW) * shortScale)) finalH := int(math.Round(float64(fitH) * shortScale))
- Model classification / typos:
- The patch-based switch contains "o4-mini" (likely a typo / nonstandard prefix). Confirm intended names (gpt-4o-mini, gpt-4.1-mini, etc.) and normalize detection (use exact model strings or regex). Right now the logic is brittle and may misclassify models.
- Constants and multipliers:
- The per-model base/tile values and patch multipliers (85/170, 2833/5667, 1.62/2.46/1.72, gpt-5 values, etc.) align broadly with docs/examples, but these are fragile and model-specific.
- Recommendation: centralize these into a single lookup (map/struct) such as:
{ modelPattern -> {isPatch bool, multiplier float64, baseTokens int, tileTokens int, sourceURL string} }
— add small unit tests with representative widths/heights for each entry and a comment linking to the pricing/docs (verify mini/nano numbers via the pricing calculator).
- Minor: review the patch-scaling adjustment logic (cap handling / integer-patch adjustment) for edge cases and add tests covering >1536 patches.
Result: please make the scaling fix above, consolidate model constants into a single structure (with doc/source), and add unit tests for canonical image sizes + a couple of mini/nano models. After that I can re-check and sign off.
🤖 Prompt for AI Agents
In service/token_counter.go around lines 78–244, the tile-scaling logic
incorrectly upscales images by always forcing the shortest side to 768; change
it to only scale the shortest side when minSide > 768 (use shortScale = 1.0 by
default and set to 768/minSide only when minSide > 768), and return baseTokens
if minSide == 0 as before. Replace the brittle model-detection switch with a
centralized lookup (map[string]struct{isPatch bool; multiplier float64;
baseTokens int; tileTokens int; pattern string}) that normalizes model names
(lowercase, exact prefixes/regex for mini/nano/gpt-4o vs o4 typos) and drives
both patch-vs-tile behavior and token constants. Keep the existing patch-based
cap/adjustment but review edge cases and add unit tests that assert expected
token counts for representative widths/heights (including >1536 patch cap and
small images to ensure no upscaling), and add comments linking to the referenced
docs/pricing for each lookup entry.
| fileMeta.MimeType = format | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Populate MimeType in a standard form (image/)
Decode functions return format like "jpeg"/"png". Assigning that raw to MimeType is misleading; normalize to image/.
- fileMeta.MimeType = format
+ if format != "" && !strings.HasPrefix(format, "image/") {
+ fileMeta.MimeType = "image/" + format
+ } else {
+ fileMeta.MimeType = format
+ }📝 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.
| fileMeta.MimeType = format | |
| if format != "" && !strings.HasPrefix(format, "image/") { | |
| fileMeta.MimeType = "image/" + format | |
| } else { | |
| fileMeta.MimeType = format | |
| } |
🤖 Prompt for AI Agents
In service/token_counter.go around lines 166-167, fileMeta.MimeType is being set
directly to the raw image format string (e.g., "jpeg"/"png"); change this to
normalize to a proper MIME type by prefixing with "image/" (e.g., "image/jpeg"),
ensuring the format is lowercased/trimmed first, and fall back to a safe default
like "application/octet-stream" or an empty string if format is empty or
invalid.
| if config.Width == 0 || config.Height == 0 { | ||
| // not an image | ||
| if format != "" && b64str != "" { | ||
| // file type | ||
| return 3 * baseTokens, nil | ||
| } | ||
| return 0, errors.New(fmt.Sprintf("fail to decode base64 config: %s", imageUrl.Url)) | ||
| } | ||
|
|
||
| shortSide := config.Width | ||
| otherSide := config.Height | ||
| log.Printf("format: %s, width: %d, height: %d", format, config.Width, config.Height) | ||
| // 缩放倍数 | ||
| scale := 1.0 | ||
| if config.Height < shortSide { | ||
| shortSide = config.Height | ||
| otherSide = config.Width | ||
| } | ||
|
|
||
| // 将最小变的尺寸缩小到768以下，如果大于768，则缩放到768 | ||
| if shortSide > 768 { | ||
| scale = float64(shortSide) / 768 | ||
| shortSide = 768 | ||
| } | ||
| // 将另一边按照相同的比例缩小，向上取整 | ||
| otherSide = int(math.Ceil(float64(otherSide) / scale)) | ||
| log.Printf("shortSide: %d, otherSide: %d, scale: %f", shortSide, otherSide, scale) | ||
| // 计算图片的token数量(边的长度除以512，向上取整) | ||
| tiles := (shortSide + 511) / 512 * ((otherSide + 511) / 512) | ||
| log.Printf("tiles: %d", tiles) | ||
| return 0, errors.New(fmt.Sprintf("fail to decode base64 config: %s", fileMeta.Data)) | ||
| } |
There was a problem hiding this comment.
Avoid logging or returning raw image data in errors (PII/large payload risk)
Returning fileMeta.Data can leak Base64 or URLs and bloat logs. Keep the message generic.
- if format != "" && b64str != "" {
- // file type
- return 3 * baseTokens, nil
- }
- return 0, errors.New(fmt.Sprintf("fail to decode base64 config: %s", fileMeta.Data))
+ if format != "" && b64str != "" {
+ // treat as non-image file; fall back to default cost
+ return 3 * baseTokens, nil
+ }
+ return 0, fmt.Errorf("fail to decode image config")📝 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 config.Width == 0 || config.Height == 0 { | |
| // not an image | |
| if format != "" && b64str != "" { | |
| // file type | |
| return 3 * baseTokens, nil | |
| } | |
| return 0, errors.New(fmt.Sprintf("fail to decode base64 config: %s", imageUrl.Url)) | |
| } | |
| shortSide := config.Width | |
| otherSide := config.Height | |
| log.Printf("format: %s, width: %d, height: %d", format, config.Width, config.Height) | |
| // 缩放倍数 | |
| scale := 1.0 | |
| if config.Height < shortSide { | |
| shortSide = config.Height | |
| otherSide = config.Width | |
| } | |
| // 将最小变的尺寸缩小到768以下，如果大于768，则缩放到768 | |
| if shortSide > 768 { | |
| scale = float64(shortSide) / 768 | |
| shortSide = 768 | |
| } | |
| // 将另一边按照相同的比例缩小，向上取整 | |
| otherSide = int(math.Ceil(float64(otherSide) / scale)) | |
| log.Printf("shortSide: %d, otherSide: %d, scale: %f", shortSide, otherSide, scale) | |
| // 计算图片的token数量(边的长度除以512，向上取整) | |
| tiles := (shortSide + 511) / 512 * ((otherSide + 511) / 512) | |
| log.Printf("tiles: %d", tiles) | |
| return 0, errors.New(fmt.Sprintf("fail to decode base64 config: %s", fileMeta.Data)) | |
| } | |
| if config.Width == 0 || config.Height == 0 { | |
| // not an image | |
| if format != "" && b64str != "" { | |
| // treat as non-image file; fall back to default cost | |
| return 3 * baseTokens, nil | |
| } | |
| return 0, fmt.Errorf("fail to decode image config") | |
| } |
🤖 Prompt for AI Agents
In service/token_counter.go around lines 168 to 175 the error currently includes
fileMeta.Data (raw base64/URLs); remove that sensitive/large payload from the
error and return a generic message instead (e.g. "failed to decode base64
config") or, if you need contextual info, include only safe metadata (file name,
type, or a small truncated length) but do not include the raw data itself.
| for _, file := range meta.Files { | ||
| switch file.FileType { | ||
| case types.FileTypeImage: | ||
| if info.RelayFormat == types.RelayFormatGemini { | ||
| tkm += 240 | ||
| } else { | ||
| token, err := getImageToken(file, model, info.IsStream) | ||
| if err != nil { | ||
| return 0, fmt.Errorf("error counting image token: %v", err) | ||
| } | ||
| tkm += token | ||
| } | ||
| case types.FileTypeAudio: |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Verify constants for Gemini image tokens and model coverage
A flat 240 token charge per image for Gemini may not reflect current pricing or counting rules. Please verify against latest docs. If model-dependent, branch by model instead of a fixed value.
🌐 Web query:
For Gemini models (e.g., 1.0/1.5 Pro/Flash), what are the prompt image token counting rules? Is there a fixed per-image token cost or does it depend on resolution/tiles?
💡 Result:
For Gemini models (including 1.0/1.5 Pro and Flash), the image token counting rules depend primarily on the image’s dimensions and tiling, not on a fixed token cost per image—except for very small images. Larger images are split into tiles, and each tile incurs a token cost.
Detailed rules:
- Images ≤384×384 px: Each image counts as 258 tokens regardless of exact resolution[1][9][3].
- Images >384 px in either dimension: The image is internally split into 768x768 px tiles. Each tile is counted as 258 tokens. The image is cropped and/or scaled as needed to fit into an integer number of tiles. For example, a large image could use multiple tiles, each increasing the total token count[1][3].
| Case | Token cost per image | Dependence on size? |
|---|---|---|
| ≤384×384 px | 258 tokens | No (fixed) |
| >384 px edge | 258 × (#tiles) | Yes (increases per tile)[1][3] |
Example cases (from user reports and documentation)
- 320×180 px image: 258 tokens (matches fixed small-image rule)[1].
- 1920×1080 px image: Calculated to require multiple tiles, resulting in a token count of ~1800+ (consistent with splitting into 7 tiles × 258 = 1806 tokens, not just one tile)[1].
- 4 large (1920×1080) images: Results in ~1000 input tokens total, suggesting downscaling or tiling at import[1].
Additional notes
- Gemini 1.0 Pro Vision: May use a fixed 258 tokens per image regardless of size[3], but 1.5 and later always use the tile-based method for large images.
- Token counting precedes billable costs: Cost to users may be flat per image for API billing, but this does not affect the model’s internal token accounting, which affects the maximum prompt size and may influence performance[2][3][8][9].
In summary:
- For Gemini 1.0/1.5 Pro/Flash/2.x, images ≤384 px per edge: fixed 258 tokens.
- For images >384 px on any edge: token count = 258 × (number of 768×768 tiles needed to represent image).
- There is no single fixed token cost per image for all sizes; it scales with image area via the number of tiles for larger images.
Citations:
[1] Google community, user-tested counts
[3] Gemini Pro API forum, summarizing official docs
[9] Google support, image token rule summary
[2] Community pricing discussion—pricing and token count are separate
[8] API comparison—mentions flat image input costs, but token calculation still follows tiling rules
Citations:
- 1: https://discuss.ai.google.dev/t/strange-image-token-counting-with-gemini-2-0-flash/82438
- 2: https://www.googlecloudcommunity.com/gc/AI-ML/Outdated-Gemini-Pro-image-pricing-By-tile-or-by-image/td-p/812892
- 3: https://discuss.ai.google.dev/t/gemini-pro-image-pricing-by-tile-or-fixed/40839
- 4: https://firebase.google.com/docs/ai-logic/count-tokens
- 5: https://ai.google.dev/gemini-api/docs/pricing
- 6: https://packages.debian.org/sid/python/
- 7: https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/list-token
- 8: https://www.leanware.co/insights/gemini-2-5-cost-quality-comparison
- 9: https://support.google.com/gemini/thread/359918381/reach-model-limit?hl=en
- 10: https://techcrunch.com/2025/04/04/gemini-2-5-pro-is-googles-most-expensive-ai-model-yet/
Update Gemini image token handling — replace fixed 240 with 258/tile logic
The flat 240-token constant is incorrect for Gemini images. Per current guidance, use 258 tokens for small images (≤384×384 px) and 258 × (#768×768 tiles) for larger images. Also consider model-specific variants (e.g., Gemini Pro) if they require different rules.
- Files to change:
- service/token_counter.go — replace the Gemini branch in the types.FileTypeImage case (lines ~271–283).
Suggested replacement (concise example — adapt to your actual file metadata and types):
case types.FileTypeImage:
if info.RelayFormat == types.RelayFormatGemini {
// Prefer explicit image dimensions if available; otherwise fall back to existing helper
if file.Width > 0 && file.Height > 0 {
if file.Width <= 384 && file.Height <= 384 {
tkm += 258
} else {
tilesW := (file.Width + 767) / 768
tilesH := (file.Height + 767) / 768
tkm += 258 * tilesW * tilesH
}
} else {
token, err := getImageToken(file, model, info.IsStream)
if err != nil {
return 0, fmt.Errorf("error counting image token: %v", err)
}
tkm += token
}
} else {
token, err := getImageToken(file, model, info.IsStream)
...
}
- Action items:
- Replace magic constant 240 → 258 and implement tile-based calculation.
- Prefer reusing/Extending getImageToken to handle Gemini to avoid duplicated logic.
- Add unit tests for small, tile-aligned, and multi-tile images and verify behavior for Gemini Pro/Flash variants.
🤖 Prompt for AI Agents
In service/token_counter.go around lines 271–283, the Gemini image branch uses a
hardcoded 240 tokens; replace this with the Gemini-specific rule: use 258 tokens
for images ≤384x384 and 258 * (#tiles) for larger images where tiles =
ceil(width/768) * ceil(height/768); prefer using file.Width and file.Height when
available, fall back to calling getImageToken(file, model, info.IsStream) on
missing dimensions, and ensure model-specific variants (e.g., Gemini Pro/Flash)
can be handled (either via model checks or by extending getImageToken) rather
than duplicating logic; update/remove the 240 constant, implement the tile
calculation, and add unit tests for small, single-tile, and multi-tile Gemini
images and for any Pro/Flash variants.
# Conflicts: # dto/openai_image.go
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
dto/openai_image.go (1)
14-14: Fix zero-N undercount in ImagePriceRatio and set TokenType for downstream pricing/metaIf N is omitted it defaults to 0 (uint zero value), making ImagePriceRatio 0 and undercounting pre-consume quota. Also set TokenType for consistency with other DTOs and downstream logic.
Apply this diff:
func (i *ImageRequest) GetTokenCountMeta() *types.TokenCountMeta { var sizeRatio = 1.0 var qualityRatio = 1.0 + + // Default N to 1 if omitted (uint zero value) + n := i.N + if n == 0 { + n = 1 + } @@ // not support token count for dalle return &types.TokenCountMeta{ + TokenType: types.TokenTypeImage, // Prefer Image token type if available CombineText: i.Prompt, MaxTokens: 1584, - ImagePriceRatio: sizeRatio * qualityRatio * float64(i.N), + ImagePriceRatio: sizeRatio * qualityRatio * float64(n), } }If your codebase does not define TokenTypeImage, use types.TokenTypeTokenizer instead. You can verify available token types with:
#!/bin/bash # Inspect available token type constants and their usage rg -n -C2 'TokenTypeImage|TokenTypeTokenizer' rg -n -C3 -P 'type\s+TokenType\b'Also applies to: 30-33, 54-59
🧹 Nitpick comments (8)
dto/openai_image.go (3)
36-44: Use switch for clearer size-to-ratio mappingImproves readability and maintainability without changing behavior.
- // Size - if i.Size == "256x256" { - sizeRatio = 0.4 - } else if i.Size == "512x512" { - sizeRatio = 0.45 - } else if i.Size == "1024x1024" { - sizeRatio = 1 - } else if i.Size == "1024x1792" || i.Size == "1792x1024" { - sizeRatio = 2 - } + // Size + switch i.Size { + case "256x256": + sizeRatio = 0.4 + case "512x512": + sizeRatio = 0.45 + case "1024x1024": + sizeRatio = 1 + case "1024x1792", "1792x1024": + sizeRatio = 2 + }
62-64: Avoid unused parameter warning in IsStreamRename the parameter to underscore to satisfy linters.
-func (i *ImageRequest) IsStream(c *gin.Context) bool { +func (i *ImageRequest) IsStream(_ *gin.Context) bool { return false }
71-73: Make response fields omitempty and consider Go initialism naming
- Adding omitempty prevents emitting empty fields when only one of url or b64_json is present.
- Optional: Use URL/JSON initialisms to follow Go naming conventions (may be a breaking change if referenced elsewhere).
-type ImageData struct { - Url string `json:"url"` - B64Json string `json:"b64_json"` - RevisedPrompt string `json:"revised_prompt"` -} +type ImageData struct { + Url string `json:"url,omitempty"` + B64Json string `json:"b64_json,omitempty"` + RevisedPrompt string `json:"revised_prompt,omitempty"` +}If you want to adopt initialism naming:
type ImageData struct { URL string `json:"url,omitempty"` B64JSON string `json:"b64_json,omitempty"` RevisedPrompt string `json:"revised_prompt,omitempty"` }common/str.go (5)
102-118: Harden MaskEmail: trim, use last '@', guard empty domain, and reuse domain-masking to avoid subdomain leakageCurrent implementation returns "***@" when the domain part is empty (e.g., "user@"), doesn’t trim whitespace, and always reveals the full domain (including subdomains). For consistency with your domain-masking rules and to reduce PII exposure, consider:
- Trim spaces.
- Use LastIndex to be more robust.
- Return "masked" if domain is empty.
- Mask the domain tail via maskHostForPlainDomain.
Apply this diff:
func MaskEmail(email string) string { - if email == "" { - return "***masked***" - } - - // Find the @ symbol - atIndex := strings.Index(email, "@") - if atIndex == -1 { - // No @ symbol found, return masked - return "***masked***" - } - - // Return only the domain part with @ symbol - return "***@" + email[atIndex+1:] + email = strings.TrimSpace(email) + if email == "" { + return "***masked***" + } + // Use last '@' and guard empty local/domain + atIndex := strings.LastIndex(email, "@") + if atIndex <= 0 || atIndex == len(email)-1 { + return "***masked***" + } + domain := email[atIndex+1:] + // Reuse domain masking to avoid leaking subdomains + maskedDomain := maskHostForPlainDomain(domain) + return "***@" + maskedDomain }
120-134: TLD-tail heuristic is fragile; consider using publicsuffix for correctness across registriesThe “ccTLD heuristic” (len(lastPart)==2 && len(secondLastPart)<=3) misclassifies several multi-label public suffixes beyond co.uk/com.cn (e.g., "github.io", "k12.ca.us", many *.uk). For robust masking, use golang.org/x/net/publicsuffix.
Example approach (outside this hunk):
import "golang.org/x/net/publicsuffix" func maskHostTail(parts []string) []string { host := strings.Join(parts, ".") suffix, _ := publicsuffix.PublicSuffix(host) // e.g., "co.uk" or "com" return strings.Split(suffix, ".") }If you’d rather not add a dependency, consider a curated allowlist of two-label suffixes used in your domain set.
226-231: Mask emails before bare domains to avoid leaking local partsWith the new plain-domain masking, a string like "alice@example.com" becomes "alice@***.com", leaving the local part ("alice") unmasked. Since you already introduced MaskEmail, call it before domain masking.
Apply this diff:
- // Mask domain names without protocol (like openai.com, www.openai.com) - domainPattern := regexp.MustCompile(`\b(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}\b`) - str = domainPattern.ReplaceAllStringFunc(str, func(domain string) string { - return maskHostForPlainDomain(domain) - }) + // Mask emails before bare domains to avoid leaking local parts + emailPattern := regexp.MustCompile(`[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}`) + str = emailPattern.ReplaceAllStringFunc(str, MaskEmail) + + // Mask domain names without protocol (like openai.com, www.openai.com) + domainPattern := regexp.MustCompile(`\b(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}\b`) + str = domainPattern.ReplaceAllStringFunc(str, func(domain string) string { + return maskHostForPlainDomain(domain) + })Optional: precompile these regexes at package scope for performance.
146-161: Single-label domains (e.g., "localhost") behaviorFor plain domains, you currently return the original string when there’s no dot. That’s reasonable, but if you want consistency with maskHostForURL (which returns "***" for <2 parts), consider masking single-label domains too.
Possible tweak:
- if len(parts) < 2 { - return domain - } + if len(parts) < 2 { + return "***" + }
173-224: Optional: precompile regexes once at package scopeurlPattern, domainPattern, (and the new emailPattern) are compiled per call. Moving them to package-level variables reduces allocations and CPU in hot paths (e.g., logging). Leave ipPattern as well if you move forward.
Outside this hunk, add:
var ( urlPattern = regexp.MustCompile(`(http|https)://[^\s/$.?#].[^\s]*`) domainPattern = regexp.MustCompile(`\b(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}\b`) emailPattern = regexp.MustCompile(`[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}`) ipv4Pattern = regexp.MustCompile(`\b(?:\d{1,3}\.){3}\d{1,3}\b`) )Then inside MaskSensitiveInfo, reference these vars instead of redeclaring.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
common/str.go(3 hunks)dto/openai_image.go(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
dto/openai_image.go (3)
model/model_meta.go (1)
Model(22-42)dto/openai_request.go (1)
ResponseFormat(13-16)types/request_meta.go (1)
TokenCountMeta(20-31)
🔇 Additional comments (2)
common/str.go (2)
226-231: Nice: Added plain-domain masking and removed the URL-skip pitfallThe new domain-only masking is a good addition, and the previous pitfall that skipped bare-domain matches when a URL with the same domain existed is gone. This aligns with the earlier feedback.
185-187: Unify host masking call is clean and readableThe move to a single maskHostForURL path improves consistency versus scattered per-host logic.
| // maskHostForURL collapses subdomains and keeps only masked prefix + preserved tail. | ||
| // Example: api.openai.com -> ***.com, sub.domain.co.uk -> ***.co.uk | ||
| func maskHostForURL(host string) string { | ||
| parts := strings.Split(host, ".") | ||
| if len(parts) < 2 { | ||
| return "***" | ||
| } | ||
| tail := maskHostTail(parts) | ||
| return "***." + strings.Join(tail, ".") | ||
| } | ||
|
|
There was a problem hiding this comment.
Critical: URLs with IP hosts leak the last octet and ports aren’t handled; use Hostname/Port and IP-aware masking
- Example: "http://192.168.1.1:8080/path" is masked to "http://.1/" (last octet leaked), because maskHostForURL splits by '.' and treats IPs as hostnames.
- Also, u.Host may include a port; the current approach merges the port into the last label, breaking the ccTLD heuristic and producing odd outputs.
Fixes:
- Use URL.Hostname() and URL.Port() to separate host and port.
- Teach maskHostForURL to detect IPs (via net.ParseIP) and mask accordingly:
- IPv4: "..."
- IPv6: "***" (consistent with your fallback for non-dot hosts)
Apply these diffs:
func maskHostForURL(host string) string {
- parts := strings.Split(host, ".")
+ // If host is an IP, mask accordingly
+ if ip := net.ParseIP(host); ip != nil {
+ if ip.To4() != nil {
+ return "***.***.***.***"
+ }
+ // IPv6
+ return "***"
+ }
+ parts := strings.Split(host, ".")
if len(parts) < 2 {
return "***"
}
tail := maskHostTail(parts)
return "***." + strings.Join(tail, ".")
}- host := u.Host
- if host == "" {
+ hostOnly := u.Hostname()
+ port := u.Port()
+ if hostOnly == "" {
return urlStr
}
- // Mask host with unified logic
- maskedHost := maskHostForURL(host)
+ // Mask host with unified logic (preserve port if present)
+ maskedHost := maskHostForURL(hostOnly)
+ if port != "" {
+ maskedHost += ":" + port
+ }
result := u.Scheme + "://" + maskedHostAdditionally, add this import at the top of the file:
import "net"Also applies to: 180-189
🤖 Prompt for AI Agents
In common/str.go around lines 135-145 (and similarly apply the fix at 180-189),
maskHostForURL currently splits u.Host by '.' which leaks IP octets and
mis-handles ports; change to use URL.Hostname() and URL.Port() to separate host
and port, add import "net", and update masking logic to detect IPs via
net.ParseIP: if IPv4 return "***.***.***.***" (and preserve the original port if
present), if IPv6 return "***" (plus port if present); for non-IP hosts keep the
existing subdomain-collapse behavior but operate on the hostname only and
re-append the port (":port") when returning the masked host string.
refactor: Introduce pre-consume quota and unify relay handlers
1 participant
Summary by CodeRabbit
New Features
Improvements
Stability