Merge branch 'main' into test
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -115,6 +115,23 @@ type LoadCodeAssistResponse struct {
|
||||
IneligibleTiers []*IneligibleTier `json:"ineligibleTiers,omitempty"`
|
||||
}
|
||||
|
||||
// OnboardUserRequest onboardUser 请求
|
||||
type OnboardUserRequest struct {
|
||||
TierID string `json:"tierId"`
|
||||
Metadata struct {
|
||||
IDEType string `json:"ideType"`
|
||||
Platform string `json:"platform,omitempty"`
|
||||
PluginType string `json:"pluginType,omitempty"`
|
||||
} `json:"metadata"`
|
||||
}
|
||||
|
||||
// OnboardUserResponse onboardUser 响应
|
||||
type OnboardUserResponse struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
Done bool `json:"done"`
|
||||
Response map[string]any `json:"response,omitempty"`
|
||||
}
|
||||
|
||||
// GetTier 获取账户类型
|
||||
// 优先返回 paidTier(付费订阅级别),否则返回 currentTier
|
||||
func (r *LoadCodeAssistResponse) GetTier() string {
|
||||
@@ -371,6 +388,117 @@ func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadC
|
||||
return nil, nil, lastErr
|
||||
}
|
||||
|
||||
// OnboardUser 触发账号 onboarding,并返回 project_id
|
||||
// 说明:
|
||||
// 1) 部分账号 loadCodeAssist 不会立即返回 cloudaicompanionProject;
|
||||
// 2) 这时需要调用 onboardUser 完成初始化,之后才能拿到 project_id。
|
||||
func (c *Client) OnboardUser(ctx context.Context, accessToken, tierID string) (string, error) {
|
||||
tierID = strings.TrimSpace(tierID)
|
||||
if tierID == "" {
|
||||
return "", fmt.Errorf("tier_id 为空")
|
||||
}
|
||||
|
||||
reqBody := OnboardUserRequest{TierID: tierID}
|
||||
reqBody.Metadata.IDEType = "ANTIGRAVITY"
|
||||
reqBody.Metadata.Platform = "PLATFORM_UNSPECIFIED"
|
||||
reqBody.Metadata.PluginType = "GEMINI"
|
||||
|
||||
bodyBytes, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("序列化请求失败: %w", err)
|
||||
}
|
||||
|
||||
availableURLs := BaseURLs
|
||||
var lastErr error
|
||||
|
||||
for urlIdx, baseURL := range availableURLs {
|
||||
apiURL := baseURL + "/v1internal:onboardUser"
|
||||
|
||||
for attempt := 1; attempt <= 5; attempt++ {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("创建请求失败: %w", err)
|
||||
break
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", UserAgent)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("onboardUser 请求失败: %w", err)
|
||||
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
|
||||
log.Printf("[antigravity] onboardUser URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1])
|
||||
break
|
||||
}
|
||||
return "", lastErr
|
||||
}
|
||||
|
||||
respBodyBytes, err := io.ReadAll(resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("读取响应失败: %w", err)
|
||||
}
|
||||
|
||||
if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
|
||||
log.Printf("[antigravity] onboardUser URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1])
|
||||
break
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
lastErr = fmt.Errorf("onboardUser 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
|
||||
return "", lastErr
|
||||
}
|
||||
|
||||
var onboardResp OnboardUserResponse
|
||||
if err := json.Unmarshal(respBodyBytes, &onboardResp); err != nil {
|
||||
lastErr = fmt.Errorf("onboardUser 响应解析失败: %w", err)
|
||||
return "", lastErr
|
||||
}
|
||||
|
||||
if onboardResp.Done {
|
||||
if projectID := extractProjectIDFromOnboardResponse(onboardResp.Response); projectID != "" {
|
||||
DefaultURLAvailability.MarkSuccess(baseURL)
|
||||
return projectID, nil
|
||||
}
|
||||
lastErr = fmt.Errorf("onboardUser 完成但未返回 project_id")
|
||||
return "", lastErr
|
||||
}
|
||||
|
||||
// done=false 时等待后重试(与 CLIProxyAPI 行为一致)
|
||||
select {
|
||||
case <-time.After(2 * time.Second):
|
||||
case <-ctx.Done():
|
||||
return "", ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if lastErr != nil {
|
||||
return "", lastErr
|
||||
}
|
||||
return "", fmt.Errorf("onboardUser 未返回 project_id")
|
||||
}
|
||||
|
||||
func extractProjectIDFromOnboardResponse(resp map[string]any) string {
|
||||
if len(resp) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
if v, ok := resp["cloudaicompanionProject"]; ok {
|
||||
switch project := v.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(project)
|
||||
case map[string]any:
|
||||
if id, ok := project["id"].(string); ok {
|
||||
return strings.TrimSpace(id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// ModelQuotaInfo 模型配额信息
|
||||
type ModelQuotaInfo struct {
|
||||
RemainingFraction float64 `json:"remainingFraction"`
|
||||
|
||||
@@ -1655,3 +1655,74 @@ func TestClient_FetchAvailableModels_404Fallback_RealCall(t *testing.T) {
|
||||
t.Error("应返回 fallback server 的模型 m1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractProjectIDFromOnboardResponse(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
resp map[string]any
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "nil response",
|
||||
resp: nil,
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "empty response",
|
||||
resp: map[string]any{},
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "project as string",
|
||||
resp: map[string]any{
|
||||
"cloudaicompanionProject": "my-project-123",
|
||||
},
|
||||
want: "my-project-123",
|
||||
},
|
||||
{
|
||||
name: "project as string with spaces",
|
||||
resp: map[string]any{
|
||||
"cloudaicompanionProject": " my-project-123 ",
|
||||
},
|
||||
want: "my-project-123",
|
||||
},
|
||||
{
|
||||
name: "project as map with id",
|
||||
resp: map[string]any{
|
||||
"cloudaicompanionProject": map[string]any{
|
||||
"id": "proj-from-map",
|
||||
},
|
||||
},
|
||||
want: "proj-from-map",
|
||||
},
|
||||
{
|
||||
name: "project as map without id",
|
||||
resp: map[string]any{
|
||||
"cloudaicompanionProject": map[string]any{
|
||||
"name": "some-name",
|
||||
},
|
||||
},
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "missing cloudaicompanionProject key",
|
||||
resp: map[string]any{
|
||||
"otherField": "value",
|
||||
},
|
||||
want: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got := extractProjectIDFromOnboardResponse(tc.resp)
|
||||
if got != tc.want {
|
||||
t.Fatalf("extractProjectIDFromOnboardResponse() = %q, want %q", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -282,7 +282,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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user