feat: add MiMo upstream support with Responses/Chat switching#248
Merged
Conversation
…ount-level model mapping
📝 WalkthroughWalkthroughThis PR adds complete MiMo upstream provider support to the proxy. The feature implements bidirectional protocol translation between OpenAI Responses/Chat Completions APIs and MiMo's chat protocol, including request routing, streaming response handling, per-account model mapping, and auto web search integration for Token Plan accounts. ChangesMiMo Upstream Integration
Sequence DiagramsequenceDiagram
participant Client
participant Handler as Responses/Chat Handler
participant Router as MiMo Router
participant Executor as MiMo Executor
participant Upstream as MiMo /v1/chat/completions
participant Translator as Response Translator
Client->>Handler: Responses/Chat Completions request
Handler->>Handler: io.ReadAll body, apply model mapping
Handler->>Router: Check if MiMo account or MiMo model
alt MiMo Route Detected
Router->>Executor: Request body + MimoUpstreamConfig
Executor->>Translator: Convert to MiMo format
Translator->>Upstream: POST /v1/chat/completions
Upstream->>Translator: Streaming or full response
Translator->>Executor: Responses SSE events
Executor->>Client: response.created / deltas / response.completed
else Legacy Route
Router->>Handler: Continue Codex flow
Handler->>Client: Codex response
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~65 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
proxy/mimo_executor.go (1)
353-363: 💤 Low valueConsider passing the actual model name.
The
modelfield is hardcoded as"mimo"but the actual model name is available in the calling context. Clients may benefit from seeing the specific model (e.g.,mimo-v2.5-pro).♻️ Proposed refactor
-func buildResponsesCreatedEvent(responseID string) string { +func buildResponsesCreatedEvent(responseID, model string) string { resp := map[string]any{ "id": responseID, "object": "response", "status": "in_progress", - "model": "mimo", + "model": model, "output": []any{}, }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@proxy/mimo_executor.go` around lines 353 - 363, The buildResponsesCreatedEvent function currently hardcodes "model":"mimo"; change its signature to accept a model string (e.g., func buildResponsesCreatedEvent(responseID, model string) string), set the "model" field to that parameter when building resp, and update every call site to pass the actual model name from the caller (so callers supply values like "mimo-v2.5-pro"); no other behavior changes required.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@proxy/handler.go`:
- Around line 1505-1533: At both locations (proxy/handler.go:1505-1533 and
proxy/handler.go:2852-2878) the MiMo routing paths return early after acquiring
an account and so leak the account lease; before calling HandleMimoResponsesAPI
(first site) and HandleMimoChatCompletionsAPI (second site) release the acquired
account using the same release call used elsewhere in this file (i.e., invoke
the pool/account release helper just prior to the early return), and rename the
local shadowed apiKey variable to mimoAPIKey to avoid shadowing the outer
account.APIKey; the first site (1505-1533) requires inserting the release call
then calling HandleMimoResponsesAPI with mimoCfg and returning, and the second
site (2852-2878) requires the same change before calling
HandleMimoChatCompletionsAPI.
In `@proxy/mimo_executor.go`:
- Around line 265-280: translateChatStreamChunk already emits a
response.completed event for finish_reason "stop" or "tool_calls", so the
unconditional send of buildResponsesCompletedEvent(responseID) in the
currentData flush path can cause duplicate completion events; add a completion
flag (e.g., responseCompleted bool) in the mimo_executor.go handler scope or on
the connection context, set it to true whenever translateChatStreamChunk or
ChatChunkToResponsesEvent produces a response.completed event, and before
calling fmt.Fprintf(c.Writer, "%s\n\n",
buildResponsesCompletedEvent(responseID)) check that flag and only emit the
completed event if responseCompleted is false (no other change to event
formatting); ensure the flag is updated wherever
translateChatStreamChunk/ChatChunkToResponsesEvent emits completion so
duplicates are suppressed.
- Around line 30-35: The GenerateMimoResponseID function currently ignores
errors from rand.Read; update GenerateMimoResponseID to check the returned n and
err from rand.Read(b) and handle failures — if err == nil ensure n == len(b) and
then return "resp_mimo_"+hex.EncodeToString(b); if rand.Read returns an error,
fall back to a safe alternative (for example, generate a unique ID using a
timestamp + crypto/sha256 or seed math/rand with time.Now().UnixNano() and
produce hex bytes) or return an error upstream (change signature to (string,
error)) — choose one consistent approach and implement the error-path code,
keeping the function name GenerateMimoResponseID and ensuring the fallback
produces reasonably unique IDs.
In `@proxy/mimo_translator.go`:
- Around line 869-905: injectSearchResultsToChatBody currently only detects
existence of a system message but always appends to messages.0, corrupting the
wrong item; update the function to record the actual index of the system message
during the messages.ForEach loop (e.g., store an int sysIdx initialized to -1
and set it when msg.Get("role").String() == "system"), then when sysIdx >= 0
append searchResults to "messages.<sysIdx>.content" using sjson.SetBytes; keep
the existing insertion branch unchanged for the case where no system message
exists.
---
Nitpick comments:
In `@proxy/mimo_executor.go`:
- Around line 353-363: The buildResponsesCreatedEvent function currently
hardcodes "model":"mimo"; change its signature to accept a model string (e.g.,
func buildResponsesCreatedEvent(responseID, model string) string), set the
"model" field to that parameter when building resp, and update every call site
to pass the actual model name from the caller (so callers supply values like
"mimo-v2.5-pro"); no other behavior changes required.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 99795126-43dc-4be8-bbc0-389091157db2
📒 Files selected for processing (3)
proxy/handler.goproxy/mimo_executor.goproxy/mimo_translator.go
Comment on lines
+1505
to
+1533
| // ============================================================ | ||
| // MiMo 上游路由 | ||
| // ============================================================ | ||
| // 方式1: 账号级别配置（推荐）- 支持 base_url 自动检测 | ||
| if IsMimoAccount(account.UpstreamType, account.BaseURL, account.APIKey) { | ||
| mimoCfg := &MimoUpstreamConfig{ | ||
| BaseURL: account.BaseURL, | ||
| APIKey: account.APIKey, | ||
| IsTokenPlan: IsMimoTokenPlanKey(account.APIKey), | ||
| ModelMapping: marshalModelMapping(account.ModelMapping), | ||
| } | ||
| HandleMimoResponsesAPI(c, rawBody, mimoCfg) | ||
| return | ||
| } | ||
|
|
||
| // 方式2: 根据模型名称自动检测 + 全局环境变量 | ||
| if IsMimoModel(model) { | ||
| apiKey := os.Getenv("MIMO_API_KEY") | ||
| if apiKey != "" { | ||
| mimoCfg := &MimoUpstreamConfig{ | ||
| BaseURL: os.Getenv("MIMO_BASE_URL"), | ||
| APIKey: apiKey, | ||
| IsTokenPlan: IsMimoTokenPlanKey(apiKey), | ||
| } | ||
| HandleMimoResponsesAPI(c, rawBody, mimoCfg) | ||
| return | ||
| } | ||
| } | ||
| // ============================================================ |
There was a problem hiding this comment.
Account resource leak in MiMo routing paths across both handlers. Both Responses and ChatCompletions handlers have MiMo routing blocks placed inside their retry loops after acquiring an account from the pool. When routing to MiMo upstream, the handlers return early without releasing the acquired account, causing pool lease leaks.
proxy/handler.go#L1505-L1533: Release account before callingHandleMimoResponsesAPIin both the account-level and environment variable routing paths; rename shadowedapiKeytomimoAPIKey.proxy/handler.go#L2852-L2878: Apply the same fix before callingHandleMimoChatCompletionsAPI; rename shadowedapiKeytomimoAPIKey.
📍 Affects 1 file
proxy/handler.go#L1505-L1533(this comment)proxy/handler.go#L2852-L2878
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@proxy/handler.go` around lines 1505 - 1533, At both locations
(proxy/handler.go:1505-1533 and proxy/handler.go:2852-2878) the MiMo routing
paths return early after acquiring an account and so leak the account lease;
before calling HandleMimoResponsesAPI (first site) and
HandleMimoChatCompletionsAPI (second site) release the acquired account using
the same release call used elsewhere in this file (i.e., invoke the pool/account
release helper just prior to the early return), and rename the local shadowed
apiKey variable to mimoAPIKey to avoid shadowing the outer account.APIKey; the
first site (1505-1533) requires inserting the release call then calling
HandleMimoResponsesAPI with mimoCfg and returning, and the second site
(2852-2878) requires the same change before calling
HandleMimoChatCompletionsAPI.
Owner
|
感谢 PR！在最新 完整的未定义符号清单（已去重，rebase 到最新 main 后）： A. MiMo 配置/判别函数（疑似缺少一个
|
Based on v2.3.2, fix 3 issues found in CodeRabbit review: 1. injectSearchResultsToChatBody index bug (mimo_translator.go) - Issue: Hardcoded messages.0 assumes system message is always first - Fix: Dynamically find system message actual index 2. GenerateMimoResponseID error handling (mimo_executor.go) - Issue: Ignored rand.Read error return value - Fix: Add error check, use timestamp as fallback 3. buildResponsesCreatedEvent model hardcoded (mimo_executor.go) - Issue: Model name fixed as 'mimo' - Fix: Pass model name as parameter Backward compatible, no API changes.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
proxy/mimo_translator.go (2)
1400-1404:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMissing
isTokenPlanparameter inTranslateChatCompletionsToMimocauses inconsistent auto-search behavior.
TranslateChatCompletionsToMimodoesn't receive anisTokenPlanparameter, unlikeResponsesToMimoChatwhich uses it to gate auto-search (line 181). The comment at line 1452 states "Token Plan 账号 + 搜索已启用 = 自动搜索注入" but the condition at line 1453 only checkshasWebSearch && IsSearchEnabled(). This causes Chat Completions requests to trigger auto-search for all accounts, not just Token Plan accounts.
proxy/mimo_translator.go#L1400-L1404: AddisTokenPlan boolparameter to the function signature and include it in the condition at line 1453:if hasWebSearch && isTokenPlan && IsSearchEnabled().proxy/mimo_executor.go#L147-L147: Update call site to passcfg.IsTokenPlanas the second argument.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@proxy/mimo_translator.go` around lines 1400 - 1404, The TranslateChatCompletionsToMimo function in proxy/mimo_translator.go at lines 1400-1404 lacks an isTokenPlan boolean parameter, causing inconsistent auto-search behavior across different account types. Add the isTokenPlan bool parameter to the TranslateChatCompletionsToMimo function signature, then locate the auto-search condition at line 1453 and update it to check hasWebSearch AND isTokenPlan AND IsSearchEnabled() instead of only checking hasWebSearch and IsSearchEnabled(). Additionally, update the function call site in proxy/mimo_executor.go at line 147 to pass cfg.IsTokenPlan as the second argument when calling TranslateChatCompletionsToMimo.
1452-1453:⚠️ Potential issue | 🟠 Major | ⚡ Quick winInconsistent Token Plan check for auto search.
ResponsesToMimoChat(line 181) checkshasWebSearch && isTokenPlan && IsSearchEnabled()before triggering auto search, butTranslateChatCompletionsToMimoonly checkshasWebSearch && IsSearchEnabled()despite the comment on line 1452 stating "Token Plan 账号 + 搜索已启用 = 自动搜索注入".This inconsistency means Chat Completions requests will trigger auto search for all accounts with web_search tools, not just Token Plan accounts.
🐛 Proposed fix to add isTokenPlan parameter
-func TranslateChatCompletionsToMimo(chatBody []byte) ([]byte, string, error) { +func TranslateChatCompletionsToMimo(chatBody []byte, isTokenPlan bool) ([]byte, string, error) { // ... existing code ... // Token Plan 账号 + 搜索已启用 = 自动搜索注入 - if hasWebSearch && IsSearchEnabled() { + if hasWebSearch && isTokenPlan && IsSearchEnabled() {Also update the call site in
HandleMimoChatCompletionsAPIto passcfg.IsTokenPlan.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@proxy/mimo_translator.go` around lines 1452 - 1453, The `TranslateChatCompletionsToMimo` function at line 1452-1453 is missing the `isTokenPlan` check in its auto search condition, causing inconsistency with the `ResponsesToMimoChat` function which correctly checks `hasWebSearch && isTokenPlan && IsSearchEnabled()`. Update the condition in `TranslateChatCompletionsToMimo` to include the `isTokenPlan` parameter check alongside `hasWebSearch && IsSearchEnabled()` to match the logic in `ResponsesToMimoChat`. Additionally, modify the call site in `HandleMimoChatCompletionsAPI` to pass the `isTokenPlan` parameter (likely from `cfg.IsTokenPlan`) when invoking `TranslateChatCompletionsToMimo` so that the function has access to the required value.
🧹 Nitpick comments (2)
proxy/mimo_executor.go (2)
586-589: 💤 Low valueInconsistent logging: uses
log.Printfwhile rest of file usesslog.Lines 586 and 588 use
log.Printfwhile the rest of the file consistently uses structured logging viaslog. Consider usingslog.Infoorslog.Debugfor consistency.♻️ Suggested change
- log.Printf("[MIMO-USAGE] input=%d, output=%d, total=%d, cached=%d, reasoning=%d", usageResult.InputTokens, usageResult.OutputTokens, usageResult.TotalTokens, usageResult.CachedTokens, usageResult.ReasoningTokens) + slog.Debug("mimo usage", + "input", usageResult.InputTokens, + "output", usageResult.OutputTokens, + "total", usageResult.TotalTokens, + "cached", usageResult.CachedTokens, + "reasoning", usageResult.ReasoningTokens) } else { - log.Printf("[MIMO-USAGE] lastUsageData is EMPTY") + slog.Debug("mimo usage: lastUsageData is empty") }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@proxy/mimo_executor.go` around lines 586 - 589, Replace the two log.Printf calls in the MIMO-USAGE logging block (the one logging usageResult.InputTokens, usageResult.OutputTokens, etc., and the one logging "lastUsageData is EMPTY") with slog.Info or slog.Debug calls to maintain consistency with the structured logging approach used throughout the rest of the file. Convert the format string parameters into structured key-value pairs that slog expects, preserving all the logged information.
689-757: 💤 Low valueCode duplication with
buildResponsesCompletedEventWithDatain mimo_translator.go.
buildResponsesCompletedEventWithDataWithSeqduplicates most of the logic frombuildResponsesCompletedEventWithData(lines 1137-1205 in mimo_translator.go), only adding asequence_numberfield. Consider refactoring to share the common logic.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@proxy/mimo_executor.go` around lines 689 - 757, The function buildResponsesCompletedEventWithDataWithSeq duplicates the logic of buildResponsesCompletedEventWithData (which exists in mimo_translator.go) with the only difference being the addition of a sequence_number field. Refactor by modifying buildResponsesCompletedEventWithData to accept an optional sequence number parameter, then have buildResponsesCompletedEventWithDataWithSeq call this refactored function with the sequence number provided. This way you avoid duplicating the entire response and event construction logic, and the sequence_number field is only added to the event map when a non-zero or non-nil sequence value is passed.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@proxy/mimo_executor.go`:
- Around line 565-568: The json.Unmarshal function cannot populate a
gjson.Result type as gjson.Result is not a standard JSON unmarshaler type.
Replace the json.Unmarshal(usageBytes, &usageJSON) call with
gjson.ParseBytes(usageBytes) which is the correct gjson library function to
parse bytes into a gjson.Result, assigning the return value directly to the
usageJSON variable. This will ensure buildResponsesUsageEvent receives a
properly populated gjson.Result instead of an empty/zero state.
---
Outside diff comments:
In `@proxy/mimo_translator.go`:
- Around line 1400-1404: The TranslateChatCompletionsToMimo function in
proxy/mimo_translator.go at lines 1400-1404 lacks an isTokenPlan boolean
parameter, causing inconsistent auto-search behavior across different account
types. Add the isTokenPlan bool parameter to the TranslateChatCompletionsToMimo
function signature, then locate the auto-search condition at line 1453 and
update it to check hasWebSearch AND isTokenPlan AND IsSearchEnabled() instead of
only checking hasWebSearch and IsSearchEnabled(). Additionally, update the
function call site in proxy/mimo_executor.go at line 147 to pass cfg.IsTokenPlan
as the second argument when calling TranslateChatCompletionsToMimo.
- Around line 1452-1453: The `TranslateChatCompletionsToMimo` function at line
1452-1453 is missing the `isTokenPlan` check in its auto search condition,
causing inconsistency with the `ResponsesToMimoChat` function which correctly
checks `hasWebSearch && isTokenPlan && IsSearchEnabled()`. Update the condition
in `TranslateChatCompletionsToMimo` to include the `isTokenPlan` parameter check
alongside `hasWebSearch && IsSearchEnabled()` to match the logic in
`ResponsesToMimoChat`. Additionally, modify the call site in
`HandleMimoChatCompletionsAPI` to pass the `isTokenPlan` parameter (likely from
`cfg.IsTokenPlan`) when invoking `TranslateChatCompletionsToMimo` so that the
function has access to the required value.
---
Nitpick comments:
In `@proxy/mimo_executor.go`:
- Around line 586-589: Replace the two log.Printf calls in the MIMO-USAGE
logging block (the one logging usageResult.InputTokens,
usageResult.OutputTokens, etc., and the one logging "lastUsageData is EMPTY")
with slog.Info or slog.Debug calls to maintain consistency with the structured
logging approach used throughout the rest of the file. Convert the format string
parameters into structured key-value pairs that slog expects, preserving all the
logged information.
- Around line 689-757: The function buildResponsesCompletedEventWithDataWithSeq
duplicates the logic of buildResponsesCompletedEventWithData (which exists in
mimo_translator.go) with the only difference being the addition of a
sequence_number field. Refactor by modifying
buildResponsesCompletedEventWithData to accept an optional sequence number
parameter, then have buildResponsesCompletedEventWithDataWithSeq call this
refactored function with the sequence number provided. This way you avoid
duplicating the entire response and event construction logic, and the
sequence_number field is only added to the event map when a non-zero or non-nil
sequence value is passed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 1a12a4d3-ba20-4696-a8ed-c2e6cf1f22b9
📒 Files selected for processing (2)
proxy/mimo_executor.goproxy/mimo_translator.go
Comment on lines
+565
to
+568
| var usageJSON gjson.Result | ||
| json.Unmarshal(usageBytes, &usageJSON) | ||
| seq = state.NextSeq() | ||
| emit(buildResponsesUsageEvent(usageJSON, responseID, seq)) |
There was a problem hiding this comment.
Bug: json.Unmarshal into gjson.Result doesn't work.
gjson.Result is a struct used by the gjson library for query results, not a type that json.Unmarshal can populate. The unmarshal will effectively leave usageJSON in an empty/zero state, causing buildResponsesUsageEvent to receive an invalid result.
🐛 Proposed fix
// usage 事件
if lastUsageData != "" {
- usageBytes := []byte(lastUsageData)
- var usageJSON gjson.Result
- json.Unmarshal(usageBytes, &usageJSON)
+ usageJSON := gjson.Parse(lastUsageData)
seq = state.NextSeq()
emit(buildResponsesUsageEvent(usageJSON, responseID, seq))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| var usageJSON gjson.Result | |
| json.Unmarshal(usageBytes, &usageJSON) | |
| seq = state.NextSeq() | |
| emit(buildResponsesUsageEvent(usageJSON, responseID, seq)) | |
| usageJSON := gjson.Parse(lastUsageData) | |
| seq = state.NextSeq() | |
| emit(buildResponsesUsageEvent(usageJSON, responseID, seq)) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@proxy/mimo_executor.go` around lines 565 - 568, The json.Unmarshal function
cannot populate a gjson.Result type as gjson.Result is not a standard JSON
unmarshaler type. Replace the json.Unmarshal(usageBytes, &usageJSON) call with
gjson.ParseBytes(usageBytes) which is the correct gjson library function to
parse bytes into a gjson.Result, assigning the return value directly to the
usageJSON variable. This will ensure buildResponsesUsageEvent receives a
properly populated gjson.Result instead of an empty/zero state.
Owner
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add MiMo upstream integration supporting both Responses API and Chat Completions API, with automatic client type detection and format translation.
Features
Fixes found during integration
Files changed
Summary by CodeRabbit