Merge branch 'Wei-Shaw:main' into main

This commit is contained in:
程序猿MT
2026-02-12 23:12:41 +08:00
committed by GitHub
30 changed files with 583 additions and 76 deletions

View File

@@ -65,3 +65,27 @@ func (h *AntigravityOAuthHandler) ExchangeCode(c *gin.Context) {
response.Success(c, tokenInfo)
}
// AntigravityRefreshTokenRequest represents the request for validating Antigravity refresh token
type AntigravityRefreshTokenRequest struct {
RefreshToken string `json:"refresh_token" binding:"required"`
ProxyID *int64 `json:"proxy_id"`
}
// RefreshToken validates an Antigravity refresh token and returns full token info
// POST /api/v1/admin/antigravity/oauth/refresh-token
func (h *AntigravityOAuthHandler) RefreshToken(c *gin.Context) {
var req AntigravityRefreshTokenRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.BadRequest(c, "请求无效: "+err.Error())
return
}
tokenInfo, err := h.antigravityOAuthService.ValidateRefreshToken(c.Request.Context(), req.RefreshToken, req.ProxyID)
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, tokenInfo)
}

View File

@@ -202,7 +202,7 @@ func (h *RedeemHandler) Export(c *gin.Context) {
writer := csv.NewWriter(&buf)
// Write header
if err := writer.Write([]string{"id", "code", "type", "value", "status", "used_by", "used_at", "created_at"}); err != nil {
if err := writer.Write([]string{"id", "code", "type", "value", "status", "used_by", "used_by_email", "used_at", "created_at"}); err != nil {
response.InternalError(c, "Failed to export redeem codes: "+err.Error())
return
}
@@ -213,6 +213,10 @@ func (h *RedeemHandler) Export(c *gin.Context) {
if code.UsedBy != nil {
usedBy = fmt.Sprintf("%d", *code.UsedBy)
}
usedByEmail := ""
if code.User != nil {
usedByEmail = code.User.Email
}
usedAt := ""
if code.UsedAt != nil {
usedAt = code.UsedAt.Format("2006-01-02 15:04:05")
@@ -224,6 +228,7 @@ func (h *RedeemHandler) Export(c *gin.Context) {
fmt.Sprintf("%.2f", code.Value),
code.Status,
usedBy,
usedByEmail,
usedAt,
code.CreatedAt.Format("2006-01-02 15:04:05"),
}); err != nil {

View File

@@ -27,7 +27,7 @@ type ClaudeMessage struct {
// ThinkingConfig Thinking 配置
type ThinkingConfig struct {
Type string `json:"type"` // "enabled" or "disabled"
Type string `json:"type"` // "enabled" / "adaptive" / "disabled"
BudgetTokens int `json:"budget_tokens,omitempty"` // thinking budget
}

View File

@@ -155,6 +155,7 @@ type GeminiUsageMetadata struct {
CandidatesTokenCount int `json:"candidatesTokenCount,omitempty"`
CachedContentTokenCount int `json:"cachedContentTokenCount,omitempty"`
TotalTokenCount int `json:"totalTokenCount,omitempty"`
ThoughtsTokenCount int `json:"thoughtsTokenCount,omitempty"` // thinking tokens按输出价格计费
}
// GeminiGroundingMetadata Gemini grounding 元数据Web Search

View File

@@ -64,6 +64,10 @@ const MaxTokensBudgetPadding = 1000
// Gemini 2.5 Flash thinking budget 上限
const Gemini25FlashThinkingBudgetLimit = 24576
// 对于 Antigravity 的 Claudebudget-only模型该语义最终等价为 thinkingBudget=24576。
// 这里复用相同数值以保持行为一致。
const ClaudeAdaptiveHighThinkingBudgetTokens = Gemini25FlashThinkingBudgetLimit
// ensureMaxTokensGreaterThanBudget 确保 max_tokens > budget_tokens
// Claude API 要求启用 thinking 时max_tokens 必须大于 thinking.budget_tokens
// 返回调整后的 maxTokens 和是否进行了调整
@@ -96,7 +100,7 @@ func TransformClaudeToGeminiWithOptions(claudeReq *ClaudeRequest, projectID, map
}
// 检测是否启用 thinking
isThinkingEnabled := claudeReq.Thinking != nil && claudeReq.Thinking.Type == "enabled"
isThinkingEnabled := claudeReq.Thinking != nil && (claudeReq.Thinking.Type == "enabled" || claudeReq.Thinking.Type == "adaptive")
// 只有 Gemini 模型支持 dummy thought workaround
// Claude 模型通过 Vertex/Google API 需要有效的 thought signatures
@@ -198,8 +202,7 @@ type modelInfo struct {
// modelInfoMap 模型前缀 → 模型信息映射
// 只有在此映射表中的模型才会注入身份提示词
// 注意:当前 claude-opus-4-6 会被映射到 claude-opus-4-5-thinking
// 但保留此条目以便后续 Antigravity 上游支持 4.6 时快速切换
// 注意:模型映射逻辑在网关层完成;这里仅用于按模型前缀判断是否注入身份提示词。
var modelInfoMap = map[string]modelInfo{
"claude-opus-4-5": {DisplayName: "Claude Opus 4.5", CanonicalID: "claude-opus-4-5-20250929"},
"claude-opus-4-6": {DisplayName: "Claude Opus 4.6", CanonicalID: "claude-opus-4-6"},
@@ -593,6 +596,10 @@ func maxOutputTokensLimit(model string) int {
return maxOutputTokensUpperBound
}
func isAntigravityOpus46Model(model string) bool {
return strings.HasPrefix(strings.ToLower(model), "claude-opus-4-6")
}
func buildGenerationConfig(req *ClaudeRequest) *GeminiGenerationConfig {
maxLimit := maxOutputTokensLimit(req.Model)
config := &GeminiGenerationConfig{
@@ -606,25 +613,36 @@ func buildGenerationConfig(req *ClaudeRequest) *GeminiGenerationConfig {
}
// Thinking 配置
if req.Thinking != nil && req.Thinking.Type == "enabled" {
if req.Thinking != nil && (req.Thinking.Type == "enabled" || req.Thinking.Type == "adaptive") {
config.ThinkingConfig = &GeminiThinkingConfig{
IncludeThoughts: true,
}
// - thinking.type=enabledbudget_tokens>0 用显式预算
// - thinking.type=adaptive仅在 Antigravity 的 Opus 4.6 上覆写为 24576
budget := -1
if req.Thinking.BudgetTokens > 0 {
budget := req.Thinking.BudgetTokens
budget = req.Thinking.BudgetTokens
}
if req.Thinking.Type == "adaptive" && isAntigravityOpus46Model(req.Model) {
budget = ClaudeAdaptiveHighThinkingBudgetTokens
}
// 正预算需要做上限与 max_tokens 约束;动态预算(-1直接透传给上游。
if budget > 0 {
// gemini-2.5-flash 上限
if strings.Contains(req.Model, "gemini-2.5-flash") && budget > Gemini25FlashThinkingBudgetLimit {
budget = Gemini25FlashThinkingBudgetLimit
}
config.ThinkingConfig.ThinkingBudget = budget
// 自动修正max_tokens 必须大于 budget_tokens
// 自动修正max_tokens 必须大于 budget_tokensClaude 上游要求)
if adjusted, ok := ensureMaxTokensGreaterThanBudget(config.MaxOutputTokens, budget); ok {
log.Printf("[Antigravity] Auto-adjusted max_tokens from %d to %d (must be > budget_tokens=%d)",
config.MaxOutputTokens, adjusted, budget)
config.MaxOutputTokens = adjusted
}
}
config.ThinkingConfig.ThinkingBudget = budget
}
if config.MaxOutputTokens > maxLimit {

View File

@@ -259,3 +259,93 @@ func TestBuildTools_CustomTypeTools(t *testing.T) {
})
}
}
func TestBuildGenerationConfig_ThinkingDynamicBudget(t *testing.T) {
tests := []struct {
name string
model string
thinking *ThinkingConfig
wantBudget int
wantPresent bool
}{
{
name: "enabled without budget defaults to dynamic (-1)",
model: "claude-opus-4-6-thinking",
thinking: &ThinkingConfig{Type: "enabled"},
wantBudget: -1,
wantPresent: true,
},
{
name: "enabled with budget uses the provided value",
model: "claude-opus-4-6-thinking",
thinking: &ThinkingConfig{Type: "enabled", BudgetTokens: 1024},
wantBudget: 1024,
wantPresent: true,
},
{
name: "enabled with -1 budget uses dynamic (-1)",
model: "claude-opus-4-6-thinking",
thinking: &ThinkingConfig{Type: "enabled", BudgetTokens: -1},
wantBudget: -1,
wantPresent: true,
},
{
name: "adaptive on opus4.6 maps to high budget (24576)",
model: "claude-opus-4-6-thinking",
thinking: &ThinkingConfig{Type: "adaptive", BudgetTokens: 20000},
wantBudget: ClaudeAdaptiveHighThinkingBudgetTokens,
wantPresent: true,
},
{
name: "adaptive on non-opus model keeps default dynamic (-1)",
model: "claude-sonnet-4-5-thinking",
thinking: &ThinkingConfig{Type: "adaptive"},
wantBudget: -1,
wantPresent: true,
},
{
name: "disabled does not emit thinkingConfig",
model: "claude-opus-4-6-thinking",
thinking: &ThinkingConfig{Type: "disabled", BudgetTokens: 1024},
wantBudget: 0,
wantPresent: false,
},
{
name: "nil thinking does not emit thinkingConfig",
model: "claude-opus-4-6-thinking",
thinking: nil,
wantBudget: 0,
wantPresent: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := &ClaudeRequest{
Model: tt.model,
Thinking: tt.thinking,
}
cfg := buildGenerationConfig(req)
if cfg == nil {
t.Fatalf("expected non-nil generationConfig")
}
if tt.wantPresent {
if cfg.ThinkingConfig == nil {
t.Fatalf("expected thinkingConfig to be present")
}
if !cfg.ThinkingConfig.IncludeThoughts {
t.Fatalf("expected includeThoughts=true")
}
if cfg.ThinkingConfig.ThinkingBudget != tt.wantBudget {
t.Fatalf("expected thinkingBudget=%d, got %d", tt.wantBudget, cfg.ThinkingConfig.ThinkingBudget)
}
return
}
if cfg.ThinkingConfig != nil {
t.Fatalf("expected thinkingConfig to be nil, got %+v", cfg.ThinkingConfig)
}
})
}
}

View File

@@ -280,7 +280,7 @@ func (p *NonStreamingProcessor) buildResponse(geminiResp *GeminiResponse, respon
if geminiResp.UsageMetadata != nil {
cached := geminiResp.UsageMetadata.CachedContentTokenCount
usage.InputTokens = geminiResp.UsageMetadata.PromptTokenCount - cached
usage.OutputTokens = geminiResp.UsageMetadata.CandidatesTokenCount
usage.OutputTokens = geminiResp.UsageMetadata.CandidatesTokenCount + geminiResp.UsageMetadata.ThoughtsTokenCount
usage.CacheReadInputTokens = cached
}

View File

@@ -85,7 +85,7 @@ func (p *StreamingProcessor) ProcessLine(line string) []byte {
if geminiResp.UsageMetadata != nil {
cached := geminiResp.UsageMetadata.CachedContentTokenCount
p.inputTokens = geminiResp.UsageMetadata.PromptTokenCount - cached
p.outputTokens = geminiResp.UsageMetadata.CandidatesTokenCount
p.outputTokens = geminiResp.UsageMetadata.CandidatesTokenCount + geminiResp.UsageMetadata.ThoughtsTokenCount
p.cacheReadTokens = cached
}
@@ -146,7 +146,7 @@ func (p *StreamingProcessor) emitMessageStart(v1Resp *V1InternalResponse) []byte
if v1Resp.Response.UsageMetadata != nil {
cached := v1Resp.Response.UsageMetadata.CachedContentTokenCount
usage.InputTokens = v1Resp.Response.UsageMetadata.PromptTokenCount - cached
usage.OutputTokens = v1Resp.Response.UsageMetadata.CandidatesTokenCount
usage.OutputTokens = v1Resp.Response.UsageMetadata.CandidatesTokenCount + v1Resp.Response.UsageMetadata.ThoughtsTokenCount
usage.CacheReadInputTokens = cached
}

View File

@@ -448,7 +448,12 @@ func (r *accountRepository) ListWithFilters(ctx context.Context, params paginati
q = q.Where(dbaccount.TypeEQ(accountType))
}
if status != "" {
q = q.Where(dbaccount.StatusEQ(status))
switch status {
case "rate_limited":
q = q.Where(dbaccount.RateLimitResetAtGT(time.Now()))
default:
q = q.Where(dbaccount.StatusEQ(status))
}
}
if search != "" {
q = q.Where(dbaccount.NameContainsFold(search))

View File

@@ -6,6 +6,7 @@ import (
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/ent/redeemcode"
"github.com/Wei-Shaw/sub2api/ent/user"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/service"
)
@@ -106,7 +107,12 @@ func (r *redeemCodeRepository) ListWithFilters(ctx context.Context, params pagin
q = q.Where(redeemcode.StatusEQ(status))
}
if search != "" {
q = q.Where(redeemcode.CodeContainsFold(search))
q = q.Where(
redeemcode.Or(
redeemcode.CodeContainsFold(search),
redeemcode.HasUserWith(user.EmailContainsFold(search)),
),
)
}
total, err := q.Count(ctx)

View File

@@ -10,6 +10,7 @@ import (
"time"
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/ent/apikey"
dbuser "github.com/Wei-Shaw/sub2api/ent/user"
"github.com/Wei-Shaw/sub2api/ent/userallowedgroup"
"github.com/Wei-Shaw/sub2api/ent/usersubscription"
@@ -191,6 +192,7 @@ func (r *userRepository) ListWithFilters(ctx context.Context, params pagination.
dbuser.EmailContainsFold(filters.Search),
dbuser.UsernameContainsFold(filters.Search),
dbuser.NotesContainsFold(filters.Search),
dbuser.HasAPIKeysWith(apikey.KeyContainsFold(filters.Search)),
),
)
}

View File

@@ -281,6 +281,7 @@ func registerAntigravityOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers)
{
antigravity.POST("/oauth/auth-url", h.Admin.AntigravityOAuth.GenerateAuthURL)
antigravity.POST("/oauth/exchange-code", h.Admin.AntigravityOAuth.ExchangeCode)
antigravity.POST("/oauth/refresh-token", h.Admin.AntigravityOAuth.RefreshToken)
}
}

View File

@@ -1309,7 +1309,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
return nil, s.writeClaudeError(c, http.StatusForbidden, "permission_error", fmt.Sprintf("model %s not in whitelist", claudeReq.Model))
}
// 应用 thinking 模式自动后缀:如果 thinking 开启且目标是 claude-sonnet-4-5自动改为 thinking 版本
thinkingEnabled := claudeReq.Thinking != nil && claudeReq.Thinking.Type == "enabled"
thinkingEnabled := claudeReq.Thinking != nil && (claudeReq.Thinking.Type == "enabled" || claudeReq.Thinking.Type == "adaptive")
mappedModel = applyThinkingModelSuffix(mappedModel, thinkingEnabled)
// 获取 access_token

View File

@@ -591,6 +591,75 @@ func TestHandleClaudeStreamingResponse_NormalComplete(t *testing.T) {
require.NotContains(t, body, "event: error")
}
// TestHandleGeminiStreamingResponse_ThoughtsTokenCount
// 验证Gemini 流式转发时 thoughtsTokenCount 被计入 OutputTokens
func TestHandleGeminiStreamingResponse_ThoughtsTokenCount(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
go func() {
defer func() { _ = pw.Close() }()
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":"Hello"}]}}],"usageMetadata":{"promptTokenCount":100,"candidatesTokenCount":20,"thoughtsTokenCount":50}}`)
fmt.Fprintln(pw, "")
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":" world"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":100,"candidatesTokenCount":30,"thoughtsTokenCount":80,"cachedContentTokenCount":10}}`)
fmt.Fprintln(pw, "")
}()
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
_ = pr.Close()
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.usage)
// promptTokenCount=100, cachedContentTokenCount=10 → InputTokens=90
require.Equal(t, 90, result.usage.InputTokens)
// candidatesTokenCount=30 + thoughtsTokenCount=80 → OutputTokens=110
require.Equal(t, 110, result.usage.OutputTokens)
require.Equal(t, 10, result.usage.CacheReadInputTokens)
}
// TestHandleClaudeStreamingResponse_ThoughtsTokenCount
// 验证Gemini→Claude 流式转换时 thoughtsTokenCount 被计入 OutputTokens
func TestHandleClaudeStreamingResponse_ThoughtsTokenCount(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newAntigravityTestService(&config.Config{
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
})
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
pr, pw := io.Pipe()
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
go func() {
defer func() { _ = pw.Close() }()
fmt.Fprintln(pw, `data: {"response":{"candidates":[{"content":{"parts":[{"text":"Hi"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":50,"candidatesTokenCount":10,"thoughtsTokenCount":25}}}`)
fmt.Fprintln(pw, "")
}()
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "gemini-2.5-pro")
_ = pr.Close()
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.usage)
// promptTokenCount=50 → InputTokens=50
require.Equal(t, 50, result.usage.InputTokens)
// candidatesTokenCount=10 + thoughtsTokenCount=25 → OutputTokens=35
require.Equal(t, 35, result.usage.OutputTokens)
}
// --- 流式客户端断开检测测试 ---
// TestStreamUpstreamResponse_ClientDisconnectDrainsUsage

View File

@@ -192,6 +192,43 @@ func (s *AntigravityOAuthService) RefreshToken(ctx context.Context, refreshToken
return nil, fmt.Errorf("token 刷新失败 (重试后): %w", lastErr)
}
// ValidateRefreshToken 用 refresh token 验证并获取完整的 token 信息(含 email 和 project_id
func (s *AntigravityOAuthService) ValidateRefreshToken(ctx context.Context, refreshToken string, proxyID *int64) (*AntigravityTokenInfo, error) {
var proxyURL string
if proxyID != nil {
proxy, err := s.proxyRepo.GetByID(ctx, *proxyID)
if err == nil && proxy != nil {
proxyURL = proxy.URL()
}
}
// 刷新 token
tokenInfo, err := s.RefreshToken(ctx, refreshToken, proxyURL)
if err != nil {
return nil, err
}
// 获取用户信息email
client := antigravity.NewClient(proxyURL)
userInfo, err := client.GetUserInfo(ctx, tokenInfo.AccessToken)
if err != nil {
fmt.Printf("[AntigravityOAuth] 警告: 获取用户信息失败: %v\n", err)
} else {
tokenInfo.Email = userInfo.Email
}
// 获取 project_id容错失败不阻塞
projectID, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3)
if loadErr != nil {
fmt.Printf("[AntigravityOAuth] 警告: 获取 project_id 失败(重试后): %v\n", loadErr)
tokenInfo.ProjectIDMissing = true
} else {
tokenInfo.ProjectID = projectID
}
return tokenInfo, nil
}
func isNonRetryableAntigravityOAuthError(err error) bool {
msg := err.Error()
nonRetryable := []string{

View File

@@ -101,9 +101,9 @@ func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) {
}
}
// thinking: {type: "enabled"}
// thinking: {type: "enabled" | "adaptive"}
if rawThinking, ok := req["thinking"].(map[string]any); ok {
if t, ok := rawThinking["type"].(string); ok && t == "enabled" {
if t, ok := rawThinking["type"].(string); ok && (t == "enabled" || t == "adaptive") {
parsed.ThinkingEnabled = true
}
}
@@ -161,9 +161,9 @@ func parseIntegralNumber(raw any) (int, bool) {
// Returns filtered body or original body if filtering fails (fail-safe)
// This prevents 400 errors from invalid thinking block signatures
//
// Strategy:
// - When thinking.type != "enabled": Remove all thinking blocks
// - When thinking.type == "enabled": Only remove thinking blocks without valid signatures
// 策略:
// - thinking.type 不是 "enabled"/"adaptive":移除所有 thinking 相关块
// - thinking.type "enabled"/"adaptive":仅移除缺失/无效 signature 的 thinking 块(避免 400
// (blocks with missing/empty/dummy signatures that would cause 400 errors)
func FilterThinkingBlocks(body []byte) []byte {
return filterThinkingBlocksInternal(body, false)
@@ -489,9 +489,9 @@ func FilterSignatureSensitiveBlocksForRetry(body []byte) []byte {
}
// filterThinkingBlocksInternal removes invalid thinking blocks from request
// Strategy:
// - When thinking.type != "enabled": Remove all thinking blocks
// - When thinking.type == "enabled": Only remove thinking blocks without valid signatures
// 策略:
// - thinking.type 不是 "enabled"/"adaptive":移除所有 thinking 相关块
// - thinking.type "enabled"/"adaptive":仅移除缺失/无效 signature 的 thinking 块
func filterThinkingBlocksInternal(body []byte, _ bool) []byte {
// Fast path: if body doesn't contain "thinking", skip parsing
if !bytes.Contains(body, []byte(`"type":"thinking"`)) &&
@@ -511,7 +511,7 @@ func filterThinkingBlocksInternal(body []byte, _ bool) []byte {
// Check if thinking is enabled
thinkingEnabled := false
if thinking, ok := req["thinking"].(map[string]any); ok {
if thinkType, ok := thinking["type"].(string); ok && thinkType == "enabled" {
if thinkType, ok := thinking["type"].(string); ok && (thinkType == "enabled" || thinkType == "adaptive") {
thinkingEnabled = true
}
}

View File

@@ -29,6 +29,14 @@ func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) {
require.True(t, parsed.ThinkingEnabled)
}
func TestParseGatewayRequest_ThinkingAdaptiveEnabled(t *testing.T) {
body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"adaptive"},"messages":[{"content":"hi"}]}`)
parsed, err := ParseGatewayRequest(body, "")
require.NoError(t, err)
require.Equal(t, "claude-sonnet-4-5", parsed.Model)
require.True(t, parsed.ThinkingEnabled)
}
func TestParseGatewayRequest_MaxTokens(t *testing.T) {
body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1}`)
parsed, err := ParseGatewayRequest(body, "")
@@ -209,6 +217,16 @@ func TestFilterThinkingBlocks(t *testing.T) {
input: `{"model":"claude-3-5-sonnet-20241022","messages":[{"role":"user","content":[{"type":"text","text":"Hello"},{"type":"thinking","thinking":"internal","signature":"invalid"},{"type":"text","text":"World"}]}]}`,
shouldFilter: true,
},
{
name: "does not filter signed thinking blocks when thinking adaptive",
input: `{"thinking":{"type":"adaptive"},"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"ok","signature":"sig_real_123"},{"type":"text","text":"B"}]}]}`,
shouldFilter: false,
},
{
name: "filters unsigned thinking blocks when thinking adaptive",
input: `{"thinking":{"type":"adaptive"},"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"internal","signature":""},{"type":"text","text":"B"}]}]}`,
shouldFilter: true,
},
{
name: "handles no thinking blocks",
input: `{"model":"claude-3-5-sonnet-20241022","messages":[{"role":"user","content":[{"type":"text","text":"Hello"}]}]}`,

View File

@@ -3631,7 +3631,8 @@ func requestNeedsBetaFeatures(body []byte) bool {
if tools.Exists() && tools.IsArray() && len(tools.Array()) > 0 {
return true
}
if strings.EqualFold(gjson.GetBytes(body, "thinking.type").String(), "enabled") {
thinkingType := gjson.GetBytes(body, "thinking.type").String()
if strings.EqualFold(thinkingType, "enabled") || strings.EqualFold(thinkingType, "adaptive") {
return true
}
return false

View File

@@ -2663,11 +2663,12 @@ func extractGeminiUsage(geminiResp map[string]any) *ClaudeUsage {
prompt, _ := asInt(usageMeta["promptTokenCount"])
cand, _ := asInt(usageMeta["candidatesTokenCount"])
cached, _ := asInt(usageMeta["cachedContentTokenCount"])
thoughts, _ := asInt(usageMeta["thoughtsTokenCount"])
// 注意Gemini 的 promptTokenCount 包含 cachedContentTokenCount
// 但 Claude 的 input_tokens 不包含 cache_read_input_tokens需要减去
return &ClaudeUsage{
InputTokens: prompt - cached,
OutputTokens: cand,
OutputTokens: cand + thoughts,
CacheReadInputTokens: cached,
}
}

View File

@@ -4,6 +4,8 @@ import (
"encoding/json"
"strings"
"testing"
"github.com/stretchr/testify/require"
)
// TestConvertClaudeToolsToGeminiTools_CustomType 测试custom类型工具转换
@@ -203,3 +205,70 @@ func TestEnsureGeminiFunctionCallThoughtSignatures_InsertsWhenMissing(t *testing
t.Fatalf("expected injected thoughtSignature %q, got: %s", geminiDummyThoughtSignature, s)
}
}
func TestExtractGeminiUsage_ThoughtsTokenCount(t *testing.T) {
tests := []struct {
name string
resp map[string]any
wantInput int
wantOutput int
wantCacheRead int
wantNil bool
}{
{
name: "with thoughtsTokenCount",
resp: map[string]any{
"usageMetadata": map[string]any{
"promptTokenCount": float64(100),
"candidatesTokenCount": float64(20),
"thoughtsTokenCount": float64(50),
},
},
wantInput: 100,
wantOutput: 70,
},
{
name: "with thoughtsTokenCount and cache",
resp: map[string]any{
"usageMetadata": map[string]any{
"promptTokenCount": float64(100),
"candidatesTokenCount": float64(20),
"cachedContentTokenCount": float64(30),
"thoughtsTokenCount": float64(50),
},
},
wantInput: 70,
wantOutput: 70,
wantCacheRead: 30,
},
{
name: "without thoughtsTokenCount (old model)",
resp: map[string]any{
"usageMetadata": map[string]any{
"promptTokenCount": float64(100),
"candidatesTokenCount": float64(20),
},
},
wantInput: 100,
wantOutput: 20,
},
{
name: "no usageMetadata",
resp: map[string]any{},
wantNil: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
usage := extractGeminiUsage(tt.resp)
if tt.wantNil {
require.Nil(t, usage)
return
}
require.NotNil(t, usage)
require.Equal(t, tt.wantInput, usage.InputTokens)
require.Equal(t, tt.wantOutput, usage.OutputTokens)
require.Equal(t, tt.wantCacheRead, usage.CacheReadInputTokens)
})
}
}