Skip to content

Commit 2c667a1

Browse files
fix(provider): retain upstream model for gemini compat and ws
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
1 parent bac4080 commit 2c667a1

3 files changed

Lines changed: 102 additions & 34 deletions

File tree

backend/internal/service/gemini_messages_compat_service.go

Lines changed: 39 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1028,14 +1028,15 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
10281028
}
10291029

10301030
return &ForwardResult{
1031-
RequestID: requestID,
1032-
Usage: *usage,
1033-
Model: originalModel,
1034-
Stream: req.Stream,
1035-
Duration: time.Since(startTime),
1036-
FirstTokenMs: firstTokenMs,
1037-
ImageCount: imageCount,
1038-
ImageSize: imageSize,
1031+
RequestID: requestID,
1032+
Usage: *usage,
1033+
Model: originalModel,
1034+
UpstreamModel: mappedModel,
1035+
Stream: req.Stream,
1036+
Duration: time.Since(startTime),
1037+
FirstTokenMs: firstTokenMs,
1038+
ImageCount: imageCount,
1039+
ImageSize: imageSize,
10391040
}, nil
10401041
}
10411042

@@ -1241,12 +1242,13 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
12411242
estimated := estimateGeminiCountTokens(body)
12421243
c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
12431244
return &ForwardResult{
1244-
RequestID: "",
1245-
Usage: ClaudeUsage{},
1246-
Model: originalModel,
1247-
Stream: false,
1248-
Duration: time.Since(startTime),
1249-
FirstTokenMs: nil,
1245+
RequestID: "",
1246+
Usage: ClaudeUsage{},
1247+
Model: originalModel,
1248+
UpstreamModel: mappedModel,
1249+
Stream: false,
1250+
Duration: time.Since(startTime),
1251+
FirstTokenMs: nil,
12501252
}, nil
12511253
}
12521254
setOpsUpstreamError(c, 0, safeErr, "")
@@ -1310,12 +1312,13 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
13101312
estimated := estimateGeminiCountTokens(body)
13111313
c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
13121314
return &ForwardResult{
1313-
RequestID: "",
1314-
Usage: ClaudeUsage{},
1315-
Model: originalModel,
1316-
Stream: false,
1317-
Duration: time.Since(startTime),
1318-
FirstTokenMs: nil,
1315+
RequestID: "",
1316+
Usage: ClaudeUsage{},
1317+
Model: originalModel,
1318+
UpstreamModel: mappedModel,
1319+
Stream: false,
1320+
Duration: time.Since(startTime),
1321+
FirstTokenMs: nil,
13191322
}, nil
13201323
}
13211324
// Final attempt: surface the upstream error body (passed through below) instead of a generic retry error.
@@ -1350,12 +1353,13 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
13501353
estimated := estimateGeminiCountTokens(body)
13511354
c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
13521355
return &ForwardResult{
1353-
RequestID: requestID,
1354-
Usage: ClaudeUsage{},
1355-
Model: originalModel,
1356-
Stream: false,
1357-
Duration: time.Since(startTime),
1358-
FirstTokenMs: nil,
1356+
RequestID: requestID,
1357+
Usage: ClaudeUsage{},
1358+
Model: originalModel,
1359+
UpstreamModel: mappedModel,
1360+
Stream: false,
1361+
Duration: time.Since(startTime),
1362+
FirstTokenMs: nil,
13591363
}, nil
13601364
}
13611365

@@ -1527,14 +1531,15 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
15271531
}
15281532

15291533
return &ForwardResult{
1530-
RequestID: requestID,
1531-
Usage: *usage,
1532-
Model: originalModel,
1533-
Stream: stream,
1534-
Duration: time.Since(startTime),
1535-
FirstTokenMs: firstTokenMs,
1536-
ImageCount: imageCount,
1537-
ImageSize: imageSize,
1534+
RequestID: requestID,
1535+
Usage: *usage,
1536+
Model: originalModel,
1537+
UpstreamModel: mappedModel,
1538+
Stream: stream,
1539+
Duration: time.Since(startTime),
1540+
FirstTokenMs: firstTokenMs,
1541+
ImageCount: imageCount,
1542+
ImageSize: imageSize,
15381543
}, nil
15391544
}
15401545

