Merge branch 'Wei-Shaw:main' into main
This commit is contained in:
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -64,6 +64,10 @@ const MaxTokensBudgetPadding = 1000
|
||||
// Gemini 2.5 Flash thinking budget 上限
|
||||
const Gemini25FlashThinkingBudgetLimit = 24576
|
||||
|
||||
// 对于 Antigravity 的 Claude(budget-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=enabled:budget_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_tokens(Claude 上游要求)
|
||||
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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"}]}]}`,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user