backend/internal/service/gemini_messages_compat_service_test.go

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package service
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67
"io"
@@ -15,6 +16,30 @@ import (
1516
"github.com/stretchr/testify/require"
1617
)
1718

19+
type geminiCompatHTTPUpstreamStub struct {
20+
response *http.Response
21+
err error
22+
calls int
23+
lastReq *http.Request
24+
}
25+
26+
func (s *geminiCompatHTTPUpstreamStub) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
27+
s.calls++
28+
s.lastReq = req
29+
if s.err != nil {
30+
return nil, s.err
31+
}
32+
if s.response == nil {
33+
return nil, fmt.Errorf("missing stub response")
34+
}
35+
resp := *s.response
36+
return &resp, nil
37+
}
38+
39+
func (s *geminiCompatHTTPUpstreamStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, enableTLSFingerprint bool) (*http.Response, error) {
40+
return s.Do(req, proxyURL, accountID, accountConcurrency)
41+
}
42+
1843
// TestConvertClaudeToolsToGeminiTools_CustomType 测试custom类型工具转换
1944
func TestConvertClaudeToolsToGeminiTools_CustomType(t *testing.T) {
2045
tests := []struct {
@@ -170,6 +195,42 @@ func TestGeminiHandleNativeNonStreamingResponse_DebugDisabledDoesNotEmitHeaderLo
170195
require.False(t, logSink.ContainsMessage("[GeminiAPI]"), "debug 关闭时不应输出 Gemini 响应头日志")
171196
}
172197

198+
func TestGeminiMessagesCompatServiceForward_PreservesRequestedModelAndMappedUpstreamModel(t *testing.T) {
199+
gin.SetMode(gin.TestMode)
200+
w := httptest.NewRecorder()
201+
c, _ := gin.CreateTestContext(w)
202+
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
203+
204+
httpStub := &geminiCompatHTTPUpstreamStub{
205+
response: &http.Response{
206+
StatusCode: http.StatusOK,
207+
Header: http.Header{"x-request-id": []string{"gemini-req-1"}},
208+
Body: io.NopCloser(strings.NewReader(`{"candidates":[{"content":{"parts":[{"text":"hello"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5}}`)),
209+
},
210+
}
211+
svc := &GeminiMessagesCompatService{httpUpstream: httpStub, cfg: &config.Config{}}
212+
account := &Account{
213+
ID: 1,
214+
Type: AccountTypeAPIKey,
215+
Credentials: map[string]any{
216+
"api_key": "test-key",
217+
"model_mapping": map[string]any{
218+
"claude-sonnet-4": "claude-sonnet-4-20250514",
219+
},
220+
},
221+
}
222+
body := []byte(`{"model":"claude-sonnet-4","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`)
223+
224+
result, err := svc.Forward(context.Background(), c, account, body)
225+
require.NoError(t, err)
226+
require.NotNil(t, result)
227+
require.Equal(t, "claude-sonnet-4", result.Model)
228+
require.Equal(t, "claude-sonnet-4-20250514", result.UpstreamModel)
229+
require.Equal(t, 1, httpStub.calls)
230+
require.NotNil(t, httpStub.lastReq)
231+
require.Contains(t, httpStub.lastReq.URL.String(), "/models/claude-sonnet-4-20250514:")
232+
}
233+
173234
func TestConvertClaudeMessagesToGeminiGenerateContent_AddsThoughtSignatureForToolUse(t *testing.T) {
174235
claudeReq := map[string]any{
175236
"model": "claude-haiku-4-5-20251001",

backend/internal/service/openai_ws_forwarder.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2328,6 +2328,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
23282328
RequestID: responseID,
23292329
Usage: *usage,
23302330
Model: originalModel,
2331+
UpstreamModel: mappedModel,
23312332
ServiceTier: extractOpenAIServiceTier(reqBody),
23322333
ReasoningEffort: extractOpenAIReasoningEffort(reqBody, originalModel),
23332334
Stream: reqStream,
@@ -2945,6 +2946,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
29452946
RequestID: responseID,
29462947
Usage: usage,
29472948
Model: originalModel,
2949+
UpstreamModel: mappedModel,
29482950
ServiceTier: extractOpenAIServiceTierFromBody(payload),
29492951
ReasoningEffort: extractOpenAIReasoningEffortFromBody(payload, originalModel),
29502952
Stream: reqStream,

0 commit comments

Comments
 (0)